max_match - initial implementation

Signed-off-by: Jialiang Liang <jiallian@amazon.com>

Author: RyanL1997

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -187,26 +187,36 @@ public RelNode visitRex(Rex node, CalcitePlanContext context) {
           "Rex pattern must contain at least one named capture group");
     }
 
-    // Filter matching rows on data nodes using script pushdown
     RexNode regexMatchCondition =
         context.rexBuilder.makeCall(
             SqlLibraryOperators.REGEXP_CONTAINS,
             fieldRex,
             context.rexBuilder.makeLiteral(patternStr));
     context.relBuilder.filter(regexMatchCondition);
 
-    // Extract fields from filtered data
     List<RexNode> newFields = new ArrayList<>();
     List<String> newFieldNames = new ArrayList<>();
 
     for (int i = 0; i < namedGroups.size(); i++) {
-      RexNode extractCall =
-          PPLFuncImpTable.INSTANCE.resolve(
-              context.rexBuilder,
-              BuiltinFunctionName.REX_EXTRACT,
-              fieldRex,
-              context.rexBuilder.makeLiteral(patternStr),
-              context.relBuilder.literal(i + 1));
+      RexNode extractCall;
+      if (node.getMaxMatch().isPresent() && node.getMaxMatch().get() != 1) {
+        extractCall =
+            PPLFuncImpTable.INSTANCE.resolve(
+                context.rexBuilder,
+                BuiltinFunctionName.REX_EXTRACT_MULTI,
+                fieldRex,
+                context.rexBuilder.makeLiteral(patternStr),
+                context.relBuilder.literal(i + 1),
+                context.relBuilder.literal(node.getMaxMatch().get()));
+      } else {
+        extractCall =
+            PPLFuncImpTable.INSTANCE.resolve(
+                context.rexBuilder,
+                BuiltinFunctionName.REX_EXTRACT,
+                fieldRex,
+                context.rexBuilder.makeLiteral(patternStr),
+                context.relBuilder.literal(i + 1));
+      }
       newFields.add(extractCall);
       newFieldNames.add(namedGroups.get(i));
     }

core/src/main/java/org/opensearch/sql/calcite/utils/PPLOperandTypes.java

@@ -51,6 +51,13 @@ private PPLOperandTypes() {}
       UDFOperandMetadata.wrap(
           OperandTypes.family(
               SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.INTEGER));
+  public static final UDFOperandMetadata STRING_STRING_INTEGER_INTEGER =
+      UDFOperandMetadata.wrap(
+          OperandTypes.family(
+              SqlTypeFamily.CHARACTER,
+              SqlTypeFamily.CHARACTER,
+              SqlTypeFamily.INTEGER,
+              SqlTypeFamily.INTEGER));
 
   public static final UDFOperandMetadata NUMERIC_NUMERIC_NUMERIC =
       UDFOperandMetadata.wrap(

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -216,6 +216,7 @@ public enum BuiltinFunctionName {
   POSITION(FunctionName.of("position")),
   REGEXP(FunctionName.of("regexp")),
   REX_EXTRACT(FunctionName.of("REX_EXTRACT")),
+  REX_EXTRACT_MULTI(FunctionName.of("REX_EXTRACT_MULTI")),
   REX_OFFSET(FunctionName.of("REX_OFFSET")),
   REPLACE(FunctionName.of("replace")),
   REVERSE(FunctionName.of("reverse")),

core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java

@@ -49,6 +49,7 @@
 import org.opensearch.sql.expression.function.udf.GrokFunction;
 import org.opensearch.sql.expression.function.udf.RelevanceQueryFunction;
 import org.opensearch.sql.expression.function.udf.RexExtractFunction;
+import org.opensearch.sql.expression.function.udf.RexExtractMultiFunction;
 import org.opensearch.sql.expression.function.udf.RexOffsetFunction;
 import org.opensearch.sql.expression.function.udf.SpanFunction;
 import org.opensearch.sql.expression.function.udf.condition.EarliestFunction;
@@ -383,6 +384,8 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
   public static final SqlOperator NUMBER_TO_STRING =
       new NumberToStringFunction().toUDF("NUMBER_TO_STRING");
   public static final SqlOperator REX_EXTRACT = new RexExtractFunction().toUDF("REX_EXTRACT");
+  public static final SqlOperator REX_EXTRACT_MULTI =
+      new RexExtractMultiFunction().toUDF("REX_EXTRACT_MULTI");
   public static final SqlOperator REX_OFFSET = new RexOffsetFunction().toUDF("REX_OFFSET");
 
   /**

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -163,6 +163,7 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REPLACE;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REVERSE;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REX_EXTRACT;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.REX_EXTRACT_MULTI;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REX_OFFSET;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.RIGHT;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.RINT;
@@ -786,6 +787,7 @@ void populate() {
       registerOperator(QUERY_STRING, PPLBuiltinOperators.QUERY_STRING);
       registerOperator(MULTI_MATCH, PPLBuiltinOperators.MULTI_MATCH);
       registerOperator(REX_EXTRACT, PPLBuiltinOperators.REX_EXTRACT);
+      registerOperator(REX_EXTRACT_MULTI, PPLBuiltinOperators.REX_EXTRACT_MULTI);
       registerOperator(REX_OFFSET, PPLBuiltinOperators.REX_OFFSET);
 
       // Register PPL Datetime UDF operator

core/src/main/java/org/opensearch/sql/expression/function/udf/RexExtractFunction.java

@@ -6,6 +6,8 @@
 package org.opensearch.sql.expression.function.udf;
 
 import java.util.List;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import org.apache.calcite.adapter.enumerable.NotNullImplementor;
 import org.apache.calcite.adapter.enumerable.NullPolicy;
 import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
@@ -54,8 +56,8 @@ public static String extractGroup(String text, String pattern, int groupIndex) {
     }
 
     try {
-      java.util.regex.Pattern compiledPattern = java.util.regex.Pattern.compile(pattern);
-      java.util.regex.Matcher matcher = compiledPattern.matcher(text);
+      Pattern compiledPattern = Pattern.compile(pattern);
+      Matcher matcher = compiledPattern.matcher(text);
 
       if (matcher.find() && groupIndex > 0 && groupIndex <= matcher.groupCount()) {
         return matcher.group(groupIndex);

core/src/main/java/org/opensearch/sql/expression/function/udf/RexExtractMultiFunction.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.udf;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.apache.calcite.adapter.enumerable.NotNullImplementor;
+import org.apache.calcite.adapter.enumerable.NullPolicy;
+import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
+import org.apache.calcite.linq4j.tree.Expression;
+import org.apache.calcite.linq4j.tree.Expressions;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.opensearch.sql.calcite.utils.PPLOperandTypes;
+import org.opensearch.sql.expression.function.ImplementorUDF;
+import org.opensearch.sql.expression.function.UDFOperandMetadata;
+
+/** Custom REX_EXTRACT_MULTI function for extracting multiple regex matches. */
+public final class RexExtractMultiFunction extends ImplementorUDF {
+
+  public RexExtractMultiFunction() {
+    super(new RexExtractMultiImplementor(), NullPolicy.ARG0);
+  }
+
+  @Override
+  public SqlReturnTypeInference getReturnTypeInference() {
+    return call ->
+        call.getTypeFactory()
+            .createArrayType(call.getTypeFactory().createSqlType(SqlTypeName.VARCHAR, 2000), -1);
+  }
+
+  @Override
+  public UDFOperandMetadata getOperandMetadata() {
+    return PPLOperandTypes.STRING_STRING_INTEGER_INTEGER;
+  }
+
+  private static class RexExtractMultiImplementor implements NotNullImplementor {
+
+    @Override
+    public Expression implement(
+        RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
+      Expression field = translatedOperands.get(0);
+      Expression pattern = translatedOperands.get(1);
+      Expression groupIndex = translatedOperands.get(2);
+      Expression maxMatch = translatedOperands.get(3);
+
+      return Expressions.call(
+          RexExtractMultiFunction.class,
+          "extractMultipleGroups",
+          field,
+          pattern,
+          groupIndex,
+          maxMatch);
+    }
+  }
+
+  public static List<String> extractMultipleGroups(
+      String text, String pattern, int groupIndex, int maxMatch) {
+    if (text == null || pattern == null) {
+      return null;
+    }
+
+    try {
+      Pattern compiledPattern = Pattern.compile(pattern);
+      Matcher matcher = compiledPattern.matcher(text);
+      List<String> matches = new ArrayList<>();
+
+      int matchCount = 0;
+      while (matcher.find() && (maxMatch == 0 || matchCount < maxMatch)) {
+        if (groupIndex > 0 && groupIndex <= matcher.groupCount()) {
+          String match = matcher.group(groupIndex);
+          if (match != null) {
+            matches.add(match);
+            matchCount++;
+          }
+        }
+      }
+
+      return matches.isEmpty() ? null : matches;
+    } catch (Exception e) {
+      return null;
+    }
+  }
+}