Add short circuit if no live docs are in segments (#2059)

* Add short circuit if no live docs are in segments Signed-off-by: Vijayan Balasubramanian <balasvij@amazon.com> --------- Signed-off-by: Vijayan Balasubramanian <balasvij@amazon.com>
opensearch-project · Sep 12, 2024 · 7aedefd · 7aedefd
1 parent 5d10d64
commit 7aedefd
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,6 +17,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.17...2.x)
 ### Features
 ### Enhancements
+* Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059)
 ### Bug Fixes
 ### Infrastructure
 ### Documentation

diff --git a/...ain/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java b/...ain/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsWriter.java
@@ -244,26 +244,25 @@ private <T, C> void trainAndIndex(
         final String operationName
     ) throws IOException {
         final VectorDataType vectorDataType = extractVectorDataType(fieldInfo);
-        KNNVectorValues<T> knnVectorValues = vectorValuesRetriever.apply(vectorDataType, fieldInfo, VectorProcessingContext);
-        QuantizationParams quantizationParams = quantizationService.getQuantizationParams(fieldInfo);
-        QuantizationState quantizationState = null;
         // Count the docIds
         int totalLiveDocs = getLiveDocs(vectorValuesRetriever.apply(vectorDataType, fieldInfo, VectorProcessingContext));
-        if (quantizationParams != null && totalLiveDocs > 0) {
+        if (totalLiveDocs == 0) {
+            log.debug("No live docs for field " + fieldInfo.name);
+            return;
+        }
+        QuantizationState quantizationState = null;
+        QuantizationParams quantizationParams = quantizationService.getQuantizationParams(fieldInfo);
+        if (quantizationParams != null) {
             initQuantizationStateWriterIfNecessary();
+            KNNVectorValues<T> knnVectorValues = vectorValuesRetriever.apply(vectorDataType, fieldInfo, VectorProcessingContext);
             quantizationState = quantizationService.train(quantizationParams, knnVectorValues, totalLiveDocs);
             quantizationStateWriter.writeState(fieldInfo.getFieldNumber(), quantizationState);
         }
-        NativeIndexWriter writer = (quantizationParams != null)
-            ? NativeIndexWriter.getWriter(fieldInfo, segmentWriteState, quantizationState)
-            : NativeIndexWriter.getWriter(fieldInfo, segmentWriteState);
-
-        knnVectorValues = vectorValuesRetriever.apply(vectorDataType, fieldInfo, VectorProcessingContext);
-
-        StopWatch stopWatch = new StopWatch();
-        stopWatch.start();
-        indexOperation.buildAndWrite(writer, knnVectorValues, totalLiveDocs);
-        long time_in_millis = stopWatch.totalTime().millis();
+        NativeIndexWriter writer = NativeIndexWriter.getWriter(fieldInfo, segmentWriteState, quantizationState);
+        KNNVectorValues<T> knnVectors = vectorValuesRetriever.apply(vectorDataType, fieldInfo, VectorProcessingContext);
+        StopWatch stopWatch = new StopWatch().start();
+        indexOperation.buildAndWrite(writer, knnVectors, totalLiveDocs);
+        long time_in_millis = stopWatch.stop().totalTime().millis();
         graphBuildTime.incrementBy(time_in_millis);
         log.warn("Graph build took " + time_in_millis + " ms for " + operationName);
     }