Fix result order of parse with other run time fields (opensearch-proj…

…ect#934)

Signed-off-by: Joshua Li <[email protected]>

core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java

@@ -10,6 +10,7 @@
 import com.google.common.collect.ImmutableMap.Builder;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
@@ -55,32 +56,33 @@ public boolean hasNext() {
   public ExprValue next() {
     ExprValue inputValue = input.next();
     ImmutableMap.Builder<String, ExprValue> mapBuilder = new Builder<>();
+
+    // ParseExpression will always override NamedExpression when identifier conflicts
+    // TODO needs a better implementation, see https://github.com/opensearch-project/sql/issues/458
     for (NamedExpression expr : projectList) {
       ExprValue exprValue = expr.valueOf(inputValue.bindingTuples());
-      if (namedParseExpressions.stream()
-          .noneMatch(parsed -> parsed.getNameOrAlias().equals(expr.getNameOrAlias()))) {
+      Optional<NamedExpression> optionalParseExpression = namedParseExpressions.stream()
+          .filter(parseExpr -> parseExpr.getNameOrAlias().equals(expr.getNameOrAlias()))
+          .findFirst();
+      if (optionalParseExpression.isEmpty()) {
         mapBuilder.put(expr.getNameOrAlias(), exprValue);
-      }
-    }
-    // ParseExpression will always override NamedExpression when identifier conflicts
-    // TODO needs a better implementation, see https://github.com/opensearch-project/sql/issues/458
-    for (NamedExpression expr : namedParseExpressions) {
-      if (projectList.stream()
-          .noneMatch(field -> field.getNameOrAlias().equals(expr.getNameOrAlias()))) {
         continue;
       }
+
+      NamedExpression parseExpression = optionalParseExpression.get();
       ExprValue sourceFieldValue = inputValue.bindingTuples()
-          .resolve(((ParseExpression) expr.getDelegated()).getSourceField());
+          .resolve(((ParseExpression) parseExpression.getDelegated()).getSourceField());
       if (sourceFieldValue.isMissing()) {
         // source field will be missing after stats command, read from inputValue if it exists
         // otherwise do nothing since it should not appear as a field
-        ExprValue exprValue = ExprValueUtils.getTupleValue(inputValue).get(expr.getNameOrAlias());
-        if (exprValue != null) {
-          mapBuilder.put(expr.getNameOrAlias(), exprValue);
+        ExprValue tupleValue =
+            ExprValueUtils.getTupleValue(inputValue).get(parseExpression.getNameOrAlias());
+        if (tupleValue != null) {
+          mapBuilder.put(parseExpression.getNameOrAlias(), tupleValue);
         }
       } else {
-        ExprValue parsedValue = expr.valueOf(inputValue.bindingTuples());
-        mapBuilder.put(expr.getNameOrAlias(), parsedValue);
+        ExprValue parsedValue = parseExpression.valueOf(inputValue.bindingTuples());
+        mapBuilder.put(parseExpression.getNameOrAlias(), parsedValue);
       }
     }
     return ExprTupleValue.fromExprValueMap(mapBuilder.build());

core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java

@@ -118,7 +118,26 @@ public void project_fields_with_parse_expressions() {
                     DSL.literal("action"))), DSL.named("response",
                 DSL.regex(DSL.ref("response", STRING),
                     DSL.literal("(?<action>\\w+) (?<response>\\d+)"),
-                    DSL.literal("response"))), DSL.named("ignored",
+                    DSL.literal("response"))))
+        );
+    List<ExprValue> result = execute(plan);
+
+    assertThat(
+        result,
+        allOf(
+            iterableWithSize(1),
+            hasItems(
+                ExprValueUtils.tupleValue(ImmutableMap.of("action", "GET", "response", "200")))));
+  }
+
+  @Test
+  public void project_fields_with_unused_parse_expressions() {
+    when(inputPlan.hasNext()).thenReturn(true, false);
+    when(inputPlan.next())
+        .thenReturn(ExprValueUtils.tupleValue(ImmutableMap.of("response", "GET 200")));
+    PhysicalPlan plan =
+        project(inputPlan, ImmutableList.of(DSL.named("response", DSL.ref("response", STRING))),
+            ImmutableList.of(DSL.named("ignored",
                 DSL.regex(DSL.ref("response", STRING),
                     DSL.literal("(?<action>\\w+) (?<ignored>\\d+)"),
                     DSL.literal("ignored"))))
@@ -130,7 +149,33 @@ public void project_fields_with_parse_expressions() {
         allOf(
             iterableWithSize(1),
             hasItems(
-                ExprValueUtils.tupleValue(ImmutableMap.of("action", "GET", "response", "200")))));
+                ExprValueUtils.tupleValue(ImmutableMap.of("response", "GET 200")))));
+  }
+
+  @Test
+  public void project_fields_with_parse_expressions_and_runtime_fields() {
+    when(inputPlan.hasNext()).thenReturn(true, false);
+    when(inputPlan.next())
+        .thenReturn(
+            ExprValueUtils.tupleValue(ImmutableMap.of("response", "GET 200", "eval_field", 1)));
+    PhysicalPlan plan =
+        project(inputPlan, ImmutableList.of(DSL.named("response", DSL.ref("response", STRING)),
+                DSL.named("action", DSL.ref("action", STRING)),
+                DSL.named("eval_field", DSL.ref("eval_field", INTEGER))),
+            ImmutableList.of(DSL.named("action",
+                DSL.regex(DSL.ref("response", STRING),
+                    DSL.literal("(?<action>\\w+) (?<response>\\d+)"),
+                    DSL.literal("action"))))
+        );
+    List<ExprValue> result = execute(plan);
+
+    assertThat(
+        result,
+        allOf(
+            iterableWithSize(1),
+            hasItems(
+                ExprValueUtils.tupleValue(
+                    ImmutableMap.of("response", "GET 200", "action", "GET", "eval_field", 1)))));
   }
 
   @Test

integ-test/src/test/java/org/opensearch/sql/ppl/ParseCommandIT.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.ppl;
+
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
+import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.verifyOrder;
+
+import java.io.IOException;
+import org.json.JSONObject;
+import org.junit.Test;
+
+public class ParseCommandIT extends PPLIntegTestCase {
+
+  @Override
+  public void init() throws IOException {
+    loadIndex(Index.BANK);
+  }
+
+  @Test
+  public void testParseCommand() throws IOException {
+    JSONObject result = executeQuery(
+        String.format("source=%s | parse email '.+@(?<host>.+)' | fields email, host",
+            TEST_INDEX_BANK));
+    verifyOrder(
+        result,
+        rows("[email protected]", "pyrami.com"),
+        rows("[email protected]", "netagy.com"),
+        rows("[email protected]", "quility.com"),
+        rows("[email protected]", "boink.com"),
+        rows("[email protected]", "scentric.com"),
+        rows("[email protected]", "filodyne.com"),
+        rows("[email protected]", "quailcom.com"));
+  }
+
+  @Test
+  public void testParseCommandReplaceOriginalField() throws IOException {
+    JSONObject result = executeQuery(
+        String.format("source=%s | parse email '.+@(?<email>.+)' | fields email", TEST_INDEX_BANK));
+    verifyOrder(
+        result,
+        rows("pyrami.com"),
+        rows("netagy.com"),
+        rows("quility.com"),
+        rows("boink.com"),
+        rows("scentric.com"),
+        rows("filodyne.com"),
+        rows("quailcom.com"));
+  }
+
+  @Test
+  public void testParseCommandWithOtherRunTimeFields() throws IOException {
+    JSONObject result = executeQuery(String.format("source=%s | parse email '.+@(?<host>.+)' | "
+        + "eval eval_result=1 | fields host, eval_result", TEST_INDEX_BANK));
+    verifyOrder(
+        result,
+        rows("pyrami.com", 1),
+        rows("netagy.com", 1),
+        rows("quility.com", 1),
+        rows("boink.com", 1),
+        rows("scentric.com", 1),
+        rows("filodyne.com", 1),
+        rows("quailcom.com", 1));
+  }
+}