rex - initial implementation

RyanL1997 · RyanL1997 · commit f625a52a98eb · 2025-08-23T12:42:01.000-07:00
Signed-off-by: Jialiang Liang &lt;jiallian@amazon.com&gt;
diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java
@@ -80,6 +80,7 @@
 import org.opensearch.sql.ast.tree.RareTopN;
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.RelationSubquery;
+import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Sort;
@@ -736,6 +737,12 @@ public LogicalPlan visitReverse(Reverse node, AnalysisContext context) {
         "REVERSE is supported only when " + CALCITE_ENGINE_ENABLED.getKeyValue() + "=true");
   }
 
+  @Override
+  public LogicalPlan visitRex(Rex node, AnalysisContext context) {
+    throw new UnsupportedOperationException(
+        "REX is supported only when " + CALCITE_ENGINE_ENABLED.getKeyValue() + "=true");
+  }
+
   @Override
   public LogicalPlan visitPaginate(Paginate paginate, AnalysisContext context) {
     LogicalPlan child = paginate.getChild().get(0).accept(this, context);
diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java
@@ -67,6 +67,7 @@
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
 import org.opensearch.sql.ast.tree.Relation;
+import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.RelationSubquery;
 import org.opensearch.sql.ast.tree.Rename;
 import org.opensearch.sql.ast.tree.Reverse;
@@ -249,6 +250,9 @@ public T visitReverse(Reverse node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitRex(Rex node, C context) {
+    return visitChildren(node, context);
+  }
   public T visitLambdaFunction(LambdaFunction node, C context) {
     return visitChildren(node, context);
   }
diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Rex.java b/core/src/main/java/org/opensearch/sql/ast/tree/Rex.java
@@ -0,0 +1,61 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import java.util.Optional;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Literal;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+
+/** AST node represent Rex field extraction operation. */
+@Getter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+public class Rex extends UnresolvedPlan {
+  /** Field to extract from. */
+  private final UnresolvedExpression field;
+
+  /** Pattern with named capture groups. */
+  private final Literal pattern;
+
+  /** Maximum number of matches (optional). */
+  private final Optional<Integer> maxMatch;
+
+  /** Child Plan. */
+  @Setter private UnresolvedPlan child;
+
+  public Rex(UnresolvedExpression field, Literal pattern) {
+    this(field, pattern, Optional.empty());
+  }
+
+  public Rex(UnresolvedExpression field, Literal pattern, Optional<Integer> maxMatch) {
+    this.field = field;
+    this.pattern = pattern;
+    this.maxMatch = maxMatch;
+  }
+
+  @Override
+  public Rex attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return ImmutableList.of(child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitRex(this, context);
+  }
+}
diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java
@@ -100,6 +100,7 @@
 import org.opensearch.sql.ast.tree.RareTopN;
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.Rename;
+import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.ast.tree.SubqueryAlias;
@@ -117,6 +118,7 @@
 import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
 import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.expression.parse.RegexCommonUtils;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLFuncImpTable;
 import org.opensearch.sql.utils.ParseUtils;
@@ -171,6 +173,41 @@ public RelNode visitFilter(Filter node, CalcitePlanContext context) {
     return context.relBuilder.peek();
   }
 
+  @Override
+  public RelNode visitRex(Rex node, CalcitePlanContext context) {
+    visitChildren(node, context);
+
+    RexNode fieldRex = rexVisitor.analyze(node.getField(), context);
+    String patternStr = (String) node.getPattern().getValue();
+    List<String> namedGroups = RegexCommonUtils.getNamedGroupCandidates(patternStr);
+    
+    if (namedGroups.isEmpty()) {
+      throw new IllegalArgumentException("Rex pattern must contain at least one named capture group");
+    }
+
+    // Filter matching rows on data nodes using script pushdown
+    RexNode regexMatchCondition = context.rexBuilder.makeCall(
+        org.apache.calcite.sql.fun.SqlLibraryOperators.REGEXP_CONTAINS,
+        fieldRex,
+        context.rexBuilder.makeLiteral(patternStr)
+    );
+    context.relBuilder.filter(regexMatchCondition);
+
+    // Extract fields from filtered data
+    List<RexNode> newFields = 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));
+      newFields.add(extractCall);
+    }
+
+    projectPlusOverriding(newFields, namedGroups, context);
+    return context.relBuilder.peek();
+  }
   private boolean containsSubqueryExpression(Node expr) {
     if (expr == null) {
       return false;
diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/PPLOperandTypes.java b/core/src/main/java/org/opensearch/sql/calcite/utils/PPLOperandTypes.java
@@ -47,6 +47,8 @@ private PPLOperandTypes() {}
       UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.NUMERIC_NUMERIC);
   public static final UDFOperandMetadata STRING_INTEGER =
       UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.CHARACTER, SqlTypeFamily.INTEGER));
+  public static final UDFOperandMetadata STRING_STRING_INTEGER =
+      UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.INTEGER));
 
   public static final UDFOperandMetadata NUMERIC_NUMERIC_NUMERIC =
       UDFOperandMetadata.wrap(
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java
@@ -215,6 +215,7 @@ public enum BuiltinFunctionName {
   LTRIM(FunctionName.of("ltrim")),
   POSITION(FunctionName.of("position")),
   REGEXP(FunctionName.of("regexp")),
+  REX_EXTRACT(FunctionName.of("REX_EXTRACT")),
   REPLACE(FunctionName.of("replace")),
   REVERSE(FunctionName.of("reverse")),
   RIGHT(FunctionName.of("right")),
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java
@@ -48,6 +48,7 @@
 import org.opensearch.sql.expression.function.udf.CryptographicFunction;
 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.SpanFunction;
 import org.opensearch.sql.expression.function.udf.condition.EarliestFunction;
 import org.opensearch.sql.expression.function.udf.condition.LatestFunction;
@@ -380,6 +381,7 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
       RELEVANCE_QUERY_FUNCTION_INSTANCE.toUDF("multi_match", false);
   public static final SqlOperator NUMBER_TO_STRING =
       new NumberToStringFunction().toUDF("NUMBER_TO_STRING");
+  public static final SqlOperator REX_EXTRACT = new RexExtractFunction().toUDF("REX_EXTRACT");
 
   /**
    * Returns the PPL specific operator table, creating it if necessary.
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java b/core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java
@@ -160,6 +160,7 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.RAND;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REDUCE;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REGEXP;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.REX_EXTRACT;
 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.RIGHT;
@@ -711,6 +712,7 @@ void populate() {
       registerOperator(SIMPLE_QUERY_STRING, PPLBuiltinOperators.SIMPLE_QUERY_STRING);
       registerOperator(QUERY_STRING, PPLBuiltinOperators.QUERY_STRING);
       registerOperator(MULTI_MATCH, PPLBuiltinOperators.MULTI_MATCH);
+      registerOperator(REX_EXTRACT, PPLBuiltinOperators.REX_EXTRACT);
 
       // Register PPL Datetime UDF operator
       registerOperator(TIMESTAMP, PPLBuiltinOperators.TIMESTAMP);
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/udf/RexExtractFunction.java b/core/src/main/java/org/opensearch/sql/expression/function/udf/RexExtractFunction.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.udf;
+
+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.ReturnTypes;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.opensearch.sql.calcite.utils.PPLOperandTypes;
+import org.opensearch.sql.expression.function.ImplementorUDF;
+import org.opensearch.sql.expression.function.UDFOperandMetadata;
+
+import java.util.List;
+
+/**
+ * Custom REX_EXTRACT function for extracting regex named capture groups.
+ */
+public final class RexExtractFunction extends ImplementorUDF {
+
+  public RexExtractFunction() {
+    super(new RexExtractImplementor(), NullPolicy.ARG0);
+  }
+
+  @Override
+  public SqlReturnTypeInference getReturnTypeInference() {
+    return ReturnTypes.VARCHAR_2000_NULLABLE;
+  }
+
+  @Override
+  public UDFOperandMetadata getOperandMetadata() {
+    return PPLOperandTypes.STRING_STRING_INTEGER;
+  }
+
+  private static class RexExtractImplementor 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);
+      
+      return Expressions.call(
+          RexExtractFunction.class,
+          "extractGroup",
+          field,
+          pattern,
+          groupIndex
+      );
+    }
+  }
+  public static String extractGroup(String text, String pattern, int groupIndex) {
+    if (text == null || pattern == null) {
+      return null;
+    }
+    
+    try {
+      java.util.regex.Pattern compiledPattern = java.util.regex.Pattern.compile(pattern);
+      java.util.regex.Matcher matcher = compiledPattern.matcher(text);
+      
+      if (matcher.find() && groupIndex > 0 && groupIndex <= matcher.groupCount()) {
+        return matcher.group(groupIndex);
+      }
+      
+      return null;
+    } catch (Exception e) {
+      return null;
+    }
+  }
+}
diff --git a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 b/ppl/src/main/antlr/OpenSearchPPLLexer.g4
@@ -17,6 +17,7 @@ EXPLAIN:                            'EXPLAIN';
 FROM:                               'FROM';
 WHERE:                              'WHERE';
 FIELDS:                             'FIELDS';
+FIELD:                              'FIELD';
 TABLE:                              'TABLE';  // Alias for FIELDS command
 RENAME:                             'RENAME';
 STATS:                              'STATS';
@@ -29,6 +30,7 @@ TOP:                                'TOP';
 RARE:                               'RARE';
 PARSE:                              'PARSE';
 REGEX:                              'REGEX';
+REX:                                'REX';
 PUNCT:                              'PUNCT';
 GROK:                               'GROK';
 PATTERN:                            'PATTERN';
@@ -48,6 +50,7 @@ VARIABLE_COUNT_THRESHOLD:           'VARIABLE_COUNT_THRESHOLD';
 FREQUENCY_THRESHOLD_PERCENTAGE:     'FREQUENCY_THRESHOLD_PERCENTAGE';
 METHOD:                             'METHOD';
 MAX_SAMPLE_COUNT:                   'MAX_SAMPLE_COUNT';
+MAX_MATCH:                          'MAX_MATCH';
 BUFFER_LIMIT:                       'BUFFER_LIMIT';
 LABEL:                              'LABEL';
 AGGREGATION:                        'AGGREGATION';
diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4
@@ -71,6 +71,7 @@ commands
    | expandCommand
    | flattenCommand
    | reverseCommand
+   | rexCommand
    ;
 
 commandName
@@ -103,6 +104,7 @@ commandName
    | TRENDLINE
    | EXPLAIN
    | REVERSE
+   | REX
    ;
 
 searchCommand
@@ -187,6 +189,15 @@ parseCommand
    : PARSE (source_field = expression) (pattern = stringLiteral)
    ;
 
+rexCommand
+    : REX rexExpr
+    ;
+
+rexExpr
+    : pattern=stringLiteral
+    | FIELD EQUAL field=qualifiedName pattern=stringLiteral
+    | FIELD EQUAL field=qualifiedName pattern=stringLiteral (MAX_MATCH EQUAL maxMatch=integerLiteral)?
+    ;
 patternsMethod
    : PUNCT
    | REGEX
diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java b/ppl/src/main/java/org/opensearch/sql/ppl/parser/AstBuilder.java
@@ -79,6 +79,7 @@
 import org.opensearch.sql.ast.tree.RareTopN.CommandType;
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.Rename;
+import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Sort;
 import org.opensearch.sql.ast.tree.SubqueryAlias;
@@ -764,6 +765,21 @@ public UnresolvedPlan visitAppendcolCommand(OpenSearchPPLParser.AppendcolCommand
     return new AppendCol(override, subsearch.get());
   }
 
+  @Override
+  public UnresolvedPlan visitRexCommand(OpenSearchPPLParser.RexCommandContext ctx) {
+    UnresolvedExpression field = null;
+    Literal pattern = (Literal) internalVisitExpression(ctx.rexExpr().pattern);
+    Optional<Integer> maxMatch = Optional.empty();
+
+    if (ctx.rexExpr().field != null) {
+      field = internalVisitExpression(ctx.rexExpr().field);
+    }
+    if (ctx.rexExpr().maxMatch != null) {
+      maxMatch = Optional.of(Integer.parseInt(ctx.rexExpr().maxMatch.getText()));
+    }
+
+    return new Rex(field, pattern, maxMatch);
+  }
   /** Get original text in query. */
   private String getTextInQuery(ParserRuleContext ctx) {
     Token start = ctx.getStart();
