Add more IT for all commands on map path

Signed-off-by: Chen Dai <daichen@amazon.com>

Author: dai-chen

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

@@ -7,14 +7,16 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.LinkedHashMap;
+import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import lombok.RequiredArgsConstructor;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.tree.AddTotals;
 import org.opensearch.sql.ast.tree.FillNull;
@@ -25,24 +27,20 @@
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 
 /**
- * Pre-materializes MAP dotted paths for symbol-based commands (Category A) before the command's own
- * visitor logic runs. When a command like {@code rename} or {@code fillnull} references a dotted
- * path such as {@code doc.user.name}, this class resolves it to an {@code ITEM()} expression and
- * projects it as a flat named column so the command's string-matching logic can find it.
- *
- * <p>Called from {@link CalciteRelNodeVisitor#visitChildren} after children are visited. The
- * pipeline is: {@code extractFieldOperands → resolveMapPaths → projectMapPaths}.
+ * Resolves MAP dotted paths (e.g. {@code doc.user.name}) referenced by a PPL command and projects
+ * them as flat named columns. Each dotted path that resolves to an {@code ITEM()} expression is
+ * added to the current row type so downstream command logic can reference it by name.
  */
 @RequiredArgsConstructor
 public class MapPathPreMaterializer {
 
+  private static final Logger log = LogManager.getLogger(MapPathPreMaterializer.class);
+
   /** Visitor used to resolve field expressions to Calcite {@link RexNode}. */
   private final CalciteRexNodeVisitor rexVisitor;
 
   /**
-   * Materializes MAP dotted paths referenced by the given command. For each field operand that
-   * resolves to an {@code ITEM()} access on a MAP column, projects it as a flat named column via
-   * {@code relBuilder.projectPlus()}.
+   * Resolves and projects MAP dotted paths referenced by the given command as flat named columns.
    *
    * @param plan the AST command being visited
    * @param context the current plan context with relBuilder state
@@ -54,7 +52,7 @@ public void materializePaths(UnresolvedPlan plan, CalcitePlanContext context) {
 
     List<Field> fields = extractFieldOperands(plan);
     if (!fields.isEmpty()) {
-      projectMapPaths(resolveMapPaths(fields, context), context);
+      doMaterializeMapPaths(fields, context);
     }
   }
 
@@ -63,41 +61,30 @@ private List<Field> extractFieldOperands(UnresolvedPlan node) {
       case Rename rename -> toFields(rename.getRenameList(), m -> m.getOrigin());
       case FillNull fillNull -> toFields(fillNull.getReplacementPairs(), Pair::getLeft);
       case Replace replace -> toFields(replace.getFieldList());
+      // Only fields-exclusion needs pre-materialization; inclusion resolves paths via visitProject
       case Project project -> project.isExcluded() ? toFields(project.getProjectList()) : List.of();
       case AddTotals addTotals -> toFields(addTotals.getFieldList());
       case MvCombine mvCombine -> List.of(mvCombine.getField());
       default -> List.of();
     };
   }
 
-  private Map<String, RexNode> resolveMapPaths(List<Field> fields, CalcitePlanContext context) {
-    Map<String, RexNode> paths = new LinkedHashMap<>();
+  private void doMaterializeMapPaths(List<Field> fields, CalcitePlanContext context) {
+    Set<String> existingFields =
+        new HashSet<>(context.relBuilder.peek().getRowType().getFieldNames());
+    List<RexNode> newColumns = new ArrayList<>();
     for (Field f : fields) {
       try {
         RexNode resolved = rexVisitor.analyze(f, context);
-        if (resolved.getKind() == SqlKind.ITEM) {
-          paths.put(f.getField().toString(), resolved);
+        String name = f.getField().toString();
+        if (resolved.getKind() == SqlKind.ITEM && !existingFields.contains(name)) {
+          newColumns.add(context.relBuilder.alias(resolved, name));
         }
       } catch (RuntimeException e) {
-        // Field cannot be resolved (e.g., not in schema) — skip silently
+        log.debug("Skipping field resolution for '{}': {}", f.getField(), e.getMessage(), e);
       }
     }
-    return paths;
-  }
 
-  private void projectMapPaths(Map<String, RexNode> mapPaths, CalcitePlanContext context) {
-    if (mapPaths.isEmpty()) {
-      return;
-    }
-
-    List<String> existingFields = context.relBuilder.peek().getRowType().getFieldNames();
-    List<RexNode> newColumns = new ArrayList<>();
-    mapPaths.forEach(
-        (name, itemAccess) -> {
-          if (!existingFields.contains(name)) {
-            newColumns.add(context.relBuilder.alias(itemAccess, name));
-          }
-        });
     if (!newColumns.isEmpty()) {
       context.relBuilder.projectPlus(newColumns);
     }

core/src/test/java/org/opensearch/sql/calcite/MapPathPreMaterializerTest.java

@@ -34,6 +34,7 @@
 import java.util.Set;
 import java.util.stream.Stream;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.tools.FrameworkConfig;
@@ -80,8 +81,7 @@ void setUp() {
     mockedToolsHelper
         .when(() -> CalciteToolsHelper.create(any(), any(), any()))
         .thenReturn(relBuilder);
-    when(relBuilder.getRexBuilder())
-        .thenReturn(new org.apache.calcite.rex.RexBuilder(OpenSearchTypeFactory.TYPE_FACTORY));
+    when(relBuilder.getRexBuilder()).thenReturn(new RexBuilder(OpenSearchTypeFactory.TYPE_FACTORY));
     context =
         CalcitePlanContext.create(mock(FrameworkConfig.class), SysLimit.DEFAULT, QueryType.PPL);
     lenient().when(relBuilder.size()).thenReturn(1);
@@ -129,7 +129,7 @@ void testFieldsExclusion() {
   @Test
   void testAddtotals() {
     givenMapPaths("doc.user.name")
-        .whenCommand(new AddTotals(List.of(field("doc.user.name")), java.util.Map.of()))
+        .whenCommand(new AddTotals(List.of(field("doc.user.name")), Map.of()))
         .shouldProject("doc.user.name");
   }
 
@@ -185,7 +185,7 @@ void testNoOpForNonExcludedProject() {
 
   @Test
   void testNoOpWhenRelBuilderStackEmpty() {
-    when(relBuilder.size()).thenReturn(0);
+    lenient().when(relBuilder.size()).thenReturn(0);
     givenMapPaths("doc.user.name")
         .whenCommand(rename(DUMMY_CHILD, map("doc.user.name", "u")))
         .shouldNotProject();
@@ -223,7 +223,7 @@ MapPathAssertion givenNonMapPaths(String... fieldNames) {
       for (String name : fieldNames) {
         RexNode ref = mock(RexNode.class, "ref:" + name);
         when(ref.getKind()).thenReturn(SqlKind.INPUT_REF);
-        when(rexVisitor.analyze(fieldMatching(name), eq(context))).thenReturn(ref);
+        lenient().when(rexVisitor.analyze(fieldMatching(name), eq(context))).thenReturn(ref);
       }
       return this;
     }