Introduce RelBuilder wrapper for dynamic fields

Signed-off-by: Tomoyuki Morita <[email protected]>

Author: ykmr1224

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

@@ -22,6 +22,8 @@
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.tools.RelBuilder;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.calcite.rel.RelBuilderWrapper;
+import org.opensearch.sql.calcite.rel.RelFieldBuilder;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.executor.QueryType;
@@ -31,7 +33,8 @@ public class CalcitePlanContext {
 
   public FrameworkConfig config;
   public final Connection connection;
-  public final RelBuilder relBuilder;
+  public final RelBuilderWrapper relBuilder;
+  public final RelFieldBuilder fieldBuilder;
   public final ExtendedRexBuilder rexBuilder;
   public final FunctionProperties functionProperties;
   public final QueryType queryType;
@@ -66,8 +69,10 @@ private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType
     this.sysLimit = sysLimit;
     this.queryType = queryType;
     this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
-    this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
-    this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
+    RelBuilder rawRelBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
+    this.relBuilder = new RelBuilderWrapper(rawRelBuilder);
+    this.rexBuilder = new ExtendedRexBuilder(rawRelBuilder.getRexBuilder());
+    this.fieldBuilder = new RelFieldBuilder(rawRelBuilder, this.rexBuilder);
     this.functionProperties = new FunctionProperties(QueryType.PPL);
     this.rexLambdaRefMap = new HashMap<>();
   }

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

@@ -72,6 +72,7 @@
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
 import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
+import org.opensearch.sql.calcite.rel.QualifiedNameResolver;
 import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
 import org.opensearch.sql.calcite.utils.PlanUtils;
 import org.opensearch.sql.calcite.utils.SubsearchUtils;
@@ -548,7 +549,7 @@ private RelNode resolveSubqueryPlan(
                 replacement,
                 context.relBuilder.literal(context.sysLimit.subsearchLimit()));
       }
-      PlanUtils.replaceTop(context.relBuilder, replacement);
+      PlanUtils.replaceTop(context.relBuilder.getRawRelBuilder(), replacement);
     }
     // pop the inner plan
     RelNode subqueryRel = context.relBuilder.build();

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

@@ -64,7 +64,11 @@ public static List<RelNode> buildUnifiedSchemaWithConflictResolution(
       List<RelNode> renamedNodes = new ArrayList<>();
       for (RelNode node : projectedNodes) {
         RelNode renamedNode =
-            context.relBuilder.push(node).project(context.relBuilder.fields(), uniqueNames).build();
+            context
+                .relBuilder
+                .push(node)
+                .project(context.fieldBuilder.staticFields(), uniqueNames)
+                .build();
         renamedNodes.add(renamedNode);
       }
       return renamedNodes;

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

@@ -34,6 +34,8 @@
 import org.apache.calcite.util.mapping.Mappings;
 import org.apache.commons.lang3.tuple.Pair;
 import org.immutables.value.Value;
+import org.opensearch.sql.calcite.rel.RelBuilderWrapper;
+import org.opensearch.sql.calcite.rel.RelFieldBuilder;
 
 /**
  * Planner rule that converts specific aggCall to a more efficient expressions, which includes:
@@ -74,8 +76,10 @@ public void onMatch(RelOptRuleCall call) {
 
   public void apply(RelOptRuleCall call, LogicalAggregate aggregate, LogicalProject project) {
 
-    final RelBuilder relBuilder = call.builder();
+    final RelBuilder rawRelBuilder = call.builder();
+    final RelBuilderWrapper relBuilder = new RelBuilderWrapper(rawRelBuilder);
     final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
+    final RelFieldBuilder fieldBuilder = new RelFieldBuilder(rawRelBuilder, rexBuilder);
     relBuilder.push(project.getInput());
 
     /*
@@ -213,13 +217,14 @@ public void apply(RelOptRuleCall call, LogicalAggregate aggregate, LogicalProjec
         distinctAggregateCalls.add(aggregateCall.transform(targetMapping));
       }
       // Project the used fields
-      relBuilder.project(relBuilder.fields(fieldsUsed.stream().toList()));
+      relBuilder.project(fieldBuilder.staticFields(fieldsUsed.stream().toList()));
     }
 
     /* Build the final project-aggregate-project after eliminating unused fields */
     relBuilder.aggregate(relBuilder.groupKey(newGroupSet, newGroupSets), distinctAggregateCalls);
     List<RexNode> parentProjects =
-        new ArrayList<>(relBuilder.fields(IntStream.range(0, groupSetOffset).boxed().toList()));
+        new ArrayList<>(
+            fieldBuilder.staticFields(IntStream.range(0, groupSetOffset).boxed().toList()));
     parentProjects.addAll(
         newExprOnAggCall.transform(
             (constructor, name) ->
@@ -267,7 +272,7 @@ private static RexNode convertToNewOperand(
     }
   }
 
-  private RexNode aliasMaybe(RelBuilder builder, RexNode node, String alias) {
+  private RexNode aliasMaybe(RelBuilderWrapper builder, RexNode node, String alias) {
     return alias == null ? node : builder.alias(node, alias);
   }
 

core/src/main/java/org/opensearch/sql/calcite/plan/PPLAggregateConvertRule.java

@@ -3,7 +3,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.calcite;
+package org.opensearch.sql.calcite.rel;
+
+import static org.opensearch.sql.calcite.plan.DynamicFieldsConstants.DYNAMIC_FIELDS_MAP;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -16,7 +18,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.sql.ast.expression.QualifiedName;
-import org.opensearch.sql.calcite.plan.DynamicFieldsConstants;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLFuncImpTable;
 
@@ -29,6 +31,25 @@ public class QualifiedNameResolver {
 
   private static final Logger log = LogManager.getLogger(QualifiedNameResolver.class);
 
+  /** Resolve field in a specific input */
+  public static Optional<RexNode> resolveField(
+      int inputCount, int inputOrdinal, String fieldName, CalcitePlanContext context) {
+    List<String> inputFieldNames = context.fieldBuilder.getAllFieldNames(inputCount, inputOrdinal);
+    if (inputFieldNames.contains(fieldName)) {
+      return Optional.of(context.fieldBuilder.staticField(inputCount, inputOrdinal, fieldName));
+    } else if (context.fieldBuilder.isDynamicFieldsExist()) {
+      return Optional.of(context.fieldBuilder.dynamicField(fieldName));
+    }
+    return Optional.empty();
+  }
+
+  public static RexNode resolveFieldOrThrow(
+      int inputCount, int inputOrdinal, String fieldName, CalcitePlanContext context) {
+    return resolveField(inputCount, inputOrdinal, fieldName, context)
+        .orElseThrow(
+            () -> new IllegalArgumentException(String.format("Field [%s] not found.", fieldName)));
+  }
+
   /**
    * Resolves a qualified name to a RexNode based on the current context.
    *
@@ -130,10 +151,9 @@ private static Optional<RexNode> resolveDynamicFields(
     List<Set<String>> inputFieldNames = collectInputFieldNames(context, inputCount);
 
     for (int i = 0; i < inputCount; i++) {
-      if (inputFieldNames.get(i).contains(DynamicFieldsConstants.DYNAMIC_FIELDS_MAP)) {
+      if (inputFieldNames.get(i).contains(DYNAMIC_FIELDS_MAP)) {
         String fieldName = String.join(".", parts);
-        RexNode dynamicField =
-            context.relBuilder.field(inputCount, i, DynamicFieldsConstants.DYNAMIC_FIELDS_MAP);
+        RexNode dynamicField = context.relBuilder.field_(inputCount, i, DYNAMIC_FIELDS_MAP);
         RexNode itemAccess = createItemAccess(dynamicField, fieldName, context);
         return Optional.of(itemAccess);
       }
@@ -149,7 +169,7 @@ private static Optional<RexNode> tryToResolveField(
         fieldName,
         inputCount);
     try {
-      return Optional.of(context.relBuilder.field(inputCount, alias, fieldName));
+      return Optional.of(context.relBuilder.field_(inputCount, alias, fieldName));
     } catch (IllegalArgumentException e) {
       log.debug("tryToResolveField() failed: {}", e.getMessage());
     }
@@ -171,7 +191,7 @@ private static Optional<RexNode> resolveFieldWithoutAlias(
       int foundInput = findInputContainingFieldName(inputCount, inputFieldNames, fieldName);
       log.debug("resolveFieldWithoutAlias() foundInput={}", foundInput);
       if (foundInput != -1) {
-        RexNode fieldNode = context.relBuilder.field(inputCount, foundInput, fieldName);
+        RexNode fieldNode = context.relBuilder.field_(inputCount, foundInput, fieldName);
         return Optional.of(resolveFieldAccess(context, parts, 0, length, fieldNode));
       }
     }
@@ -219,7 +239,7 @@ private static Optional<RexNode> resolveRenamedField(
       String alias = parts.get(0);
       for (String candidate : candidates) {
         try {
-          return Optional.of(context.relBuilder.field(alias, candidate));
+          return Optional.of(context.relBuilder.field_(alias, candidate));
         } catch (IllegalArgumentException e1) {
           // Indicates the field was not found.
         }
@@ -260,7 +280,7 @@ private static Optional<RexNode> resolveCorrelationField(
                   String fieldName = joinParts(parts, start, length);
                   log.debug("resolveCorrelationField() trying fieldName={}", fieldName);
                   if (fieldNameList.contains(fieldName)) {
-                    RexNode field = context.relBuilder.field(correlation, fieldName);
+                    RexNode field = context.relBuilder.field_(correlation, fieldName);
                     return resolveFieldAccess(context, parts, start, length, field);
                   }
                 }