Support aggregation/window commands with dynamic fields

common/src/main/java/org/opensearch/sql/common/utils/DebugUtils.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.common.utils;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/** Utility class for debugging operations. */
+public class DebugUtils {
+
+  private static void print(String format, Object... args) {
+    System.out.println(String.format(format, args));
+  }
+
+  public static <T> T debug(T obj, String message) {
+    print("### %s: %s (at %s)", message, stringify(obj), getCalledFrom(1));
+    return obj;
+  }
+
+  public static <T> T debug(T obj) {
+    print("### %s (at %s)", stringify(obj), getCalledFrom(1));
+    return obj;
+  }
+
+  private static String getCalledFrom(int pos) {
+    RuntimeException e = new RuntimeException();
+    StackTraceElement item = e.getStackTrace()[pos + 1];
+    return item.getClassName() + "." + item.getMethodName() + ":" + item.getLineNumber();
+  }
+
+  private static String stringify(Collection<?> items) {
+    if (items == null) {
+      return "null";
+    }
+
+    if (items.isEmpty()) {
+      return "()";
+    }
+
+    String result = items.stream().map(i -> stringify(i)).collect(Collectors.joining(","));
+
+    return "(" + result + ")";
+  }
+
+  private static String stringify(Map<?, ?> map) {
+    if (map == null) {
+      return "[[null]]";
+    }
+
+    if (map.isEmpty()) {
+      return "[[EMPTY]]";
+    }
+
+    String result =
+        map.entrySet().stream()
+            .map(entry -> entry.getKey() + ": " + stringify(entry.getValue()))
+            .collect(Collectors.joining(","));
+    return "{" + result + "}";
+  }
+
+  private static String stringify(Object obj) {
+    if (obj instanceof Collection) {
+      return stringify((Collection) obj);
+    } else if (obj instanceof Map) {
+      return stringify((Map) obj);
+    }
+    return String.valueOf(obj);
+  }
+}

common/src/main/java/org/opensearch/sql/common/utils/JsonUtils.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.common.utils;
+
+import lombok.experimental.UtilityClass;
+
+@UtilityClass
+public class JsonUtils {
+  /**
+   * Utility method to build JSON string from multiple strings with single-quotes. This is just for
+   * ease of read and maintain in tests. sjson("{", "'key': 'name'", "}") -> "{\n \"key\":
+   * \"name\"\n}"
+   *
+   * @param lines lines using single-quote instead for double-quote
+   * @return sting joined inputs and replaces single-quotes with double-quotes
+   */
+  public static String sjson(String... lines) {
+    StringBuilder builder = new StringBuilder();
+    for (String line : lines) {
+      builder.append(replaceQuote(line));
+      builder.append("\n");
+    }
+    return builder.toString();
+  }
+
+  private static String replaceQuote(String line) {
+    return line.replace("'", "\"");
+  }
+
+  /**
+   * Utility method to build multiline string from list of strings. Last line will also have new
+   * line at the end.
+   *
+   * @param lines input lines
+   * @return string contains lines
+   */
+  public static String lines(String... lines) {
+    StringBuilder builder = new StringBuilder();
+    for (String line : lines) {
+      builder.append(line);
+      builder.append("\n");
+    }
+    return builder.toString();
+  }
+}

common/src/test/java/org/opensearch/sql/common/utils/JsonUtilsTest.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.common.utils;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+
+public class JsonUtilsTest {
+
+  @Test
+  public void testSjsonWithSingleLine() {
+    String result = JsonUtils.sjson("{'key': 'value'}");
+    assertEquals("{\"key\": \"value\"}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithMultipleLines() {
+    String result = JsonUtils.sjson("{", "  'name': 'John',", "  'age': 30", "}");
+    assertEquals("{\n  \"name\": \"John\",\n  \"age\": 30\n}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithEmptyString() {
+    String result = JsonUtils.sjson("");
+    assertEquals("\n", result);
+  }
+
+  @Test
+  public void testSjsonWithNoQuotes() {
+    String result = JsonUtils.sjson("no quotes here");
+    assertEquals("no quotes here\n", result);
+  }
+
+  @Test
+  public void testSjsonWithMixedQuotes() {
+    String result = JsonUtils.sjson("'single' and \"double\" quotes");
+    assertEquals("\"single\" and \"double\" quotes\n", result);
+  }
+
+  @Test
+  public void testSjsonWithMultipleSingleQuotes() {
+    String result = JsonUtils.sjson("'key1': 'value1', 'key2': 'value2'");
+    assertEquals("\"key1\": \"value1\", \"key2\": \"value2\"\n", result);
+  }
+
+  @Test
+  public void testSjsonWithNestedJson() {
+    String result = JsonUtils.sjson("{", "  'outer': {", "    'inner': 'value'", "  }", "}");
+    assertEquals("{\n  \"outer\": {\n    \"inner\": \"value\"\n  }\n}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithArrays() {
+    String result = JsonUtils.sjson("{", "  'items': ['item1', 'item2', 'item3']", "}");
+    assertEquals("{\n  \"items\": [\"item1\", \"item2\", \"item3\"]\n}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithSpecialCharacters() {
+    String result = JsonUtils.sjson("{'key': 'value with \\'escaped\\' quotes'}");
+    assertEquals("{\"key\": \"value with \\\"escaped\\\" quotes\"}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithEmptyArray() {
+    String result = JsonUtils.sjson();
+    assertEquals("", result);
+  }
+
+  @Test
+  public void testSjsonWithNullValues() {
+    String result = JsonUtils.sjson("{'key': null}");
+    assertEquals("{\"key\": null}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithNumbers() {
+    String result =
+        JsonUtils.sjson("{", "  'integer': 42,", "  'float': 3.14,", "  'negative': -10", "}");
+    assertEquals("{\n  \"integer\": 42,\n  \"float\": 3.14,\n  \"negative\": -10\n}\n", result);
+  }
+
+  @Test
+  public void testSjsonWithBooleans() {
+    String result = JsonUtils.sjson("{'active': true, 'deleted': false}");
+    assertEquals("{\"active\": true, \"deleted\": false}\n", result);
+  }
+
+  @Test
+  public void testSjsonPreservesWhitespace() {
+    String result = JsonUtils.sjson("  {'key':  'value'}  ");
+    assertEquals("  {\"key\":  \"value\"}  \n", result);
+  }
+
+  @Test
+  public void testSjsonWithComplexJson() {
+    String result =
+        JsonUtils.sjson(
+            "{",
+            "  'user': {",
+            "    'name': 'Alice',",
+            "    'email': '[email protected]',",
+            "    'roles': ['admin', 'user'],",
+            "    'active': true,",
+            "    'loginCount': 42",
+            "  }",
+            "}");
+    String expected =
+        "{\n"
+            + "  \"user\": {\n"
+            + "    \"name\": \"Alice\",\n"
+            + "    \"email\": \"[email protected]\",\n"
+            + "    \"roles\": [\"admin\", \"user\"],\n"
+            + "    \"active\": true,\n"
+            + "    \"loginCount\": 42\n"
+            + "  }\n"
+            + "}\n";
+    assertEquals(expected, result);
+  }
+}

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;
@@ -1897,10 +1898,10 @@ 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());
-    List<UnresolvedExpression> fieldList =
-        node.getFields().stream().map(f -> (UnresolvedExpression) f).toList();
-    groupExprList.addAll(fieldList);
+    List<UnresolvedExpression> groupExprList = new ArrayList<>();
+    node.getGroupExprList().forEach(exp -> groupExprList.add(exp));
+    node.getFields().forEach(field -> groupExprList.add(field));
+    groupExprList.forEach(expr -> projectDynamicField(expr, context));
     List<UnresolvedExpression> aggExprList =
         List.of(AstDSL.alias(countFieldName, AstDSL.aggregate("count", null)));
     aggregateWithTrimming(groupExprList, aggExprList, context);
@@ -2048,6 +2049,9 @@ public RelNode visitTimechart(
       org.opensearch.sql.ast.tree.Timechart node, CalcitePlanContext context) {
     visitChildren(node, context);
 
+    projectDynamicFieldAsString(node.getBinExpression(), context);
+    projectDynamicFieldAsString(node.getByField(), context);
+
     // Extract parameters
     UnresolvedExpression spanExpr = node.getBinExpression();
 
@@ -2111,9 +2115,7 @@ public RelNode visitTimechart(
       List<RexNode> outputFields = context.fieldBuilder.staticFields();
       List<RexNode> reordered = new ArrayList<>();
       reordered.add(context.fieldBuilder.staticField("@timestamp")); // timestamp first
-      reordered.add(
-          context.fieldBuilder.staticField(
-              byFieldName)); // byField second. TODO: allow dynamic fields
+      reordered.add(context.fieldBuilder.staticField(byFieldName)); // byField second.
       reordered.add(outputFields.get(outputFields.size() - 1)); // value function last
       context.relBuilder.project(reordered);
 
@@ -2145,6 +2147,43 @@ 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, 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) {
+    UnaryOperator<RexNode> noWrap = node -> node;
+    projectDynamicField(exp, context, noWrap);
+  }
+
+  private void projectDynamicField(
+      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 alias =
+                    context.relBuilder.alias(nodeWrapper.apply(node), field.getField().toString());
+                context.relBuilder.projectPlus(alias);
+              }
+              return null;
+            }
+          },
+          context);
+    }
+  }
+
   /** Build top categories query - simpler approach that works better with OTHER handling */
   private RelNode buildTopCategoriesQuery(
       RelNode completeResults, int limit, CalcitePlanContext context) {
@@ -2349,6 +2388,7 @@ public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
         .ifPresent(
             sortField -> {
               SortOption sortOption = analyzeSortOption(sortField.getFieldArgs());
+              projectDynamicFieldAsString(sortField, context);
               RexNode field = rexVisitor.analyze(sortField, context);
               if (sortOption == DEFAULT_DESC) {
                 context.relBuilder.sort(context.relBuilder.desc(field));
@@ -2362,7 +2402,9 @@ public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
     node.getComputations()
         .forEach(
             trendlineComputation -> {
+              projectDynamicField(trendlineComputation.getDataField(), context);
               RexNode field = rexVisitor.analyze(trendlineComputation.getDataField(), context);
+
               context.relBuilder.filter(context.relBuilder.isNotNull(field));
 
               WindowFrame windowFrame =

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/utils/PlanUtils.java

@@ -182,15 +182,15 @@ static RexNode makeOver(
         return variance(context, field, partitions, rows, lowerBound, upperBound, false, false);
       case ROW_NUMBER:
         return withOver(
-            context.relBuilder.aggregateCall(SqlStdOperatorTable.ROW_NUMBER),
+            makeAggCall(context, functionName, false, null, List.of()),
             partitions,
             orderKeys,
             true,
             lowerBound,
             upperBound);
       case NTH_VALUE:
         return withOver(
-            context.relBuilder.aggregateCall(SqlStdOperatorTable.NTH_VALUE, field, argList.get(0)),
+            makeAggCall(context, functionName, false, field, argList.subList(0, 1)),
             partitions,
             orderKeys,
             true,
@@ -215,7 +215,12 @@ private static RexNode sumOver(
       RexWindowBound lowerBound,
       RexWindowBound upperBound) {
     return withOver(
-        ctx.relBuilder.sum(operation), partitions, List.of(), rows, lowerBound, upperBound);
+        makeAggCall(ctx, BuiltinFunctionName.SUM, false, operation, List.of()),
+        partitions,
+        List.of(),
+        rows,
+        lowerBound,
+        upperBound);
   }
 
   private static RexNode countOver(