Merge try elastic#1

dnhatn · Sep 13, 2024 · 2454d8c · 2454d8c
1 parent ed88894
commit 2454d8c
Show file tree

Hide file tree

Showing 7 changed files with 1,068 additions and 10 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/codec/tsdb/ES87TSDBDocValuesConsumer.java b/server/src/main/java/org/elasticsearch/index/codec/tsdb/ES87TSDBDocValuesConsumer.java
@@ -45,7 +45,7 @@
 import static org.elasticsearch.index.codec.tsdb.ES87TSDBDocValuesFormat.DIRECT_MONOTONIC_BLOCK_SHIFT;
 import static org.elasticsearch.index.codec.tsdb.ES87TSDBDocValuesFormat.SORTED_SET;
 
-final class ES87TSDBDocValuesConsumer extends DocValuesConsumer {
+class ES87TSDBDocValuesConsumer extends DocValuesConsumer {
 
     IndexOutput data, meta;
     final int maxDoc;
@@ -98,14 +98,19 @@ public SortedNumericDocValues getSortedNumeric(FieldInfo field) throws IOExcepti
     private long[] writeField(FieldInfo field, DocValuesProducer valuesProducer, long maxOrd) throws IOException {
         int numDocsWithValue = 0;
         long numValues = 0;
-
         SortedNumericDocValues values = valuesProducer.getSortedNumeric(field);
-        for (int doc = values.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = values.nextDoc()) {
-            numDocsWithValue++;
-            final int count = values.docValueCount();
-            numValues += count;
+        if (values instanceof WithDocValuesCount wc) {
+            numDocsWithValue = wc.numDocsWithValues();
+            numValues = wc.totalValuesCount();
+        } else if (DocValues.unwrapSingleton(values) instanceof WithDocValuesCount wc) {
+            numDocsWithValue = wc.numDocsWithValues();
+            numValues = wc.totalValuesCount();
+        } else {
+            for (int doc = values.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = values.nextDoc()) {
+                numDocsWithValue++;
+                numValues += values.docValueCount();
+            }
         }
-
         if (numDocsWithValue == 0) { // meta[-2, 0]: No documents with values
             meta.writeLong(-2); // docsWithFieldOffset
             meta.writeLong(0L); // docsWithFieldLength

diff --git a/server/src/main/java/org/elasticsearch/index/codec/tsdb/ES87TSDBDocValuesFormat.java b/server/src/main/java/org/elasticsearch/index/codec/tsdb/ES87TSDBDocValuesFormat.java
@@ -48,7 +48,7 @@ public ES87TSDBDocValuesFormat() {
 
     @Override
     public DocValuesConsumer fieldsConsumer(SegmentWriteState state) throws IOException {
-        return new ES87TSDBDocValuesConsumer(state, DATA_CODEC, DATA_EXTENSION, META_CODEC, META_EXTENSION);
+        return new ES87TSDBDocValuesMerger(state, DATA_CODEC, DATA_EXTENSION, META_CODEC, META_EXTENSION);
     }
 
     @Override