[FLINK-37789][planner] Support to validate ML_PREDICT expression

Author: lihaosky

flink-table/flink-sql-parser/src/main/codegen/data/Parser.tdd

@@ -574,6 +574,7 @@
     "JSON_EXECUTION_PLAN"
     "PLAN_ADVICE"
     "METADATA"
+    "MODEL"
     "OVERWRITE"
     "OVERWRITING"
     "PARTITIONED"

flink-table/flink-sql-parser/src/main/java/org/apache/calcite/sql/SqlExplicitModelCall.java

@@ -32,4 +32,8 @@ public SqlExplicitModelCall(
             @Nullable SqlLiteral functionQualifier) {
         super(operator, operandList, pos, functionQualifier);
     }
+
+    public SqlIdentifier getModelIdentifier() {
+        return operand(0);
+    }
 }

flink-table/flink-sql-parser/src/main/java/org/apache/calcite/sql/SqlExplicitModelOperator.java

@@ -17,10 +17,15 @@
 
 package org.apache.calcite.sql;
 
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql.validate.SqlValidatorScope;
 import org.apache.calcite.util.ImmutableNullableList;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import static org.apache.calcite.util.Static.RESOURCE;
+
 /** SqlExplicitModelOperator is a SQL operator that represents an explicit model. */
 public class SqlExplicitModelOperator extends SqlPrefixOperator {
 
@@ -37,4 +42,20 @@ public SqlCall createCall(
         return new SqlExplicitModelCall(
                 this, ImmutableNullableList.copyOf(operands), pos, functionQualifier);
     }
+
+    @Override
+    public void validateCall(
+            SqlCall call,
+            SqlValidator validator,
+            SqlValidatorScope scope,
+            SqlValidatorScope operandScope) {
+        throw SqlUtil.newContextException(
+                call.pos, RESOURCE.objectNotFound(call.operand(0).toString()));
+    }
+
+    @Override
+    public RelDataType deriveType(SqlValidator validator, SqlValidatorScope scope, SqlCall call) {
+        throw SqlUtil.newContextException(
+                call.pos, RESOURCE.objectNotFound(call.operand(0).toString()));
+    }
 }

flink-table/flink-sql-parser/src/test/java/org/apache/flink/sql/parser/FlinkSqlParserImplTest.java

@@ -3376,6 +3376,14 @@ void testModelInFunctionWithoutTable() {
                                 + "FROM TABLE(`FUNC`((TABLE `MY_TABLE`), MODEL `CAT`.`DB`.`MY_MODEL`))");
     }
 
+    @Test
+    void testModelInFunctionNamedArgs() {
+        sql("select * from table(ml_predict(INPUT => TABLE my_table, model => MODEL my_model))")
+                .ok(
+                        "SELECT *\n"
+                                + "FROM TABLE(`ML_PREDICT`(`INPUT` => (TABLE `MY_TABLE`), `MODEL` => (MODEL `MY_MODEL`)))");
+    }
+
     /*
      * This test was backported from Calcite 1.38 (CALCITE-6266).
      * Remove it together with upgrade to Calcite 1.38.

flink-table/flink-table-planner/src/main/java/org/apache/calcite/sql/SqlModelCall.java

@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.calcite.sql;
+
+import org.apache.flink.table.planner.catalog.CatalogSchemaModel;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql.validate.SqlValidatorScope;
+
+import static java.util.Objects.requireNonNull;
+
+/** SqlModelCall to fetch and reference model based on identifier. */
+public class SqlModelCall extends SqlBasicCall {
+
+    private final CatalogSchemaModel model;
+
+    public SqlModelCall(SqlExplicitModelCall modelCall, CatalogSchemaModel model) {
+        super(
+                new SqlModelOperator(model),
+                modelCall.getOperandList(),
+                modelCall.getParserPosition(),
+                modelCall.getFunctionQuantifier());
+        this.model = requireNonNull(model);
+    }
+
+    @Override
+    public void validate(SqlValidator validator, SqlValidatorScope scope) {
+        // Do nothing here, override to avoid identifier validation which will be treated as column
+    }
+
+    public RelDataType getInputType(SqlValidator validator) {
+        return model.getInputRowType(validator.getTypeFactory());
+    }
+
+    public RelDataType getOutputType(SqlValidator validator) {
+        return model.getOutputRowType(validator.getTypeFactory());
+    }
+
+    /**
+     * A custom SqlOperator to handle model identifier.
+     *
+     * <p>It is used to derive the type of the model based on the identifier.
+     */
+    private static class SqlModelOperator extends SqlPrefixOperator {
+
+        CatalogSchemaModel model;
+
+        private SqlModelOperator(CatalogSchemaModel model) {
+            super("MODEL", SqlKind.OTHER_FUNCTION, 2, null, null, null);
+            this.model = model;
+        }
+
+        @Override
+        public RelDataType deriveType(
+                SqlValidator validator, SqlValidatorScope scope, SqlCall call) {
+            return model.getOutputRowType(validator.getTypeFactory());
+        }
+    }
+}

flink-table/flink-table-planner/src/main/java/org/apache/calcite/sql/validate/ProcedureNamespace.java

@@ -18,6 +18,7 @@
 
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.table.planner.calcite.FlinkSqlCallBinding;
+import org.apache.flink.table.planner.functions.sql.ml.SqlMLTableFunction;
 
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.SqlCall;
@@ -61,7 +62,7 @@ public RelDataType validateImpl(RelDataType targetRowType) {
         final SqlOperator operator = call.getOperator();
         final SqlCallBinding callBinding = new FlinkSqlCallBinding(validator, scope, call);
         final SqlCall permutedCall = callBinding.permutedCall();
-        if (operator instanceof SqlWindowTableFunction) {
+        if (operator instanceof SqlWindowTableFunction || operator instanceof SqlMLTableFunction) {
             permutedCall.validate(validator, scope);
         }
 

flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/calcite/FlinkCalciteSqlValidator.java

@@ -26,7 +26,10 @@
 import org.apache.flink.table.catalog.Column;
 import org.apache.flink.table.catalog.ResolvedSchema;
 import org.apache.flink.table.data.TimestampData;
+import org.apache.flink.table.planner.catalog.CatalogSchemaModel;
 import org.apache.flink.table.planner.catalog.CatalogSchemaTable;
+import org.apache.flink.table.planner.functions.sql.ml.SqlMLTableFunction;
+import org.apache.flink.table.planner.plan.FlinkCalciteCatalogReader;
 import org.apache.flink.table.planner.plan.utils.FlinkRexUtil;
 import org.apache.flink.table.planner.utils.ShortcutUtils;
 import org.apache.flink.table.types.logical.DecimalType;
@@ -43,12 +46,14 @@
 import org.apache.calcite.sql.SqlAsOperator;
 import org.apache.calcite.sql.SqlBasicCall;
 import org.apache.calcite.sql.SqlCall;
+import org.apache.calcite.sql.SqlExplicitModelCall;
 import org.apache.calcite.sql.SqlFunction;
 import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlIdentifier;
 import org.apache.calcite.sql.SqlJoin;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlLiteral;
+import org.apache.calcite.sql.SqlModelCall;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.SqlNodeList;
 import org.apache.calcite.sql.SqlOperator;
@@ -371,7 +376,21 @@ protected void addToSelectList(
         final SqlBasicCall call = (SqlBasicCall) node;
         final SqlOperator operator = call.getOperator();
 
-        if (operator instanceof SqlWindowTableFunction) {
+        if (node instanceof SqlExplicitModelCall) {
+            // Convert it so that model can be accessed in planner. SqlExplicitModelCall
+            // from parser can't access model.
+            SqlExplicitModelCall modelCall = (SqlExplicitModelCall) node;
+            SqlIdentifier modelIdentifier = modelCall.getModelIdentifier();
+            FlinkCalciteCatalogReader catalogReader =
+                    (FlinkCalciteCatalogReader) getCatalogReader();
+            CatalogSchemaModel model = catalogReader.getModel(modelIdentifier.names);
+            if (model != null) {
+                return new SqlModelCall(modelCall, model);
+            }
+        }
+
+        // TODO (FLINK-37819): add test for SqlMLTableFunction
+        if (operator instanceof SqlWindowTableFunction || operator instanceof SqlMLTableFunction) {
             if (tableArgs.stream().allMatch(Objects::isNull)) {
                 return rewritten;
             }

flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/sql/FlinkSqlOperatorTable.java

@@ -21,6 +21,7 @@
 import org.apache.flink.table.api.TableException;
 import org.apache.flink.table.planner.calcite.FlinkTypeFactory;
 import org.apache.flink.table.planner.functions.sql.internal.SqlAuxiliaryGroupAggFunction;
+import org.apache.flink.table.planner.functions.sql.ml.SqlMLPredictTableFunction;
 import org.apache.flink.table.planner.plan.type.FlinkReturnTypes;
 import org.apache.flink.table.planner.plan.type.NumericExceptFirstOperandChecker;
 
@@ -1341,6 +1342,9 @@ public List<SqlGroupedWindowFunction> getAuxiliaryFunctions() {
     public static final SqlFunction CUMULATE = new SqlCumulateTableFunction();
     public static final SqlFunction SESSION = new SqlSessionTableFunction();
 
+    // MODEL TABLE FUNCTIONS
+    public static final SqlFunction ML_PREDICT = new SqlMLPredictTableFunction();
+
     // Catalog Functions
     public static final SqlFunction CURRENT_DATABASE =
             BuiltInSqlFunction.newBuilder()

flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/functions/sql/ml/SqlMLPredictTableFunction.java

@@ -0,0 +1,116 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.table.planner.functions.sql.ml;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlOperandCountRange;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlOperatorBinding;
+import org.apache.calcite.sql.type.SqlOperandCountRanges;
+import org.apache.calcite.sql.type.SqlOperandMetadata;
+import org.apache.calcite.sql.type.SqlTypeName;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * SqlMlPredictTableFunction implements an operator for prediction.
+ *
+ * <p>It allows four parameters:
+ *
+ * <ol>
+ *   <li>a table
+ *   <li>a model name
+ *   <li>a descriptor to provide a column name from the input table
+ *   <li>an optional config map
+ * </ol>
+ */
+public class SqlMLPredictTableFunction extends SqlMLTableFunction {
+
+    public SqlMLPredictTableFunction() {
+        super("ML_PREDICT", new PredictOperandMetadata());
+    }
+
+    /**
+     * {@inheritDoc}
+     *
+     * <p>Overrides because the first parameter of table-value function windowing is an explicit
+     * TABLE parameter, which is not scalar.
+     */
+    @Override
+    public boolean argumentMustBeScalar(int ordinal) {
+        return ordinal != 0;
+    }
+
+    @Override
+    protected RelDataType inferRowType(SqlOperatorBinding opBinding) {
+        // TODO: FLINK-37780 output type based on table schema and model output schema
+        // model output schema to be available after integrated with SqlExplicitModelCall
+        return opBinding.getOperandType(1);
+    }
+
+    private static class PredictOperandMetadata implements SqlOperandMetadata {
+        private static final List<String> PARAM_NAMES =
+                List.of(PARAM_INPUT, PARAM_MODEL, PARAM_COLUMN, PARAM_CONFIG);
+        private static final List<String> MANDATORY_PARAM_NAMES =
+                List.of(PARAM_INPUT, PARAM_MODEL, PARAM_COLUMN);
+
+        PredictOperandMetadata() {}
+
+        @Override
+        public List<RelDataType> paramTypes(RelDataTypeFactory typeFactory) {
+            return Collections.nCopies(
+                    PARAM_NAMES.size(), typeFactory.createSqlType(SqlTypeName.ANY));
+        }
+
+        @Override
+        public List<String> paramNames() {
+            return PARAM_NAMES;
+        }
+
+        @Override
+        public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure) {
+            // TODO: FLINK-37780 Check operand types after integrated with SqlExplicitModelCall in
+            // validator
+            return false;
+        }
+
+        @Override
+        public SqlOperandCountRange getOperandCountRange() {
+            return SqlOperandCountRanges.between(MANDATORY_PARAM_NAMES.size(), PARAM_NAMES.size());
+        }
+
+        @Override
+        public Consistency getConsistency() {
+            return Consistency.NONE;
+        }
+
+        @Override
+        public boolean isOptional(int i) {
+            return i > getOperandCountRange().getMin() && i <= getOperandCountRange().getMax();
+        }
+
+        @Override
+        public String getAllowedSignatures(SqlOperator op, String opName) {
+            return opName
+                    + "(TABLE table_name, MODEL model_name, DESCRIPTOR(input_columns), [MAP[]]";
+        }
+    }
+}