optimize reduce logic

Signed-off-by: xinyual <xinyual@amazon.com>

Author: xinyual

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

@@ -27,6 +27,7 @@
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexLambda;
 import org.apache.calcite.rex.RexLambdaRef;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlIntervalQualifier;
@@ -430,12 +431,10 @@ private List<RelDataType> modifyLambdaTypeByFunction(
       case "REDUCE": // For reduce case, the first type is acc should be any since it is the output
         // of accumulator lambda function
         if (originalType.size() == 2) {
-          if (defaultTypeForReduceAcc == null
-              || defaultTypeForReduceAcc.equals(originalType.get(1))) {
-            return List.of(originalType.get(1), originalType.get(0));
+          if (defaultTypeForReduceAcc != null) {
+            return List.of(defaultTypeForReduceAcc, originalType.get(0));
           }
-          return List.of(TYPE_FACTORY.createSqlType(SqlTypeName.ANY, true), originalType.get(0));
-
+          return List.of(originalType.get(1), originalType.get(0));
         } else {
           return List.of(originalType.get(2));
         }
@@ -444,6 +443,19 @@ private List<RelDataType> modifyLambdaTypeByFunction(
     }
   }
 
+  private List<RexNode> castArgument(
+      List<RexNode> originalArguments, String functionName, ExtendedRexBuilder rexBuilder) {
+    switch (functionName.toUpperCase(Locale.ROOT)) {
+      case "REDUCE":
+        RexLambda call = (RexLambda) originalArguments.get(2);
+        originalArguments.set(
+            1, rexBuilder.makeCast(call.getType(), originalArguments.get(1), true, true));
+        return originalArguments;
+      default:
+        return originalArguments;
+    }
+  }
+
   @Override
   public RexNode visitFunction(Function node, CalcitePlanContext context) {
     List<UnresolvedExpression> args = node.getFuncArgs();
@@ -469,6 +481,9 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
         arguments.add(analyze(arg, context));
       }
     }
+
+    arguments = castArgument(arguments, node.getFuncName(), context.rexBuilder);
+
     RexNode resolvedNode =
         PPLFuncImpTable.INSTANCE.resolve(
             context.rexBuilder, node.getFuncName(), arguments.toArray(new RexNode[0]));

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

@@ -245,7 +245,7 @@ public static ExprType convertSqlTypeNameToExprType(SqlTypeName sqlTypeName) {
       case ARRAY -> ARRAY;
       case MAP -> STRUCT;
       case GEOMETRY -> GEO_POINT;
-      case NULL, ANY -> UNDEFINED;
+      case NULL, ANY, OTHER -> UNDEFINED;
       default -> UNKNOWN;
     };
   }

core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ArrayFunctionImpl.java

@@ -7,6 +7,7 @@
 
 import static org.apache.calcite.sql.type.SqlTypeUtil.createArrayType;
 
+import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -82,7 +83,20 @@ public static Object internalCast(Object... args) {
     SqlTypeName targetType = (SqlTypeName) args[args.length - 1];
     List<Object> result;
     switch (targetType) {
-      case DOUBLE, DECIMAL:
+      case DECIMAL:
+        result =
+            originalList.stream()
+                .map(
+                    num -> {
+                      if (num instanceof BigDecimal) {
+                        return (BigDecimal) num;
+                      } else {
+                        return BigDecimal.valueOf(((Number) num).doubleValue());
+                      }
+                    })
+                .collect(Collectors.toList());
+        break;
+      case DOUBLE:
         result =
             originalList.stream()
                 .map(i -> (Object) ((Number) i).doubleValue())

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

@@ -398,7 +398,9 @@ private static List<ExprType> getExprTypes(SqlTypeFamily family) {
             // Integer is mapped to BIGINT in family.getDefaultConcreteType
           case INTEGER -> List.of(
               OpenSearchTypeFactory.TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
-          case ANY, IGNORE -> List.of(
+            // case DATETIME_INTERVAL ->
+            // SqlTypeName.INTERVAL_TYPES.stream().map(OpenSearchTypeFactory.TYPE_FACTORY::createSqlIntervalType).toList();
+          case ANY, IGNORE, DATETIME_INTERVAL -> List.of(
               OpenSearchTypeFactory.TYPE_FACTORY.createSqlType(SqlTypeName.ANY));
           default -> {
             RelDataType type = family.getDefaultConcreteType(OpenSearchTypeFactory.TYPE_FACTORY);

core/src/test/java/org/opensearch/sql/calcite/CalciteRexNodeVisitorTest.java

@@ -152,6 +152,39 @@ public void testPrepareLambdaForReduce() {
         SqlTypeName.VARCHAR);
   }
 
+  @Test
+  public void testPrepareLambdaForReduceWithDefaultType() {
+    when(componentType.getSqlTypeName()).thenReturn(SqlTypeName.DOUBLE);
+    when(arrayArg.getType()).thenReturn(arraySqlType);
+    when(arraySqlType.getComponentType()).thenReturn(componentType);
+    when(extraArg.getType()).thenReturn(extraType);
+
+    List<RexNode> previousArguments = List.of(arrayArg, extraArg);
+    when(functionArg1.toString()).thenReturn("acc");
+    when(functionArg2.toString()).thenReturn("arg1");
+    when(lambdaFunction.getFuncArgs()).thenReturn(List.of(functionArg1, functionArg2));
+
+    CalcitePlanContext lambdaContext =
+        visitor.prepareLambdaContext(
+            context,
+            lambdaFunction,
+            previousArguments,
+            "reduce",
+            TYPE_FACTORY.createSqlType(SqlTypeName.BIGINT));
+
+    assertNotNull(lambdaContext);
+    assertNotNull(lambdaContext.getRexLambdaRefMap());
+    assertEquals(2, lambdaContext.getRexLambdaRefMap().size());
+    assertTrue(lambdaContext.getRexLambdaRefMap().containsKey("arg1"));
+    assertTrue(lambdaContext.getRexLambdaRefMap().containsKey("acc"));
+    assertEquals(
+        lambdaContext.getRexLambdaRefMap().get("arg1").getType().getSqlTypeName(),
+        SqlTypeName.DOUBLE);
+    assertEquals(
+        lambdaContext.getRexLambdaRefMap().get("acc").getType().getSqlTypeName(),
+        SqlTypeName.BIGINT);
+  }
+
   @Test
   public void testPrepareLambdaForReduceFinalizerFunction() {
     when(arrayArg.getType()).thenReturn(arraySqlType);

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalciteArrayFunctionIT.java

@@ -211,6 +211,21 @@ public void testReduce2() {
     verifyDataRows(actual, rows(1230));
   }
 
+  @Test
+  public void testReduce3() {
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | where age=28 | eval array = array(1.0, 2.0, 3.0), result3 ="
+                    + " reduce(array, age, (acc, x) -> acc * 1.0 + x, acc -> acc * 10.0) | fields"
+                    + " result3 | head 1",
+                TEST_INDEX_BANK));
+
+    verifySchema(actual, schema("result3", "double"));
+
+    verifyDataRows(actual, rows(340));
+  }
+
   @Test
   public void testReduceWithUDF() {
     JSONObject actual =