Specialize SingletonOrdinalsBuilder to handle dense cases.

Author: martijnvg

server/src/main/java/org/elasticsearch/index/codec/tsdb/es819/ES819TSDBDocValuesProducer.java

@@ -373,7 +373,7 @@ public long cost() {
             @Override
             public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
                 if (ords instanceof BulkNumericDocValues b) {
-                    try (var builder = factory.singletonOrdinalsBuilder(this, docs.count() - offset)) {
+                    try (var builder = factory.singletonOrdinalsBuilder(this, docs.count() - offset, true)) {
                         return b.readOrdinals(builder, docs, offset);
                     }
                 } else {

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

@@ -663,7 +663,7 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset) throw
             if (ordinals instanceof BulkSortedDocValues bulkDv && bulkDv.supportsBlockRead()) {
                 return bulkDv.read(factory, docs, offset);
             }
-            try (var builder = factory.singletonOrdinalsBuilder(ordinals, docs.count() - offset)) {
+            try (var builder = factory.singletonOrdinalsBuilder(ordinals, docs.count() - offset, false)) {
                 for (int i = offset; i < docs.count(); i++) {
                     int doc = docs.get(i);
                     if (doc < ordinals.docID()) {

server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java

@@ -430,7 +430,7 @@ interface BlockFactory {
         /**
          * Build a reader for reading {@link SortedDocValues}
          */
-        SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count);
+        SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count, boolean isDense);
 
         /**
          * Build a reader for reading {@link SortedSetDocValues}

test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java

@@ -268,7 +268,11 @@ public BlockLoader.Block constantBytes(BytesRef value, int count) {
             }
 
             @Override
-            public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int expectedCount) {
+            public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(
+                SortedDocValues ordinals,
+                int expectedCount,
+                boolean isDense
+            ) {
                 class SingletonOrdsBuilder extends TestBlock.Builder implements BlockLoader.SingletonOrdinalsBuilder {
                     private SingletonOrdsBuilder() {
                         super(expectedCount);

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/DelegatingBlockLoaderFactory.java

@@ -87,8 +87,8 @@ public BlockLoader.Builder nulls(int expectedCount) {
     }
 
     @Override
-    public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count) {
-        return new SingletonOrdinalsBuilder(factory, ordinals, count);
+    public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count, boolean isDense) {
+        return new SingletonOrdinalsBuilder(factory, ordinals, count, isDense);
     }
 
     @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/SingletonOrdinalsBuilder.java

@@ -33,16 +33,23 @@ public class SingletonOrdinalsBuilder implements BlockLoader.SingletonOrdinalsBu
     private int maxOrd = Integer.MIN_VALUE;
     private final int[] ords;
     private int count;
+    private final boolean isDense;
 
     public SingletonOrdinalsBuilder(BlockFactory blockFactory, SortedDocValues docValues, int count) {
+        this(blockFactory, docValues, count, false);
+    }
+
+    public SingletonOrdinalsBuilder(BlockFactory blockFactory, SortedDocValues docValues, int count, boolean isDense) {
         this.blockFactory = blockFactory;
         this.docValues = docValues;
         blockFactory.adjustBreaker(ordsSize(count));
         this.ords = new int[count];
+        this.isDense = isDense;
     }
 
     @Override
     public SingletonOrdinalsBuilder appendNull() {
+        assert isDense == false;
         ords[count++] = -1; // real ords can't be < 0, so we use -1 as null
         return this;
     }
@@ -78,9 +85,11 @@ private BytesRefBlock tryBuildConstantBlock() {
         if (minOrd != maxOrd) {
             return null;
         }
-        for (int ord : ords) {
-            if (ord == -1) {
-                return null;
+        if (isDense == false) {
+            for (int ord : ords) {
+                if (ord == -1) {
+                    return null;
+                }
             }
         }
         final BytesRef v;
@@ -116,33 +125,61 @@ BytesRefBlock buildOrdinal() {
         try {
             int[] newOrds = new int[valueCount];
             Arrays.fill(newOrds, -1);
-            for (int ord : ords) {
-                if (ord != -1) {
+            // Re-mapping ordinals to be more space-efficient:
+            if (isDense) {
+                for (int ord : ords) {
                     newOrds[ord - minOrd] = 0;
                 }
+            } else {
+                for (int ord : ords) {
+                    if (ord != -1) {
+                        newOrds[ord - minOrd] = 0;
+                    }
+                }
             }
             // resolve the ordinals and remaps the ordinals
-            int nextOrd = -1;
-            try (BytesRefVector.Builder bytesBuilder = blockFactory.newBytesRefVectorBuilder(Math.min(valueCount, ords.length))) {
-                for (int i = 0; i < newOrds.length; i++) {
-                    if (newOrds[i] != -1) {
-                        newOrds[i] = ++nextOrd;
-                        bytesBuilder.appendBytesRef(docValues.lookupOrd(i + minOrd));
+            try {
+                int nextOrd = -1;
+                BytesRef firstTerm = minOrd != Integer.MAX_VALUE ? docValues.lookupOrd(minOrd) : null;
+                int estimatedSize;
+                if (firstTerm != null) {
+                    estimatedSize = Math.min(valueCount, ords.length) * firstTerm.length;
+                } else {
+                    estimatedSize = Math.min(valueCount, ords.length);
+                }
+                try (BytesRefVector.Builder bytesBuilder = blockFactory.newBytesRefVectorBuilder(estimatedSize)) {
+                    if (firstTerm != null) {
+                        newOrds[0] = ++nextOrd;
+                        bytesBuilder.appendBytesRef(firstTerm);
+                    }
+                    for (int i = firstTerm != null ? 1 : 0; i < newOrds.length; i++) {
+                        if (newOrds[i] != -1) {
+                            newOrds[i] = ++nextOrd;
+                            bytesBuilder.appendBytesRef(docValues.lookupOrd(i + minOrd));
+                        }
                     }
+                    bytesVector = bytesBuilder.build();
                 }
-                bytesVector = bytesBuilder.build();
             } catch (IOException e) {
                 throw new UncheckedIOException("error resolving ordinals", e);
             }
-            try (IntBlock.Builder ordinalsBuilder = blockFactory.newIntBlockBuilder(ords.length)) {
-                for (int ord : ords) {
-                    if (ord == -1) {
-                        ordinalsBuilder.appendNull();
-                    } else {
-                        ordinalsBuilder.appendInt(newOrds[ord - minOrd]);
+            if (isDense) {
+                // Reusing ords array and overwrite all slots with re-mapped ordinals
+                for (int i = 0; i < ords.length; i++) {
+                    ords[i] = newOrds[ords[i] - minOrd];
+                }
+                ordinalBlock = blockFactory.newIntArrayVector(ords, ords.length).asBlock();
+            } else {
+                try (IntBlock.Builder ordinalsBuilder = blockFactory.newIntBlockBuilder(ords.length)) {
+                    for (int ord : ords) {
+                        if (ord == -1) {
+                            ordinalsBuilder.appendNull();
+                        } else {
+                            ordinalsBuilder.appendInt(newOrds[ord - minOrd]);
+                        }
                     }
+                    ordinalBlock = ordinalsBuilder.build();
                 }
-                ordinalBlock = ordinalsBuilder.build();
             }
             final OrdinalBytesRefBlock result = new OrdinalBytesRefBlock(ordinalBlock, bytesVector);
             bytesVector = null;

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/read/TimeSeriesExtractFieldOperator.java

@@ -208,7 +208,7 @@ public BlockLoader.Block constantBytes(BytesRef value, int count) {
         }
 
         @Override
-        public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count) {
+        public BlockLoader.SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count, boolean isDense) {
             throw new UnsupportedOperationException("must not be used by column readers");
         }
     }