Support decimal literal with Calcite (opensearch-project#3673)

Author: LantaoJin

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -6,6 +6,7 @@
 package org.opensearch.sql.ast.dsl;
 
 import com.google.common.collect.ImmutableList;
+import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -213,6 +214,14 @@ public static Literal doubleLiteral(Double value) {
     return literal(value, DataType.DOUBLE);
   }
 
+  public static Literal decimalLiteral(Double value) {
+    return literal(BigDecimal.valueOf(value), DataType.DECIMAL);
+  }
+
+  public static Literal decimalLiteral(BigDecimal value) {
+    return literal(value, DataType.DECIMAL);
+  }
+
   public static Literal stringLiteral(String value) {
     return literal(value, DataType.STRING);
   }

core/src/main/java/org/opensearch/sql/ast/expression/DataType.java

@@ -9,7 +9,7 @@
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.data.type.ExprCoreType;
 
-/** The DataType defintion in AST. Question, could we use {@link ExprCoreType} directly in AST? */
+/** The DataType definition in AST. Question, could we use {@link ExprCoreType} directly in AST? */
 @RequiredArgsConstructor
 public enum DataType {
   TYPE_ERROR(ExprCoreType.UNKNOWN),
@@ -28,6 +28,11 @@ public enum DataType {
   TIMESTAMP(ExprCoreType.TIMESTAMP),
   INTERVAL(ExprCoreType.INTERVAL),
 
+  // the decimal DataType is only used for building decimal literal,
+  // so it still maps to double core type until we support decimal type.
+  // ref https://github.com/opensearch-project/sql/issues/3619
+  DECIMAL(ExprCoreType.DOUBLE),
+
   IP(ExprCoreType.IP);
 
   @Getter private final ExprCoreType coreType;

core/src/main/java/org/opensearch/sql/ast/expression/Literal.java

@@ -6,10 +6,10 @@
 package org.opensearch.sql.ast.expression;
 
 import com.google.common.collect.ImmutableList;
+import java.math.BigDecimal;
 import java.util.List;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 
 /**
@@ -18,12 +18,20 @@
  */
 @Getter
 @EqualsAndHashCode(callSuper = false)
-@RequiredArgsConstructor
 public class Literal extends UnresolvedExpression {
 
   private final Object value;
   private final DataType type;
 
+  public Literal(Object value, DataType dataType) {
+    if (dataType == DataType.DECIMAL && value instanceof Double) {
+      this.value = BigDecimal.valueOf((Double) value);
+    } else {
+      this.value = value;
+    }
+    this.type = dataType;
+  }
+
   @Override
   public List<UnresolvedExpression> getChild() {
     return ImmutableList.of();

core/src/main/java/org/opensearch/sql/ast/tree/AD.java

@@ -20,7 +20,7 @@
 @Getter
 @Setter
 @ToString
-@EqualsAndHashCode(callSuper = true)
+@EqualsAndHashCode(callSuper = false)
 @RequiredArgsConstructor
 @AllArgsConstructor
 public class AD extends UnresolvedPlan {

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

@@ -122,6 +122,8 @@ public RexNode visitLiteral(Literal node, CalcitePlanContext context) {
         return rexBuilder.makeApproxLiteral(
             new BigDecimal(Double.toString((Double) value)),
             typeFactory.createSqlType(SqlTypeName.DOUBLE));
+      case DECIMAL:
+        return rexBuilder.makeExactLiteral((BigDecimal) value);
       case BOOLEAN:
         return rexBuilder.makeLiteral((Boolean) value);
       case DATE:
@@ -132,9 +134,6 @@ public RexNode visitLiteral(Literal node, CalcitePlanContext context) {
       case TIMESTAMP:
         return rexBuilder.makeTimestampLiteral(
             new TimestampString(value.toString()), RelDataType.PRECISION_NOT_SPECIFIED);
-      case INTERVAL:
-        //                return rexBuilder.makeIntervalLiteral(BigDecimal.valueOf((long)
-        // node.getValue()));
       default:
         throw new UnsupportedOperationException("Unsupported literal type: " + node.getType());
     }

core/src/main/java/org/opensearch/sql/calcite/utils/MathUtils.java

@@ -5,11 +5,17 @@
 
 package org.opensearch.sql.calcite.utils;
 
+import java.math.BigDecimal;
+
 public class MathUtils {
   public static boolean isIntegral(Number n) {
     return n instanceof Byte || n instanceof Short || n instanceof Integer || n instanceof Long;
   }
 
+  public static boolean isDecimal(Number n) {
+    return n instanceof BigDecimal;
+  }
+
   /**
    * Converts a long value to the least restrictive integral type based on the types of two input
    * numbers.

core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java

@@ -223,7 +223,7 @@ public static ExprType convertSqlTypeNameToExprType(SqlTypeName sqlTypeName) {
       case INTEGER -> INTEGER;
       case BIGINT -> LONG;
       case FLOAT, REAL -> FLOAT;
-      case DOUBLE -> DOUBLE;
+      case DOUBLE, DECIMAL -> DOUBLE; // TODO the decimal is only used for literal
       case CHAR, VARCHAR -> STRING;
       case BOOLEAN -> BOOLEAN;
       case DATE -> DATE;

core/src/main/java/org/opensearch/sql/data/model/ExprValueUtils.java

@@ -9,6 +9,7 @@
 import static org.opensearch.sql.utils.ExpressionUtils.PATH_SEP;
 
 import inet.ipaddr.IPAddress;
+import java.math.BigDecimal;
 import java.sql.Date;
 import java.sql.Time;
 import java.sql.Timestamp;
@@ -134,6 +135,10 @@ public static ExprValue fromObjectValue(Object o) {
         return LITERAL_NULL;
       }
       return doubleValue(d);
+    } else if (o instanceof BigDecimal d) {
+      // TODO fallback decimal to double in v2
+      // until https://github.com/opensearch-project/sql/issues/3619 fixed.
+      return new ExprDoubleValue(d);
     } else if (o instanceof String) {
       return stringValue((String) o);
     } else if (o instanceof Float f) {

core/src/main/java/org/opensearch/sql/executor/OpenSearchTypeSystem.java

@@ -18,11 +18,33 @@ private OpenSearchTypeSystem() {}
 
   @Override
   public RelDataType deriveAvgAggType(RelDataTypeFactory typeFactory, RelDataType argumentType) {
-    if (SqlTypeName.INT_TYPES.contains(argumentType.getSqlTypeName())) {
+    if (SqlTypeName.DECIMAL == argumentType.getSqlTypeName()) {
+      return typeFactory.createTypeWithNullability(highPrecision(typeFactory, argumentType), false);
+    } else if (SqlTypeName.INT_TYPES.contains(argumentType.getSqlTypeName())) {
       return typeFactory.createTypeWithNullability(
           typeFactory.createSqlType(SqlTypeName.DOUBLE), false);
     } else {
       return argumentType;
     }
   }
+
+  /**
+   * Compute a higher precision version of a type.
+   *
+   * @return If type is a DECIMAL type, return a type with the precision and scale +4, to align the
+   *     behaviour with Apache Spark.
+   */
+  public static RelDataType highPrecision(
+      final RelDataTypeFactory typeFactory, final RelDataType type) {
+    if (type.getSqlTypeName() == SqlTypeName.DECIMAL) {
+      return typeFactory.createSqlType(
+          type.getSqlTypeName(),
+          Math.min(
+              type.getPrecision() + 4,
+              typeFactory.getTypeSystem().getMaxPrecision(SqlTypeName.DECIMAL)),
+          Math.min(
+              type.getScale() + 4, typeFactory.getTypeSystem().getMaxScale(SqlTypeName.DECIMAL)));
+    }
+    return type;
+  }
 }

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -196,7 +196,7 @@ void populate() {
       registerOperator(ADD, SqlStdOperatorTable.PLUS);
       registerOperator(SUBTRACT, SqlStdOperatorTable.MINUS);
       registerOperator(MULTIPLY, SqlStdOperatorTable.MULTIPLY);
-      // registerOperator(DIVIDE, SqlStdOperatorTable.DIVIDE);
+      registerOperator(TRUNCATE, SqlStdOperatorTable.TRUNCATE);
       registerOperator(ASCII, SqlStdOperatorTable.ASCII);
       registerOperator(LENGTH, SqlStdOperatorTable.CHAR_LENGTH);
       registerOperator(LOWER, SqlStdOperatorTable.LOWER);