fix reduce and add doc

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

Author: xinyual

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

@@ -20,6 +20,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.type.RelDataType;
@@ -381,15 +382,17 @@ public CalcitePlanContext prepareLambdaContext(
       CalcitePlanContext context,
       LambdaFunction node,
       List<RexNode> previousArgument,
-      String functionName) {
+      String functionName,
+      @Nullable RelDataType defaultTypeForReduceAcc) {
     try {
       CalcitePlanContext lambdaContext = context.clone();
       List<RelDataType> candidateType = new ArrayList<>();
       candidateType.add(
           ((ArraySqlType) previousArgument.get(0).getType())
               .getComponentType()); // The first argument should be array type
       candidateType.addAll(previousArgument.stream().skip(1).map(RexNode::getType).toList());
-      candidateType = modifyLambdaTypeByFunction(functionName, candidateType);
+      candidateType =
+          modifyLambdaTypeByFunction(functionName, candidateType, defaultTypeForReduceAcc);
       List<QualifiedName> argNames = node.getFuncArgs();
       Map<String, RexLambdaRef> lambdaTypes = new HashMap<>();
       int candidateIndex;
@@ -421,12 +424,19 @@ public CalcitePlanContext prepareLambdaContext(
    *     reduce has special logic.
    */
   private List<RelDataType> modifyLambdaTypeByFunction(
-      String functionName, List<RelDataType> originalType) {
+      String functionName,
+      List<RelDataType> originalType,
+      @Nullable RelDataType defaultTypeForReduceAcc) {
     switch (functionName.toUpperCase(Locale.ROOT)) {
       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) {
-          return List.of(originalType.get(1), originalType.get(0));
+          if (defaultTypeForReduceAcc == null
+              || defaultTypeForReduceAcc.equals(originalType.get(1))) {
+            return List.of(originalType.get(1), originalType.get(0));
+          }
+          return List.of(TYPE_FACTORY.createSqlType(SqlTypeName.ANY, true), originalType.get(0));
+
         } else {
           return List.of(originalType.get(2));
         }
@@ -442,8 +452,20 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
     for (UnresolvedExpression arg : args) {
       if (arg instanceof LambdaFunction) {
         CalcitePlanContext lambdaContext =
-            prepareLambdaContext(context, (LambdaFunction) arg, arguments, node.getFuncName());
-        arguments.add(analyze(arg, lambdaContext));
+            prepareLambdaContext(
+                context, (LambdaFunction) arg, arguments, node.getFuncName(), null);
+        RexNode lambdaNode = analyze(arg, lambdaContext);
+        if (node.getFuncName().equalsIgnoreCase("reduce")) { // analyze again with calculate type
+          lambdaContext =
+              prepareLambdaContext(
+                  context,
+                  (LambdaFunction) arg,
+                  arguments,
+                  node.getFuncName(),
+                  lambdaNode.getType());
+          lambdaNode = analyze(arg, lambdaContext);
+        }
+        arguments.add(lambdaNode);
       } else {
         arguments.add(analyze(arg, context));
       }

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

@@ -60,21 +60,6 @@ public SqlReturnTypeInference getReturnTypeInference() {
         finalReturnType = mergedReturnType;
       }
       return finalReturnType;
-
-      /*
-      RelDataTypeFactory typeFactory = sqlOperatorBinding.getTypeFactory();
-      RexCallBinding rexCallBinding = (RexCallBinding) sqlOperatorBinding;
-      List<RexNode> operands = rexCallBinding.operands();
-      RelDataType mergedReturnType =
-              ((RexLambda) operands.get(2)).getExpression().getType();
-      if (operands.size() > 3) {
-        RelDataType reduceReturnType =
-                ((RexLambda) operands.get(3)).getExpression().getType();
-        return typeFactory.leastRestrictive(List.of(mergedReturnType, reduceReturnType));
-      }
-      return mergedReturnType;
-
-       */
     };
   }
 

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

@@ -85,7 +85,7 @@ public void testPrepareLambdaForBasicLambda() {
     when(lambdaFunction.getFuncArgs()).thenReturn(List.of(functionArg1));
 
     CalcitePlanContext lambdaContext =
-        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "forall");
+        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "forall", null);
 
     assertNotNull(lambdaContext);
     assertNotNull(lambdaContext.getRexLambdaRefMap());
@@ -108,7 +108,7 @@ public void testPrepareLambdaForTransform() {
     when(lambdaFunction.getFuncArgs()).thenReturn(List.of(functionArg1, functionArg2));
 
     CalcitePlanContext lambdaContext =
-        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "transform");
+        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "transform", null);
 
     assertNotNull(lambdaContext);
     assertNotNull(lambdaContext.getRexLambdaRefMap());
@@ -137,7 +137,7 @@ public void testPrepareLambdaForReduce() {
     when(lambdaFunction.getFuncArgs()).thenReturn(List.of(functionArg1, functionArg2));
 
     CalcitePlanContext lambdaContext =
-        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "reduce");
+        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "reduce", null);
 
     assertNotNull(lambdaContext);
     assertNotNull(lambdaContext.getRexLambdaRefMap());
@@ -165,7 +165,7 @@ public void testPrepareLambdaForReduceFinalizerFunction() {
     when(lambdaFunction.getFuncArgs()).thenReturn(List.of(functionArg1));
 
     CalcitePlanContext lambdaContext =
-        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "reduce");
+        visitor.prepareLambdaContext(context, lambdaFunction, previousArguments, "reduce", null);
 
     assertNotNull(lambdaContext);
     assertNotNull(lambdaContext.getRexLambdaRefMap());