Minor refactoring

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

Author: ykmr1224

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

@@ -38,6 +38,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.calcite.plan.RelOptTable;
@@ -80,7 +81,6 @@
 import org.opensearch.sql.ast.expression.AllFieldsExcludeMeta;
 import org.opensearch.sql.ast.expression.Argument;
 import org.opensearch.sql.ast.expression.Argument.ArgumentMap;
-import org.opensearch.sql.ast.expression.Cast;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Function;
 import org.opensearch.sql.ast.expression.Let;
@@ -1900,7 +1900,6 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
     // 1. group the group-by list + field list and add a count() aggregation
     List<UnresolvedExpression> groupExprList = new ArrayList<>();
     node.getGroupExprList().forEach(exp -> groupExprList.add(exp));
-    // need alias for dynamic fields
     node.getFields().forEach(field -> groupExprList.add(field));
     groupExprList.forEach(expr -> projectDynamicField(expr, context));
     List<UnresolvedExpression> aggExprList =
@@ -2050,9 +2049,8 @@ public RelNode visitTimechart(
       org.opensearch.sql.ast.tree.Timechart node, CalcitePlanContext context) {
     visitChildren(node, context);
 
-    projectDynamicFieldAsString(
-        node.getBinExpression(), context); // spanExpr would become static field.
-    projectDynamicFieldAsString(node.getByField(), context); // byField would become static field.
+    projectDynamicFieldAsString(node.getBinExpression(), context);
+    projectDynamicFieldAsString(node.getByField(), context);
 
     // Extract parameters
     UnresolvedExpression spanExpr = node.getBinExpression();
@@ -2095,8 +2093,7 @@ public RelNode visitTimechart(
 
     // Extract parameters for byField case
     UnresolvedExpression byField = node.getByField();
-
-    String byFieldName = ((Field) node.getByField()).getField().toString();
+    String byFieldName = ((Field) byField).getField().toString();
     String valueFunctionName = getValueFunctionName(node.getAggregateFunction());
 
     int limit = Optional.ofNullable(node.getLimit()).orElse(10);
@@ -2150,25 +2147,34 @@ public RelNode visitTimechart(
     }
   }
 
+  /**
+   * Project dynamic field to static field and cast to string to make it easier to handle. It does
+   * nothing if exp does not refer dynamic field.
+   */
   private void projectDynamicFieldAsString(UnresolvedExpression exp, CalcitePlanContext context) {
-    projectDynamicField(exp, true, context);
+    projectDynamicField(exp, context, node -> context.rexBuilder.castToString(node));
   }
 
+  /**
+   * Project dynamic field to static field to make it easier to handle. It does nothing if exp does
+   * not refer dynamic field.
+   */
   private void projectDynamicField(UnresolvedExpression exp, CalcitePlanContext context) {
-    projectDynamicField(exp, false, context);
+    UnaryOperator<RexNode> noWrap = node -> node;
+    projectDynamicField(exp, context, noWrap);
   }
 
   private void projectDynamicField(
-      UnresolvedExpression exp, boolean castNeeded, CalcitePlanContext context) {
+      UnresolvedExpression exp, CalcitePlanContext context, UnaryOperator<RexNode> nodeWrapper) {
     if (exp != null) {
       exp.accept(
           new AbstractNodeVisitor<Void, CalcitePlanContext>() {
             @Override
             public Void visitField(Field field, CalcitePlanContext context) {
               RexNode node = rexVisitor.analyze(field, context);
               if (node.isA(SqlKind.ITEM)) {
-                RexNode casted = castNeeded ? castToString(node, context) : node;
-                RexNode alias = context.relBuilder.alias(casted, field.getField().toString());
+                RexNode alias =
+                    context.relBuilder.alias(nodeWrapper.apply(node), field.getField().toString());
                 context.relBuilder.projectPlus(alias);
               }
               return null;
@@ -2178,36 +2184,6 @@ public Void visitField(Field field, CalcitePlanContext context) {
     }
   }
 
-  private RexNode castToString(RexNode node, CalcitePlanContext context) {
-    RelDataType stringType = context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR);
-    RelDataType nullableStringType =
-        context.rexBuilder.getTypeFactory().createTypeWithNullability(stringType, true);
-    return context.rexBuilder.makeCast(nullableStringType, node, true, true);
-  }
-
-  /** Add cast and alias to Field. Needed for when the field is resolved as dynamic field access. */
-  private UnresolvedExpression castAndAliasToFieldAccess(UnresolvedExpression exp) {
-    if (exp instanceof Field) {
-      Field f = (Field) exp;
-      return AstDSL.alias(f.getField().toString(), castToString(f));
-    } else {
-      return castToString(exp);
-    }
-  }
-
-  private UnresolvedExpression castToString(UnresolvedExpression exp) {
-    return new Cast(exp, AstDSL.stringLiteral("STRING"));
-  }
-
-  /** Add alias to Field. Needed for when the field is resolved as dynamic field access. */
-  private UnresolvedExpression addAliasToFieldAccess(UnresolvedExpression exp) {
-    if (exp instanceof Field f) {
-      return AstDSL.alias(f.getField().toString(), f);
-    } else {
-      return exp;
-    }
-  }
-
   /** Build top categories query - simpler approach that works better with OTHER handling */
   private RelNode buildTopCategoriesQuery(
       RelNode completeResults, int limit, CalcitePlanContext context) {

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

@@ -163,4 +163,11 @@ else if ((SqlTypeUtil.isApproximateNumeric(sourceType) || SqlTypeUtil.isDecimal(
     }
     return super.makeCast(pos, type, exp, matchNullability, safe, format);
   }
+
+  /** Cast node to string */
+  public RexNode castToString(RexNode node) {
+    RelDataType stringType = getTypeFactory().createSqlType(SqlTypeName.VARCHAR);
+    RelDataType nullableStringType = getTypeFactory().createTypeWithNullability(stringType, true);
+    return makeCast(nullableStringType, node, true, true);
+  }
 }

core/src/main/java/org/opensearch/sql/calcite/rel/RelFieldBuilder.java

@@ -11,7 +11,6 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
-import org.apache.calcite.avatica.SqlType;
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexInputRef;
@@ -45,16 +44,6 @@ public List<String> getAllFieldNames() {
     return getAllFieldNames(0);
   }
 
-  public boolean isFieldSpecificType(String fieldName) {
-    List<RelDataTypeField> fields = relBuilder.peek().getRowType().getFieldList();
-    for (RelDataTypeField field : fields) {
-      if (field.getName().equals(fieldName)) {
-        return !field.getType().getSqlTypeName().equals(SqlType.ANY);
-      }
-    }
-    return false;
-  }
-
   public List<String> getAllFieldNames(int inputCount, int inputOrdinal) {
     return getAllFieldNames(getStackPosition(inputCount, inputOrdinal));
   }