From bdb1ae5ad1be435920f1a980633066ee6711475f Mon Sep 17 00:00:00 2001
From: Luigi Dell'Aquila <luigi.dellaquila@gmail.com>
Date: Tue, 3 Oct 2023 09:47:39 +0200
Subject: [PATCH] ESQL: Add block factory to GROK command (#100139)

Add DriverContext and BlockFactory to GROK command to properly handle
circuit breaking See
https://github.com/elastic/elasticsearch/issues/99826
---
 .../operator/ColumnExtractOperator.java       | 48 +++++++++++--------
 .../operator/ColumnExtractOperatorTests.java  |  3 +-
 .../src/main/resources/grok.csv-spec          |  3 +-
 3 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ColumnExtractOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ColumnExtractOperator.java
index 0ccf575fc030d..58bf9e097bec3 100644
--- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ColumnExtractOperator.java
+++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ColumnExtractOperator.java
@@ -27,7 +27,7 @@ public record Factory(
 
         @Override
         public Operator get(DriverContext driverContext) {
-            return new ColumnExtractOperator(types, inputEvalSupplier.get(driverContext), evaluatorSupplier.get());
+            return new ColumnExtractOperator(types, inputEvalSupplier.get(driverContext), evaluatorSupplier.get(), driverContext);
         }
 
         @Override
@@ -39,15 +39,18 @@ public String describe() {
     private final ElementType[] types;
     private final EvalOperator.ExpressionEvaluator inputEvaluator;
     private final ColumnExtractOperator.Evaluator evaluator;
+    private final DriverContext driverContext;
 
     public ColumnExtractOperator(
         ElementType[] types,
-        EvalOperator.ExpressionEvaluator inputEvaluator,
-        ColumnExtractOperator.Evaluator evaluator
+        ExpressionEvaluator inputEvaluator,
+        Evaluator evaluator,
+        DriverContext driverContext
     ) {
         this.types = types;
         this.inputEvaluator = inputEvaluator;
         this.evaluator = evaluator;
+        this.driverContext = driverContext;
     }
 
     @Override
@@ -55,28 +58,33 @@ protected Page process(Page page) {
         int rowsCount = page.getPositionCount();
 
         Block.Builder[] blockBuilders = new Block.Builder[types.length];
-        for (int i = 0; i < types.length; i++) {
-            blockBuilders[i] = types[i].newBlockBuilder(rowsCount);
-        }
+        try {
+            for (int i = 0; i < types.length; i++) {
+                blockBuilders[i] = types[i].newBlockBuilder(rowsCount, driverContext.blockFactory());
+            }
 
-        try (Block.Ref ref = inputEvaluator.eval(page)) {
-            BytesRefBlock input = (BytesRefBlock) ref.block();
-            BytesRef spare = new BytesRef();
-            for (int row = 0; row < rowsCount; row++) {
-                if (input.isNull(row)) {
-                    for (int i = 0; i < blockBuilders.length; i++) {
-                        blockBuilders[i].appendNull();
+            try (Block.Ref ref = inputEvaluator.eval(page)) {
+                BytesRefBlock input = (BytesRefBlock) ref.block();
+                BytesRef spare = new BytesRef();
+                for (int row = 0; row < rowsCount; row++) {
+                    if (input.isNull(row)) {
+                        for (int i = 0; i < blockBuilders.length; i++) {
+                            blockBuilders[i].appendNull();
+                        }
+                        continue;
                     }
-                    continue;
+                    evaluator.computeRow(input, row, blockBuilders, spare);
+                }
+
+                Block[] blocks = new Block[blockBuilders.length];
+                for (int i = 0; i < blockBuilders.length; i++) {
+                    blocks[i] = blockBuilders[i].build();
                 }
-                evaluator.computeRow(input, row, blockBuilders, spare);
-            }
 
-            Block[] blocks = new Block[blockBuilders.length];
-            for (int i = 0; i < blockBuilders.length; i++) {
-                blocks[i] = blockBuilders[i].build();
+                return page.appendBlocks(blocks);
             }
-            return page.appendBlocks(blocks);
+        } finally {
+            Releasables.closeExpectNoException(blockBuilders);
         }
     }
 
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java
index 8e0be216ed477..49f3382f8ad2d 100644
--- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java
+++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ColumnExtractOperatorTests.java
@@ -89,7 +89,6 @@ protected void assertSimpleOutput(List<Page> input, List<Page> results) {
 
     @Override
     protected ByteSizeValue smallEnoughToCircuitBreak() {
-        assumeTrue("doesn't use big arrays so can't break", false);
-        return null;
+        return ByteSizeValue.ofBytes(between(1, 32));
     }
 }
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/grok.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/grok.csv-spec
index 9dc9444de0155..75dfdd783a7cd 100644
--- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/grok.csv-spec
+++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/grok.csv-spec
@@ -83,7 +83,8 @@ Anneke Preusig        | Anneke       | Preusig
 ;
 
 
-grokStats
+# AwaitsFix https://github.com/elastic/elasticsearch/issues/99826
+grokStats-Ignore
 from employees | eval x = concat(gender, " foobar") | grok x "%{WORD:a} %{WORD:b}" | stats n = max(emp_no) by a | keep a, n | sort a asc;
 
 a:keyword  | n:integer