ESQL: Fix usage of already released null block in ValueSourceReaderOperator (elastic#126411) (elastic#126480)

* Add yaml test with reproducer
* Fix the bug
* Make ComputeBlockLoaderFactory Releasable

(cherry picked from commit ffd4913)

# Conflicts:
#	x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

Author: alex-spies

docs/changelog/126411.yaml

@@ -0,0 +1,6 @@
+pr: 126411
+summary: Fix usage of already released null block in `ValueSourceReaderOperator`
+area: ES|QL
+type: bug
+issues:
+ - 125850

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

@@ -220,9 +220,8 @@ private void loadFromSingleLeaf(Block[] blocks, int shard, int segment, BlockLoa
         positionFieldWork(shard, segment, firstDoc);
         StoredFieldsSpec storedFieldsSpec = StoredFieldsSpec.NO_REQUIREMENTS;
         List<RowStrideReaderWork> rowStrideReaders = new ArrayList<>(fields.length);
-        ComputeBlockLoaderFactory loaderBlockFactory = new ComputeBlockLoaderFactory(blockFactory, docs.count());
         LeafReaderContext ctx = ctx(shard, segment);
-        try {
+        try (ComputeBlockLoaderFactory loaderBlockFactory = new ComputeBlockLoaderFactory(blockFactory, docs.count())) {
             for (int f = 0; f < fields.length; f++) {
                 FieldWork field = fields[f];
                 BlockLoader.ColumnAtATimeReader columnAtATime = field.columnAtATime(ctx);
@@ -345,27 +344,28 @@ void run() throws IOException {
                 builders[f] = new Block.Builder[shardContexts.size()];
                 converters[f] = new BlockLoader[shardContexts.size()];
             }
-            ComputeBlockLoaderFactory loaderBlockFactory = new ComputeBlockLoaderFactory(blockFactory, docs.getPositionCount());
-            int p = forwards[0];
-            int shard = shards.getInt(p);
-            int segment = segments.getInt(p);
-            int firstDoc = docs.getInt(p);
-            positionFieldWork(shard, segment, firstDoc);
-            LeafReaderContext ctx = ctx(shard, segment);
-            fieldsMoved(ctx, shard);
-            verifyBuilders(loaderBlockFactory, shard);
-            read(firstDoc, shard);
-            for (int i = 1; i < forwards.length; i++) {
-                p = forwards[i];
-                shard = shards.getInt(p);
-                segment = segments.getInt(p);
-                boolean changedSegment = positionFieldWorkDocGuarteedAscending(shard, segment);
-                if (changedSegment) {
-                    ctx = ctx(shard, segment);
-                    fieldsMoved(ctx, shard);
-                }
+            try (ComputeBlockLoaderFactory loaderBlockFactory = new ComputeBlockLoaderFactory(blockFactory, docs.getPositionCount())) {
+                int p = forwards[0];
+                int shard = shards.getInt(p);
+                int segment = segments.getInt(p);
+                int firstDoc = docs.getInt(p);
+                positionFieldWork(shard, segment, firstDoc);
+                LeafReaderContext ctx = ctx(shard, segment);
+                fieldsMoved(ctx, shard);
                 verifyBuilders(loaderBlockFactory, shard);
-                read(docs.getInt(p), shard);
+                read(firstDoc, shard);
+                for (int i = 1; i < forwards.length; i++) {
+                    p = forwards[i];
+                    shard = shards.getInt(p);
+                    segment = segments.getInt(p);
+                    boolean changedSegment = positionFieldWorkDocGuarteedAscending(shard, segment);
+                    if (changedSegment) {
+                        ctx = ctx(shard, segment);
+                        fieldsMoved(ctx, shard);
+                    }
+                    verifyBuilders(loaderBlockFactory, shard);
+                    read(docs.getInt(p), shard);
+                }
             }
             for (int f = 0; f < target.length; f++) {
                 for (int s = 0; s < shardContexts.size(); s++) {
@@ -614,7 +614,7 @@ public String toString() {
         }
     }
 
-    private static class ComputeBlockLoaderFactory implements BlockLoader.BlockFactory {
+    private static class ComputeBlockLoaderFactory implements BlockLoader.BlockFactory, Releasable {
         private final BlockFactory factory;
         private final int pageSize;
         private Block nullBlock;
@@ -683,12 +683,18 @@ public BlockLoader.Builder nulls(int expectedCount) {
         public Block constantNulls() {
             if (nullBlock == null) {
                 nullBlock = factory.newConstantNullBlock(pageSize);
-            } else {
-                nullBlock.incRef();
             }
+            nullBlock.incRef();
             return nullBlock;
         }
 
+        @Override
+        public void close() {
+            if (nullBlock != null) {
+                nullBlock.close();
+            }
+        }
+
         @Override
         public BytesRefBlock constantBytes(BytesRef value) {
             return factory.newConstantBytesRefBlockWith(value, pageSize);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -782,7 +782,14 @@ public enum Cap {
         /**
          * Support for sorting when aggregate_metric_doubles are present
          */
-        AGGREGATE_METRIC_DOUBLE_SORTING(AGGREGATE_METRIC_DOUBLE_FEATURE_FLAG);
+        AGGREGATE_METRIC_DOUBLE_SORTING(AGGREGATE_METRIC_DOUBLE_FEATURE_FLAG),
+
+        /**
+         * When creating constant null blocks in {@link org.elasticsearch.compute.lucene.ValuesSourceReaderOperator}, we also handed off
+         * the ownership of that block - but didn't account for the fact that the caller might close it, leading to double releases
+         * in some union type queries. C.f. https://github.com/elastic/elasticsearch/issues/125850
+         */
+        FIX_DOUBLY_RELEASED_NULL_BLOCKS_IN_VALUESOURCEREADER;
 
         private final boolean enabled;
 

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/100_bug_fix.yml

@@ -338,3 +338,58 @@
       esql.query:
         body:
           query: 'FROM test_grok | KEEP name | WHERE last_name == "Facello" | EVAL name = concat("1 ", last_name) | GROK name "%{NUMBER:foo} %{WORD:foo}"'
+---
+"union types with null blocks from missing fields #125850":
+  - requires:
+      test_runner_features: [allowed_warnings_regex, capabilities]
+      capabilities:
+        - method: POST
+          path: /_query
+          parameters: []
+          capabilities: [fix_doubly_released_null_blocks_in_valuesourcereader]
+      reason: "fixed handing out already closed null block references in ValueSourceReader"
+  - do:
+      indices.create:
+        index: test1
+        body:
+          mappings:
+            properties:
+              truefalse1 :
+                type : boolean
+              truefalse2 :
+                type: boolean
+  - do:
+      indices.create:
+        index: test2
+        body:
+          mappings:
+            properties:
+              truefalse1 :
+                type : keyword
+              truefalse2 :
+                type: keyword
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - { "index": { "_index": "test1" } }
+          - { "truefalse1": null}
+          - { "index": { "_index": "test2" } }
+          - { "truefalse1": null }
+
+  - do:
+      allowed_warnings_regex:
+        - "No limit defined, adding default limit of \\[.*\\]"
+
+      esql.query:
+        body:
+          query: 'FROM test* | eval t1 = truefalse1::boolean, t2 = truefalse2::boolean | keep t1, t2'
+  - match: { columns.0.name: t1 }
+  - match: { columns.0.type: boolean }
+  - match: { columns.1.name: t2 }
+  - match: { columns.1.type: boolean }
+  - length: { values: 2 }
+  - match: { values.0.0: null }
+  - match: { values.0.1: null }
+  - match: { values.1.0: null }
+  - match: { values.1.1: null }