Check what a flat_object allows while the plan is analyzed, as a list of allowed uses

The check walked the finished plan and rejected every node that was not a
read. Nothing else in this package works that way -- MapPathPreMaterializer,
SchemaUnifier, ForeachPlanner and QualifiedNameResolver are all helpers the
visitor calls while building -- and a plan pass pays for it: it has to trace
a column back through a projection to find the leaf it came from, and rebuild
"attributes.duration_ms" out of a generated $f2, because Calcite renames
columns and rebuilds expressions.

The rule is now a list of the uses a flat_object allows, checked where a
field reference resolves. Every strategy in QualifiedNameResolver converges
on one point, so a command not on the list is refused without an entry of
its own -- including one added later -- and the name in the message is the
one the query wrote. Reading a field or a leaf, optionally cast, and typeof
are the list; a column a read went into is remembered so a later command
using it is refused too, named by the leaf rather than by the alias.

Two things had to be arranged around this. The materializer resolves fields
inside catch (RuntimeException | AssertionError), so it runs with the read
permission open rather than having a refusal swallowed there and the field
quietly ignored. And expand already refuses anything that is not a plain
column, in its own words, so its resolution is a read and that check decides.

Behaviour is unchanged over 22 queries covering every allowed and refused
form, with one message corrected: dedup said "Cannot filter by", the name of
what it lowers to, and now says "Cannot dedup by".

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

Author: RyanL1997

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

@@ -10,6 +10,7 @@
 import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -106,6 +107,19 @@ public class CalcitePlanContext {
    */
   @Getter private final Set<String> flatObjectFields = new LinkedHashSet<>();
 
+  /**
+   * Columns a command read a {@code flat_object} field or one of its leaves into, by the name the
+   * command gave them. A later command using such a column is refused the same way, since pushdown
+   * fuses the projection back into whatever reads it.
+   */
+  @Getter private final Map<String, String> flatObjectLeafColumns = new LinkedHashMap<>();
+
+  /** Whether the expression being resolved is one of the reads a flat_object field allows. */
+  @Getter @Setter private boolean flatObjectReadAllowed = false;
+
+  /** What the command being analyzed would do with a flat_object field, for the refusal message. */
+  @Getter @Setter private String flatObjectUse = "use";
+
   private final Stack<RexCorrelVariable> correlVar = new Stack<>();
   private final Stack<List<RexNode>> windowPartitions = new Stack<>();
 

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

@@ -181,6 +181,7 @@
 import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
 import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType;
 import org.opensearch.sql.calcite.utils.BinUtils;
+import org.opensearch.sql.calcite.utils.FlatObjectScope;
 import org.opensearch.sql.calcite.utils.JoinAndLookupUtils;
 import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
 import org.opensearch.sql.calcite.utils.PPLHintUtils;
@@ -240,15 +241,18 @@ public RelNode analyze(UnresolvedPlan unresolved, CalcitePlanContext context) {
     RelNode plan = unresolved.accept(this, context);
     // Here rather than in QueryService.analyze, so the unified query pipeline -- which calls this
     // method directly -- is held to the same rules.
-    FlatObjectScopeValidator.validate(plan, context.getFlatObjectFields());
     return plan;
   }
 
   @Override
   public RelNode visitChildren(Node node, CalcitePlanContext context) {
     RelNode result = super.visitChildren(node, context);
     if (node instanceof UnresolvedPlan plan) {
-      mapPathMaterializer.materializePaths(plan, context);
+      // Materializing a dotted path is a read, so it is allowed -- and it has to be, because the
+      // materializer wraps field resolution in catch (RuntimeException | AssertionError), which
+      // would swallow a refusal raised there and leave the command quietly ignoring the field.
+      // What the command then does with the materialized column is checked where it resolves it.
+      FlatObjectScope.runAsRead(context, () -> mapPathMaterializer.materializePaths(plan, context));
     }
     return result;
   }
@@ -369,7 +373,9 @@ public RelNode visitFilter(Filter node, CalcitePlanContext context) {
       context.relBuilder.variable(v::set);
       context.pushCorrelVar(v.get());
     }
-    RexNode condition = rexVisitor.analyze(node.getCondition(), context);
+    RexNode condition =
+        FlatObjectScope.as(
+            "filter by", context, () -> rexVisitor.analyze(node.getCondition(), context));
     if (containsSubqueryExpression) {
       context.relBuilder.filter(ImmutableList.of(v.get().id), condition);
       context.popCorrelVar();
@@ -612,7 +618,9 @@ private List<RexNode> expandProjectFields(
             }
             matchingFields.forEach(f -> expandedFields.add(context.relBuilder.field(f)));
           } else if (addedFields.add(fieldName)) {
-            RexNode resolved = rexVisitor.analyze(field, context);
+            RexNode resolved =
+                FlatObjectScope.allowRead(
+                    field, fieldName, context, () -> rexVisitor.analyze(field, context));
             /*
              * Dotted path access is resolved into ITEM(map, path) function call without aliasing.
              * Re-apply the alias so the projected column retains the user-visible name.
@@ -801,6 +809,7 @@ public RelNode visitRename(Rename node, CalcitePlanContext context) {
         String newName =
             WildcardRenameUtils.applyWildcardTransformation(
                 sourcePattern, targetPattern, fieldName);
+        FlatObjectScope.trackRename(fieldName, newName, context);
         if (newNames.contains(newName) && !newName.equals(fieldName)) {
           removeFieldIfExists(newName, newNames, context);
         }
@@ -829,23 +838,27 @@ private void removeFieldIfExists(
   public RelNode visitSort(org.opensearch.sql.ast.tree.Sort node, CalcitePlanContext context) {
     visitChildren(node, context);
     List<RexNode> sortList =
-        node.getSortList().stream()
-            .map(
-                expr -> {
-                  RexNode sortField = rexVisitor.analyze(expr, context);
-                  SortOption sortOption = analyzeSortOption(expr.getFieldArgs());
-                  // Default is ASC
-                  if (sortOption.getSortOrder() == DESC) {
-                    sortField = context.relBuilder.desc(sortField);
-                  }
-                  if (sortOption.getNullOrder() == NULL_LAST) {
-                    sortField = context.relBuilder.nullsLast(sortField);
-                  } else {
-                    sortField = context.relBuilder.nullsFirst(sortField);
-                  }
-                  return sortField;
-                })
-            .collect(Collectors.toList());
+        FlatObjectScope.as(
+            "sort by",
+            context,
+            () ->
+                node.getSortList().stream()
+                    .map(
+                        expr -> {
+                          RexNode sortField = rexVisitor.analyze(expr, context);
+                          SortOption sortOption = analyzeSortOption(expr.getFieldArgs());
+                          // Default is ASC
+                          if (sortOption.getSortOrder() == DESC) {
+                            sortField = context.relBuilder.desc(sortField);
+                          }
+                          if (sortOption.getNullOrder() == NULL_LAST) {
+                            sortField = context.relBuilder.nullsLast(sortField);
+                          } else {
+                            sortField = context.relBuilder.nullsFirst(sortField);
+                          }
+                          return sortField;
+                        })
+                    .collect(Collectors.toList()));
     context.relBuilder.sort(sortList);
     // Apply count parameter as limit
     if (node.getCount() != 0) {
@@ -1292,7 +1305,12 @@ public RelNode visitEval(Eval node, CalcitePlanContext context) {
                 context.relBuilder.variable(v::set);
                 context.pushCorrelVar(v.get());
               }
-              RexNode eval = rexVisitor.analyze(expr, context);
+              RexNode eval =
+                  FlatObjectScope.allowRead(
+                      expr,
+                      expr.getVar().getField().toString(),
+                      context,
+                      () -> rexVisitor.analyze(expr, context));
               if (containsSubqueryExpression) {
                 // RelBuilder.projectPlus doesn't have a parameter with variablesSet:
                 // projectPlus(Iterable<CorrelationId> variablesSet, RexNode... nodes)
@@ -1802,7 +1820,8 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     if (!bucketNullable) {
       nonNullGroupMask.set(0, nGroup);
     }
-    visitAggregation(node, context, nonNullGroupMask, true, false);
+    FlatObjectScope.runAs(
+        "stats by", context, () -> visitAggregation(node, context, nonNullGroupMask, true, false));
     return context.relBuilder.peek();
   }
 
@@ -2273,6 +2292,7 @@ public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
 
   @Override
   public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
+    context.setFlatObjectUse("dedup by");
     visitChildren(node, context);
     List<Argument> options = node.getOptions();
     Integer allowedDuplication = (Integer) options.get(0).getValue().getValue();
@@ -4484,7 +4504,10 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
 
     // 2. Get the field to expand and an optional alias.
     Field arrayField = expand.getField();
-    RexNode resolved = rexVisitor.analyze(arrayField, context);
+    // Resolve as a read: the check below already refuses a flat_object leaf, by the same rule it
+    // uses for an object leaf or an expression, and says so in expand's own words.
+    RexNode resolved =
+        FlatObjectScope.allowReadOf(context, () -> rexVisitor.analyze(arrayField, context));
     if (!(resolved instanceof RexInputRef arrayFieldRex)) {
       // Anything that is not a column of the input -- a leaf of an object or a flat_object field,
       // an expression -- has no column to correlate the expansion with. Raised the way the