Add stats tracking for semantic field

Signed-off-by: Bo Zhang <[email protected]>

Author: bzhangam

CHANGELOG.md

@@ -19,6 +19,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Support custom weights in RRF normalization processor ([#1322](https://github.com/opensearch-project/neural-search/pull/1322))
 - [Stats] Add stats tracking for semantic highlighting ([#1327](https://github.com/opensearch-project/neural-search/pull/1327))
 - [Stats] Add stats for text embedding processor with different settings ([#1332](https://github.com/opensearch-project/neural-search/pull/1332))
+- [Stats] Add stats tracking for semantic field ([#1362](https://github.com/opensearch-project/neural-search/pull/1362))
+
 ### 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/semantic/SemanticFieldProcessor.java

@@ -20,6 +20,8 @@
 import org.opensearch.neuralsearch.processor.chunker.Chunker;
 import org.opensearch.neuralsearch.processor.chunker.FixedTokenLengthChunker;
 import org.opensearch.neuralsearch.processor.dto.SemanticFieldInfo;
+import org.opensearch.neuralsearch.stats.events.EventStatName;
+import org.opensearch.neuralsearch.stats.events.EventStatsManager;
 import org.opensearch.neuralsearch.util.TokenWeightUtil;
 import org.opensearch.neuralsearch.util.prune.PruneType;
 import org.opensearch.neuralsearch.util.prune.PruneUtils;
@@ -118,6 +120,7 @@ public IngestDocument execute(IngestDocument ingestDocument) throws Exception {
      */
     @Override
     public void execute(IngestDocument ingestDocument, BiConsumer<IngestDocument, Exception> handler) {
+        EventStatsManager.increment(EventStatName.SEMANTIC_FIELD_PROCESSOR_EXECUTIONS);
         try {
             unflattenIngestDoc(ingestDocument);
             // Collect all the semantic field info based on the path of semantic fields found in the index mapping
@@ -173,11 +176,23 @@ private void process(
     ) {
         setModelInfo(ingestDocument, semanticFieldInfoList);
 
+        recordChunking(List.of(semanticFieldInfoList));
         chunk(ingestDocument, semanticFieldInfoList);
 
         generateAndSetEmbedding(ingestDocument, semanticFieldInfoList, handler);
     }
 
+    private void recordChunking(@NonNull Collection<List<SemanticFieldInfo>> semanticFieldInfoLists) {
+        for (List<SemanticFieldInfo> semanticFieldInfoList : semanticFieldInfoLists) {
+            for (SemanticFieldInfo semanticFieldInfo : semanticFieldInfoList) {
+                if (semanticFieldInfo.getChunkingEnabled()) {
+                    EventStatsManager.increment(EventStatName.SEMANTIC_FIELD_PROCESSOR_CHUNKING_EXECUTIONS);
+                    break;
+                }
+            }
+        }
+    }
+
     private void setModelInfo(@NonNull final IngestDocument ingestDocument, @NonNull final List<SemanticFieldInfo> semanticFieldInfoList) {
         final Map<String, Map<String, Object>> modelIdToInfoMap = new HashMap<>();
         for (final Map.Entry<String, MLModel> entry : modelIdToModelMap.entrySet()) {
@@ -386,6 +401,7 @@ private void collectSemanticFieldInfo(
 
     @Override
     public void subBatchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers, Consumer<List<IngestDocumentWrapper>> handler) {
+        EventStatsManager.increment(EventStatName.SEMANTIC_FIELD_PROCESSOR_EXECUTIONS);
         if (ingestDocumentWrappers == null || ingestDocumentWrappers.isEmpty()) {
             handler.accept(ingestDocumentWrappers);
             return;
@@ -451,7 +467,7 @@ private void batchProcess(
         @NonNull final Map<IngestDocumentWrapper, List<SemanticFieldInfo>> docToSemanticFieldInfoMap,
         @NonNull final Consumer<List<IngestDocumentWrapper>> handler
     ) {
-
+        recordChunking(docToSemanticFieldInfoMap.values());
         for (Map.Entry<IngestDocumentWrapper, List<SemanticFieldInfo>> entry : docToSemanticFieldInfoMap.entrySet()) {
             final IngestDocumentWrapper ingestDocumentWrapper = entry.getKey();
             final IngestDocument ingestDocument = entry.getKey().getIngestDocument();

src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java

@@ -17,7 +17,6 @@
 import static org.opensearch.neuralsearch.common.MinClusterVersionUtil.isClusterOnOrAfterMinReqVersionForSemanticFieldType;
 import static org.opensearch.neuralsearch.common.VectorUtil.vectorAsListToArray;
 import static org.opensearch.neuralsearch.constants.MappingConstants.PATH_SEPARATOR;
-import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.DEFAULT_SEMANTIC_INFO_FIELD_NAME_SUFFIX;
 import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.EMBEDDING_FIELD_NAME;
 import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.CHUNKS_FIELD_NAME;
 import static org.opensearch.neuralsearch.processor.TextImageEmbeddingProcessor.INPUT_IMAGE;
@@ -83,6 +82,8 @@
 import org.opensearch.neuralsearch.common.MinClusterVersionUtil;
 import org.opensearch.neuralsearch.mapper.SemanticFieldMapper;
 import org.opensearch.neuralsearch.query.dto.NeuralQueryBuildStage;
+import org.opensearch.neuralsearch.stats.events.EventStatName;
+import org.opensearch.neuralsearch.stats.events.EventStatsManager;
 import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil;
 
 import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor;
@@ -551,6 +552,7 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws
      * @throws IOException can be thrown by parser
      */
     public static NeuralQueryBuilder fromXContent(XContentParser parser) throws IOException {
+        EventStatsManager.increment(EventStatName.NEURAL_QUERY_COUNT);
         final Builder builder = new Builder();
         if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
             throw new ParsingException(parser.getTokenLocation(), "Token must be START_OBJECT");
@@ -707,6 +709,7 @@ private QueryBuilder rewriteQueryForSemanticField(@NonNull final NeuralQueryTarg
         final String chunksPath = targetFieldConfig.getChunksPath();
 
         if (KNNVectorFieldMapper.CONTENT_TYPE.equals(embeddingFieldType)) {
+            EventStatsManager.increment(EventStatName.NEURAL_QUERY_AGAINST_SEMANTIC_DENSE_COUNT);
             if (modelIdToVectorSupplierMap == null
                 || modelIdToVectorSupplierMap.get(searchModelId) == null
                 || modelIdToVectorSupplierMap.get(searchModelId).get() == null) {
@@ -722,6 +725,7 @@ private QueryBuilder rewriteQueryForSemanticField(@NonNull final NeuralQueryTarg
                 return neuralKNNQueryBuilder;
             }
         } else if (RankFeaturesFieldMapper.CONTENT_TYPE.equals(embeddingFieldType)) {
+            EventStatsManager.increment(EventStatName.NEURAL_QUERY_AGAINST_SEMANTIC_SPARSE_COUNT);
             Supplier<Map<String, Float>> queryTokensSupplier = queryTokensMapSupplier;
             // If the raw token is not provided or no search analyzer provided
             // then try to find the token generated by the ml model
@@ -797,7 +801,7 @@ private QueryBuilder rewriteQueryAgainstKnnField(QueryRewriteContext queryRewrit
             if (vectorSupplier().get() == null) {
                 return this;
             }
-
+            EventStatsManager.increment(EventStatName.NEURAL_QUERY_AGAINST_DENSE_COUNT);
             return createKNNQueryBuilder(fieldName(), vectorSupplier.get());
         }
 
@@ -961,15 +965,6 @@ private String getErrorMessageWithBaseErrorForSemantic(@NonNull final String err
         return "Failed to rewrite the neural query against the semantic field " + fieldName + ". " + errorMessage;
     }
 
-    private String getSemanticInfoFieldPath(SemanticFieldMapper.SemanticFieldType semanticFieldType) {
-        final String[] paths = semanticFieldType.name().split("\\.");
-        final String semanticInfoFieldName = semanticFieldType.getSemanticParameters().getSemanticInfoFieldName();
-        paths[paths.length - 1] = semanticInfoFieldName == null
-            ? paths[paths.length - 1] + DEFAULT_SEMANTIC_INFO_FIELD_NAME_SUFFIX
-            : semanticInfoFieldName;
-        return String.join(PATH_SEPARATOR, paths) + PATH_SEPARATOR + CHUNKS_FIELD_NAME;
-    }
-
     private QueryBuilder inferenceForSemanticField(
         @NonNull final QueryRewriteContext queryRewriteContext,
         @NonNull final Set<String> modelIdsFromTargetFields,

src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java

@@ -28,6 +28,8 @@
 import org.apache.lucene.search.Query;
 import org.opensearch.OpenSearchException;
 import org.opensearch.Version;
+import org.opensearch.neuralsearch.stats.events.EventStatName;
+import org.opensearch.neuralsearch.stats.events.EventStatsManager;
 import org.opensearch.transport.client.Client;
 import org.opensearch.common.SetOnce;
 import org.opensearch.common.collect.Tuple;
@@ -250,6 +252,7 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws
      * @throws IOException can be thrown by parser
      */
     public static NeuralSparseQueryBuilder fromXContent(XContentParser parser) throws IOException {
+        EventStatsManager.increment(EventStatName.NEURAL_SPARSE_QUERY_COUNT);
         NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder();
         if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
             throw new ParsingException(parser.getTokenLocation(), "First token of " + NAME + "query must be START_OBJECT");

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

@@ -31,11 +31,30 @@ public enum EventStatName implements StatName {
         EventStatType.TIMESTAMPED_EVENT_COUNTER
     ),
     TEXT_CHUNKING_DELIMITER_EXECUTIONS("text_chunking_delimiter_executions", "processors.ingest", EventStatType.TIMESTAMPED_EVENT_COUNTER),
+    SEMANTIC_FIELD_PROCESSOR_EXECUTIONS("semantic_field_executions", "processors.ingest", EventStatType.TIMESTAMPED_EVENT_COUNTER),
+    SEMANTIC_FIELD_PROCESSOR_CHUNKING_EXECUTIONS(
+        "semantic_field_chunking_executions",
+        "processors.ingest",
+        EventStatType.TIMESTAMPED_EVENT_COUNTER
+    ),
     SEMANTIC_HIGHLIGHTING_REQUEST_COUNT(
         "semantic_highlighting_request_count",
         "semantic_highlighting",
         EventStatType.TIMESTAMPED_EVENT_COUNTER
-    );
+    ),
+    NEURAL_QUERY_COUNT("neural_query_requests", "query.neural", EventStatType.TIMESTAMPED_EVENT_COUNTER),
+    NEURAL_QUERY_AGAINST_DENSE_COUNT("neural_query_against_dense_requests", "query.neural", EventStatType.TIMESTAMPED_EVENT_COUNTER),
+    NEURAL_QUERY_AGAINST_SEMANTIC_DENSE_COUNT(
+        "neural_query_against_semantic_dense_requests",
+        "query.neural",
+        EventStatType.TIMESTAMPED_EVENT_COUNTER
+    ),
+    NEURAL_QUERY_AGAINST_SEMANTIC_SPARSE_COUNT(
+        "neural_query_against_semantic_sparse_requests",
+        "query.neural",
+        EventStatType.TIMESTAMPED_EVENT_COUNTER
+    ),
+    NEURAL_SPARSE_QUERY_COUNT("neural_sparse_requests", "query.neural_sparse", EventStatType.TIMESTAMPED_EVENT_COUNTER);
 
     private final String nameString;
     private final String path;

src/test/java/org/opensearch/neuralsearch/processor/semantic/SemanticFieldProcessorTests.java

@@ -54,6 +54,7 @@
 import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.CHUNKING;
 import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.MODEL_ID;
 import static org.opensearch.neuralsearch.processor.TextChunkingProcessorTests.getAnalysisRegistry;
+import org.opensearch.neuralsearch.util.TestUtils;
 
 public class SemanticFieldProcessorTests extends OpenSearchTestCase {
     @Mock
@@ -80,6 +81,9 @@ public class SemanticFieldProcessorTests extends OpenSearchTestCase {
 
     @Before
     public void setup() {
+        // Initialize EventStatsManager for tests
+        TestUtils.initializeEventStatsManager();
+
         MockitoAnnotations.openMocks(this);
         // mock env
         final Settings settings = Settings.builder()

src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderRewriteTests.java

@@ -38,6 +38,7 @@
 import org.opensearch.neuralsearch.processor.TextInferenceRequest;
 import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils;
 import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil;
+import org.opensearch.neuralsearch.util.TestUtils;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.transport.client.Client;
 
@@ -99,6 +100,8 @@ public class NeuralQueryBuilderRewriteTests extends OpenSearchTestCase {
     public void setup() throws Exception {
         mlClient = mock(MLCommonsClientAccessor.class);
         NeuralQueryBuilder.initialize(mlClient);
+        // Initialize EventStatsManager for tests
+        TestUtils.initializeEventStatsManager();
     }
 
     @SneakyThrows

src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java

@@ -60,6 +60,7 @@
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.knn.index.query.rescore.RescoreContext;
 import org.opensearch.neuralsearch.common.MinClusterVersionUtil;
+import org.opensearch.neuralsearch.util.TestUtils;
 import org.opensearch.test.OpenSearchTestCase;
 
 import lombok.SneakyThrows;
@@ -90,7 +91,10 @@ public class NeuralQueryBuilderTests extends OpenSearchTestCase {
     private static final QueryBuilder ADDITIONAL_TEST_FILTER = new TermQueryBuilder(TERM_QUERY_FIELD_NAME, TERM_QUERY_FIELD_VALUE);
 
     @Before
-    public void setup() throws Exception {}
+    public void setup() throws Exception {
+        // Initialize EventStatsManager for tests
+        TestUtils.initializeEventStatsManager();
+    }
 
     @SneakyThrows
     public void testFromXContent_whenBuiltWithDefaults_thenBuildSuccessfully() {

src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilderTests.java

@@ -49,6 +49,7 @@
 import org.opensearch.index.analysis.AnalyzerScope;
 import org.opensearch.index.analysis.IndexAnalyzers;
 import org.opensearch.index.analysis.NamedAnalyzer;
+import org.opensearch.neuralsearch.util.TestUtils;
 import org.opensearch.transport.client.Client;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.SetOnce;
@@ -87,8 +88,11 @@ public class NeuralSparseQueryBuilderTests extends OpenSearchTestCase {
     private static final Supplier<Map<String, Float>> QUERY_TOKENS_SUPPLIER = () -> Map.of("hello", 1.f, "world", 2.f);
 
     @Before
-    public void setupClusterServiceToCurrentVersion() {
+    public void setup() {
         setUpClusterService(Version.CURRENT);
+
+        // Initialize EventStatsManager for tests
+        TestUtils.initializeEventStatsManager();
     }
 
     @SneakyThrows