Define operand type inference for transform function

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

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

@@ -11,6 +11,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.IntStream;
 import org.apache.calcite.adapter.enumerable.NotNullImplementor;
 import org.apache.calcite.adapter.enumerable.NullPolicy;
@@ -22,18 +23,22 @@
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlLambda;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.SqlOperandCountRange;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlUtil;
 import org.apache.calcite.sql.type.ArraySqlType;
 import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
 import org.apache.calcite.sql.type.SqlOperandCountRanges;
+import org.apache.calcite.sql.type.SqlOperandTypeInference;
 import org.apache.calcite.sql.type.SqlReturnTypeInference;
 import org.apache.calcite.sql.type.SqlSingleOperandTypeChecker;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.type.SqlTypeUtil;
+import org.apache.calcite.sql.validate.SqlLambdaScope;
+import org.apache.calcite.sql.validate.SqlValidator;
 import org.opensearch.sql.expression.function.ImplementorUDF;
 import org.opensearch.sql.expression.function.UDFOperandMetadata;
 
@@ -58,6 +63,35 @@ public SqlReturnTypeInference getReturnTypeInference() {
     };
   }
 
+  @Override
+  public SqlOperandTypeInference getOperandTypeInference() {
+    // Pass the element type of TRANSFORM's first argument as the type of the first argument of the
+    // lambda function.
+    return (callBinding, returnType, operandTypes) -> {
+      RelDataType arrayType = callBinding.getOperandType(0);
+      operandTypes[0] = arrayType;
+      if (callBinding.operand(1) instanceof SqlLambda lambdaNode) {
+        SqlValidator validator = callBinding.getValidator();
+        if (validator.getLambdaScope(lambdaNode) instanceof SqlLambdaScope lambdaScope) {
+          RelDataType elementType = arrayType.getComponentType();
+          Map<String, RelDataType> paramTypes = lambdaScope.getParameterTypes();
+          List<SqlNode> params = lambdaNode.getParameters();
+          // First parameter: array element type. Leave it as is (typically ANY) if element type is
+          // null
+          if (!params.isEmpty() && elementType != null) {
+            paramTypes.put(params.get(0).toString(), elementType);
+          }
+          // Second parameter (if exists): INTEGER (for index)
+          if (params.size() > 1) {
+            RelDataType intType = callBinding.getTypeFactory().createSqlType(SqlTypeName.INTEGER);
+            paramTypes.put(params.get(1).toString(), intType);
+          }
+          operandTypes[1] = SqlTypeUtil.deriveType(callBinding, lambdaNode);
+        }
+      }
+    };
+  }
+
   @Override
   public UDFOperandMetadata getOperandMetadata() {
     // Only checks the first two arguments as it allows arbitrary number of arguments to follow them

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

@@ -12,6 +12,7 @@
 import org.apache.calcite.sql.SqlOperandCountRange;
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.sql.type.InferTypes;
+import org.apache.calcite.sql.type.SqlOperandTypeInference;
 import org.apache.calcite.sql.type.SqlReturnTypeInference;
 import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
 
@@ -37,6 +38,16 @@ default SqlKind getKind() {
     return SqlKind.OTHER_FUNCTION;
   }
 
+  /**
+   * Define the strategy to infer unknown types of the operands of an operator call.
+   *
+   * @return SqlOperandTypeInference the specified operand type inference. Default to {@link
+   *     InferTypes#ANY_NULLABLE}
+   */
+  default SqlOperandTypeInference getOperandTypeInference() {
+    return InferTypes.ANY_NULLABLE;
+  }
+
   default SqlUserDefinedFunction toUDF(String functionName) {
     return toUDF(functionName, true);
   }
@@ -57,7 +68,7 @@ default SqlUserDefinedFunction toUDF(String functionName, boolean isDeterministi
         udfLtrimIdentifier,
         getKind(),
         getReturnTypeInference(),
-        InferTypes.ANY_NULLABLE,
+        getOperandTypeInference(),
         getOperandMetadata(),
         getFunction()) {
       @Override

docs/user/ppl/functions/collection.md

@@ -162,8 +162,8 @@ Usage: `transform(array, function)` transform the element of array one by one us
 Argument type: array:ARRAY, function:LAMBDA
 Return type: ARRAY
 Example
-<!-- TODO: To be fixed with https://github.com/opensearch-project/sql/issues/4972 -->
-```ppl ignore
+
+```ppl
 source=people
 | eval array = array(1, -2, 3), result = transform(array, x -> x + 2)
 | fields result
@@ -180,8 +180,8 @@ fetched rows / total rows = 1/1
 | [3,0,5] |
 +---------+
 ```
-<!-- TODO: To be fixed with https://github.com/opensearch-project/sql/issues/4972 -->
-```ppl ignore
+
+```ppl
 source=people
 | eval array = array(1, -2, 3), result = transform(array, (x, i) -> x + i)
 | fields result
@@ -814,8 +814,8 @@ Usage: mvmap(array, expression) iterates over each element of a multivalue array
 Argument type: array: ARRAY, expression: EXPRESSION
 Return type: ARRAY
 Example
-<!-- TODO: To be fixed with https://github.com/opensearch-project/sql/issues/4972 -->
-```ppl ignore
+
+```ppl
 source=people
 | eval array = array(1, 2, 3), result = mvmap(array, array * 10)
 | fields result
@@ -832,8 +832,8 @@ fetched rows / total rows = 1/1
 | [10,20,30] |
 +------------+
 ```
-<!-- TODO: To be fixed with https://github.com/opensearch-project/sql/issues/4972 -->
-```ppl ignore
+
+```ppl
 source=people
 | eval array = array(1, 2, 3), result = mvmap(array, array + 5)
 | fields result
@@ -854,8 +854,8 @@ fetched rows / total rows = 1/1
 Note: For nested expressions like ``mvmap(mvindex(arr, 1, 3), arr * 2)``, the field name (``arr``) is extracted from the first argument and must match the field referenced in the expression.
 
 The expression can also reference other single-value fields:
-<!-- TODO: To be fixed with https://github.com/opensearch-project/sql/issues/4972 -->
-```ppl ignore
+
+```ppl
 source=people
 | eval array = array(1, 2, 3), multiplier = 10, result = mvmap(array, array * multiplier)
 | fields result