Centralize max string length for functions, and add Replace tests for groups

Author: ivancea

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/function/scalar/ScalarFunction.java

@@ -13,6 +13,7 @@
 import java.util.List;
 
 import static java.util.Collections.emptyList;
+import static org.elasticsearch.common.unit.ByteSizeUnit.MB;
 
 /**
  * A {@code ScalarFunction} is a {@code Function} that takes values from some
@@ -22,6 +23,14 @@
  */
 public abstract class ScalarFunction extends Function {
 
+    /**
+     * Limit for the BytesRef return of functions.
+     * <p>
+     *     To be used when there's no CircuitBreaking, as an arbitrary measure to limit memory usage.
+     * </p>
+     */
+    public static final long MAX_BYTES_REF_RESULT_SIZE = MB.toBytes(1);
+
     protected ScalarFunction(Source source) {
         super(source, emptyList());
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/Repeat.java

@@ -40,8 +40,6 @@
 public class Repeat extends EsqlScalarFunction implements OptionalArgument {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Repeat", Repeat::new);
 
-    static final long MAX_REPEATED_LENGTH = MB.toBytes(1);
-
     private final Expression str;
     private final Expression number;
 
@@ -123,9 +121,9 @@ static BytesRef process(
 
     static BytesRef processInner(BreakingBytesRefBuilder scratch, BytesRef str, int number) {
         int repeatedLen = str.length * number;
-        if (repeatedLen > MAX_REPEATED_LENGTH) {
+        if (repeatedLen > MAX_BYTES_REF_RESULT_SIZE) {
             throw new IllegalArgumentException(
-                "Creating repeated strings with more than [" + MAX_REPEATED_LENGTH + "] bytes is not supported"
+                "Creating repeated strings with more than [" + MAX_BYTES_REF_RESULT_SIZE + "] bytes is not supported"
             );
         }
         scratch.grow(repeatedLen);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/Replace.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.compute.ann.Fixed;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
@@ -41,8 +42,6 @@
 public class Replace extends EsqlScalarFunction {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Replace", Replace::new);
 
-    static final long MAX_RESULT_LENGTH = MB.toBytes(1);
-
     private final Expression str;
     private final Expression regex;
     private final Expression newStr;
@@ -138,7 +137,6 @@ static BytesRef process(BytesRef str, BytesRef regex, BytesRef newStr) {
         if (str == null) {
             return null;
         }
-
         if (regex == null || newStr == null) {
             return str;
         }
@@ -160,8 +158,8 @@ private static BytesRef safeReplace(BytesRef strBytesRef, Pattern regex, BytesRe
         do {
             m.appendReplacement(result, newStr);
 
-            if (result.length() > MAX_RESULT_LENGTH) {
-                throw new IllegalArgumentException("Creating strings with more than [" + MAX_RESULT_LENGTH + "] bytes is not supported");
+            if (result.length() > MAX_BYTES_REF_RESULT_SIZE) {
+                throw new IllegalArgumentException("Creating strings with more than [" + MAX_BYTES_REF_RESULT_SIZE + "] bytes is not supported");
             }
         } while (m.find());
         m.appendTail(result);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/RepeatStaticTests.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.compute.test.TestBlockFactory;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
+import org.elasticsearch.xpack.esql.core.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
@@ -45,14 +46,14 @@ public class RepeatStaticTests extends ESTestCase {
 
     public void testAlmostTooBig() {
         String str = randomAlphaOfLength(1);
-        int number = (int) Repeat.MAX_REPEATED_LENGTH;
+        int number = (int) ScalarFunction.MAX_BYTES_REF_RESULT_SIZE;
         String repeated = process(str, number);
         assertThat(repeated, equalTo(str.repeat(number)));
     }
 
     public void testTooBig() {
         String str = randomAlphaOfLength(1);
-        int number = (int) Repeat.MAX_REPEATED_LENGTH + 1;
+        int number = (int) ScalarFunction.MAX_BYTES_REF_RESULT_SIZE + 1;
         String repeated = process(str, number);
         assertNull(repeated);
         assertWarnings(

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ReplaceTests.java

@@ -12,6 +12,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase;
@@ -22,7 +23,6 @@
 import java.util.function.Supplier;
 import java.util.regex.PatternSyntaxException;
 
-import static org.elasticsearch.xpack.esql.expression.function.scalar.string.Replace.MAX_RESULT_LENGTH;
 import static org.hamcrest.Matchers.equalTo;
 
 public class ReplaceTests extends AbstractScalarFunctionTestCase {
@@ -79,6 +79,24 @@ public static Iterable<Object[]> parameters() {
             )
         );
 
+        // Groups
+        suppliers.add(fixedCase("Full group", "Cats are awesome", ".+", "<$0>", "<Cats are awesome>"));
+        suppliers.add(fixedCase(
+            "Nested groups",
+            "A cat is great, a cat is awesome",
+            "\\b([Aa] (\\w+)) is (\\w+)\\b",
+            "$1$2",
+            "A catcat, a catcat"
+        ));
+        suppliers.add(fixedCase(
+            "Multiple groups",
+            "Cats are awesome",
+            "(\\w+) (.+)",
+            "$0 -> $1 and dogs $2",
+            "Cats are awesome -> Cats and dogs are awesome"
+        ));
+
+        // Errors
         suppliers.add(new TestCaseSupplier("syntax error", List.of(DataType.KEYWORD, DataType.KEYWORD, DataType.KEYWORD), () -> {
             String text = randomAlphaOfLength(10);
             String invalidRegex = "[";
@@ -106,7 +124,7 @@ public static Iterable<Object[]> parameters() {
         }));
 
         suppliers.add(new TestCaseSupplier("result too big", List.of(DataType.KEYWORD, DataType.KEYWORD, DataType.KEYWORD), () -> {
-            String textAndNewStr = randomAlphaOfLength((int) (MAX_RESULT_LENGTH / 10));
+            String textAndNewStr = randomAlphaOfLength((int) (ScalarFunction.MAX_BYTES_REF_RESULT_SIZE / 10));
             String regex = ".";
             return new TestCaseSupplier.TestCase(
                 List.of(
@@ -121,7 +139,7 @@ public static Iterable<Object[]> parameters() {
                 .withWarning(
                     "Line 1:1: java.lang.IllegalArgumentException: "
                         + "Creating strings with more than ["
-                        + MAX_RESULT_LENGTH
+                        + ScalarFunction.MAX_BYTES_REF_RESULT_SIZE
                         + "] bytes is not supported"
                 );
         }));