Make basic aggregation working (part 2) : adding tests (#3355)

* Add more aggregation tests Signed-off-by: Lantao Jin <ltjin@amazon.com> * delete unrelavant code Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com>
opensearch-project · Feb 27, 2025 · c2fe0ad · c2fe0ad
1 parent db22f1c
commit c2fe0ad
Show file tree

Hide file tree

Showing 10 changed files with 605 additions and 103 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java b/core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java
@@ -5,10 +5,11 @@
 
 package org.opensearch.sql.calcite;
 
+import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
+
 import java.sql.Connection;
 import java.util.function.BiFunction;
 import lombok.Getter;
-import org.apache.calcite.adapter.java.JavaTypeFactory;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.calcite.tools.RelBuilder;
@@ -24,10 +25,10 @@ public class CalcitePlanContext {
 
   @Getter private boolean isResolvingJoinCondition = false;
 
-  private CalcitePlanContext(FrameworkConfig config, JavaTypeFactory typeFactory) {
+  private CalcitePlanContext(FrameworkConfig config) {
     this.config = config;
-    this.connection = CalciteToolsHelper.connect(config, typeFactory);
-    this.relBuilder = CalciteToolsHelper.create(config, typeFactory, connection);
+    this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
+    this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
     this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
   }
 
@@ -41,10 +42,6 @@ public RexNode resolveJoinCondition(
   }
 
   public static CalcitePlanContext create(FrameworkConfig config) {
-    return new CalcitePlanContext(config, null);
-  }
-
-  public static CalcitePlanContext create(FrameworkConfig config, JavaTypeFactory typeFactory) {
-    return new CalcitePlanContext(config, typeFactory);
+    return new CalcitePlanContext(config);
   }
 }
diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java
@@ -196,12 +196,6 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
       groupByList.add(spanRex);
       // add span's group alias field (most recent added expression)
     }
-    //        List<RexNode> aggList = node.getAggExprList().stream()
-    //            .map(expr -> rexVisitor.analyze(expr, context))
-    //            .collect(Collectors.toList());
-    //        relBuilder.aggregate(relBuilder.groupKey(groupByList),
-    //            aggList.stream().map(rex -> (MyAggregateCall) rex)
-    //                .map(MyAggregateCall::getCall).collect(Collectors.toList()));
     context.relBuilder.aggregate(context.relBuilder.groupKey(groupByList), aggList);
     return context.relBuilder.peek();
   }

diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java
@@ -18,7 +18,7 @@
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlIntervalQualifier;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
-import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.calcite.sql.parser.SqlParserUtil;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.DateString;
 import org.apache.calcite.util.TimeString;
@@ -39,7 +39,6 @@
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.expression.Xor;
 import org.opensearch.sql.calcite.utils.BuiltinFunctionUtils;
-import org.opensearch.sql.calcite.utils.DataTypeTransformer;
 
 public class CalciteRexNodeVisitor extends AbstractNodeVisitor<RexNode, CalcitePlanContext> {
 
@@ -198,13 +197,9 @@ public RexNode visitSpan(Span node, CalcitePlanContext context) {
     RelDataTypeFactory typeFactory = context.rexBuilder.getTypeFactory();
     SpanUnit unit = node.getUnit();
     if (isTimeBased(unit)) {
-      String datetimeUnitString = DataTypeTransformer.translate(unit);
-      RexNode interval =
-          context.rexBuilder.makeIntervalLiteral(
-              new BigDecimal(value.toString()),
-              new SqlIntervalQualifier(datetimeUnitString, SqlParserPos.ZERO));
-      // TODO not supported yet
-      return interval;
+      SqlIntervalQualifier intervalQualifier = context.rexBuilder.createIntervalUntil(unit);
+      long millis = SqlParserUtil.intervalToMillis(value.toString(), intervalQualifier);
+      return context.rexBuilder.makeIntervalLiteral(new BigDecimal(millis), intervalQualifier);
     } else {
       // if the unit is not time base - create a math expression to bucket the span partitions
       return context.rexBuilder.makeCall(

diff --git a/core/src/main/java/org/opensearch/sql/calcite/ExtendedRexBuilder.java b/core/src/main/java/org/opensearch/sql/calcite/ExtendedRexBuilder.java
@@ -5,9 +5,13 @@
 
 package org.opensearch.sql.calcite;
 
+import org.apache.calcite.avatica.util.TimeUnit;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlIntervalQualifier;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.parser.SqlParserPos;
+import org.opensearch.sql.ast.expression.SpanUnit;
 
 public class ExtendedRexBuilder extends RexBuilder {
 
@@ -22,4 +26,49 @@ public RexNode coalesce(RexNode... nodes) {
   public RexNode equals(RexNode n1, RexNode n2) {
     return this.makeCall(SqlStdOperatorTable.EQUALS, n1, n2);
   }
+
+  public SqlIntervalQualifier createIntervalUntil(SpanUnit unit) {
+    TimeUnit timeUnit;
+    switch (unit) {
+      case MILLISECOND:
+      case MS:
+        timeUnit = TimeUnit.MILLISECOND;
+        break;
+      case SECOND:
+      case S:
+        timeUnit = TimeUnit.SECOND;
+        break;
+      case MINUTE:
+      case m:
+        timeUnit = TimeUnit.MINUTE;
+        break;
+      case HOUR:
+      case H:
+        timeUnit = TimeUnit.HOUR;
+        break;
+      case DAY:
+      case D:
+        timeUnit = TimeUnit.DAY;
+        break;
+      case WEEK:
+      case W:
+        timeUnit = TimeUnit.WEEK;
+        break;
+      case MONTH:
+      case M:
+        timeUnit = TimeUnit.MONTH;
+        break;
+      case QUARTER:
+      case Q:
+        timeUnit = TimeUnit.QUARTER;
+        break;
+      case YEAR:
+      case Y:
+        timeUnit = TimeUnit.YEAR;
+        break;
+      default:
+        timeUnit = TimeUnit.EPOCH;
+    }
+    return new SqlIntervalQualifier(timeUnit, timeUnit, SqlParserPos.ZERO);
+  }
 }
diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/DataTypeTransformer.java b/core/src/main/java/org/opensearch/sql/calcite/utils/DataTypeTransformer.java
diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java b/core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java
@@ -12,6 +12,7 @@
 import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE;
 import static org.opensearch.sql.data.type.ExprCoreType.FLOAT;
 import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
+import static org.opensearch.sql.data.type.ExprCoreType.INTERVAL;
 import static org.opensearch.sql.data.type.ExprCoreType.IP;
 import static org.opensearch.sql.data.type.ExprCoreType.LONG;
 import static org.opensearch.sql.data.type.ExprCoreType.SHORT;
@@ -151,6 +152,13 @@ public static ExprType convertRelDataTypeToExprType(RelDataType type) {
         return TIMESTAMP;
       case GEOMETRY:
         return IP;
+      case INTERVAL_YEAR:
+      case INTERVAL_MONTH:
+      case INTERVAL_DAY:
+      case INTERVAL_HOUR:
+      case INTERVAL_MINUTE:
+      case INTERVAL_SECOND:
+        return INTERVAL;
       case ARRAY:
         return ARRAY;
       case MAP:

diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java
@@ -8,8 +8,6 @@
 
 package org.opensearch.sql.executor;
 
-import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
-
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.List;
@@ -80,8 +78,7 @@ public void execute(
               (PrivilegedAction<Void>)
                   () -> {
                     final FrameworkConfig config = buildFrameworkConfig();
-                    final CalcitePlanContext context =
-                        CalcitePlanContext.create(config, TYPE_FACTORY);
+                    final CalcitePlanContext context = CalcitePlanContext.create(config);
                     executePlanByCalcite(analyze(plan, context), context, listener);
                     return null;
                   });