[FLINK-38258][model] support model args in ptf and add ml_predict builtin ptf #26924

lihaosky · 2025-08-18T17:06:41Z

What is the purpose of the change

Add model argument support in ptf
Add ml_predict builtin ptf definition

Brief change log

Add ModelSemantics in CallContext
Add model in StaticArgument
Add disableSystemArguments to control whether add uid and rowtime in ptf input/output
Add ml_predict builtin ptf
Add and update tests

Verifying this change

Add unit test for ml_predict type inference
Updated test for ml_predict function

Does this pull request potentially affect one of the following parts:

Dependencies (does it add or upgrade a dependency): (no)
The public API, i.e., is any changed class annotated with @Public(Evolving): (no)
The serializers: (no)
The runtime per-record code paths (performance sensitive): (no)
Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: (no)
The S3 file system connector: (no)

Documentation

Does this pull request introduce a new feature? (no)
If yes, how is the feature documented? (JavaDocs)

flinkbot · 2025-08-18T17:13:38Z

CI report:

94e4553 Azure: SUCCESS

Bot commands

The @flinkbot bot supports the following commands:

@flinkbot run azure re-run the last Azure build

davidradl · 2025-08-27T13:03:49Z

...k-table-common/src/main/java/org/apache/flink/table/types/inference/SystemTypeInference.java

@@ -147,7 +159,9 @@ private static void checkScalarArgsOnly(List<StaticArgument> defaultArgs) {
        checkPassThroughColumns(declaredArgs);

        final List<StaticArgument> newStaticArgs = new ArrayList<>(declaredArgs);
-        newStaticArgs.addAll(PROCESS_TABLE_FUNCTION_SYSTEM_ARGS);


I am curious what the system arguments mean. Is this something that the user needs to be aware of? I do not see this phrase in the Flip and there is no more information in the Jira. I suggest including a description and motivation behind this piece. It appears to be a type of static arg that will be added if the boolean flag is on, but I am not sure when this would/should be used.

This is to control whether uid, ontime field will be added to ptf input. This is currently used by ml_predict because it doesn't need uid and ontime field. It's not exposed to PTF function user can define. Yes. I can add more description if this approach makes sense.

This in not entirely correct. A user-defined PTF can implement a TypeInference and avoid system args, but this is kind of second-level API.

twalthr

Great job @lihaosky! I added some comments to improve the contribution a bit but nothing major.

.../main/java/org/apache/flink/table/expressions/resolver/rules/ResolveCallByArgumentsRule.java

twalthr · 2025-09-02T15:13:40Z

...k-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinition.java

@@ -298,6 +298,11 @@ public Builder staticArguments(StaticArgument... staticArguments) {
            return this;
        }

+        public Builder allowSystemArguments(boolean allowSystemArguments) {


let's call this disableSystemArguments and in TypeInference. By default, this then can be false.

...-table/flink-table-common/src/main/java/org/apache/flink/table/functions/ModelSemantics.java

twalthr · 2025-09-02T15:15:15Z

...-table/flink-table-common/src/main/java/org/apache/flink/table/functions/ModelSemantics.java

+    DataType inputDataType();
+
+    /**
+     * Output data type produced by the passed model. Extracting type from PTF class definition is


Suggested change

* Output data type produced by the passed model. Extracting type from PTF class definition is

* Output data type produced by the passed model.

twalthr · 2025-09-02T15:20:28Z

...-table/flink-table-common/src/main/java/org/apache/flink/table/functions/ModelTypeUtils.java

+        if (tableSemantics == null) {
+            if (throwOnFailure) {
+                throw new ValidationException(
+                        "First argument must be a table for ML_PREDICT function.");


An input type strategy is optional if static arguments have been declared. You can assume that this check has been done already. An input type strategy might only be useful if you want to do additional validation, like validateTableAndDescriptorArguments below

twalthr · 2025-09-03T14:37:00Z

...mmon/src/test/java/org/apache/flink/table/types/inference/strategies/ModelSemanticsMock.java

+import org.apache.flink.table.types.DataType;
+
+/** Mock implementation of {@link ModelSemantics} for testing purposes. */
+public class ModelSemanticsMock implements ModelSemantics {


move next to CallContextMock into utils package

twalthr · 2025-09-03T14:38:21Z

...src/main/java/org/apache/flink/table/planner/functions/inference/CallBindingCallContext.java

@@ -62,6 +65,7 @@ public final class CallBindingCallContext extends AbstractSqlCallContext {
    private final List<DataType> argumentDataTypes;
    private final @Nullable DataType outputType;
    private final @Nullable List<StaticArgument> staticArguments;
+    private final SqlValidator validator;


validator does not really fit in here, can we avoid it? SqlModelCall should have resolved types already. I think TableArgCall has the same? we should synchronize the two if possible.

twalthr · 2025-09-03T14:40:04Z

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

@@ -1327,7 +1326,7 @@ public List<SqlGroupedWindowFunction> getAuxiliaryFunctions() {
    public static final SqlFunction SESSION = new SqlSessionTableFunction();

    // MODEL TABLE FUNCTIONS
-    public static final SqlFunction ML_PREDICT = new SqlMLPredictTableFunction();
+    // public static final SqlFunction ML_PREDICT = new SqlMLPredictTableFunction();


...table-planner/src/main/java/org/apache/flink/table/planner/plan/QueryOperationConverter.java

twalthr · 2025-09-03T14:45:51Z

...r/src/test/java/org/apache/flink/table/planner/plan/stream/sql/ProcessTableFunctionTest.java

+        util.addTemporarySystemFunction("f", NoSystemArgsTableFunction.class);
+        assertThatThrownBy(() -> util.verifyRelPlan("SELECT * FROM f(r => TABLE t, i => 1);"))
+                .satisfies(
+                        anyCauseMatches("Disabling uid/time attributes is not supported for PTF."));


Suggested change

anyCauseMatches("Disabling uid/time attributes is not supported for PTF."));

anyCauseMatches("Disabling system arguments is not supported for user-defined PTF yet."));

lihaosky · 2025-09-04T17:01:09Z

...main/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java

+        if (tableSemantics == null) {
+            if (throwOnFailure) {
+                throw new ValidationException(
+                        "First argument must be a table for ML_PREDICT function.");
+            } else {
+                return Optional.empty();
+            }
+        }
+
+        // Check that second argument is a model
+        ModelSemantics modelSemantics = callContext.getModelSemantics(1).orElse(null);
+        if (modelSemantics == null) {
+            if (throwOnFailure) {
+                throw new ValidationException(
+                        "Second argument must be a model for ML_PREDICT function.");
+            } else {
+                return Optional.empty();
+            }
+        }
+
+        // Check that third argument is a descriptor with column names
+        Optional<ColumnList> descriptorColumns = callContext.getArgumentValue(2, ColumnList.class);
+        if (descriptorColumns.isEmpty()) {
+            if (throwOnFailure) {
+                throw new ValidationException(
+                        "Third argument must be a descriptor with simple column names for ML_PREDICT function.");
+            } else {
+                return Optional.empty();
+            }


@twalthr , I added these checks back since I have tests in MLPredictInputTypeStrategyTest testing invalid argument as well. Also tableSemantics etc are needed below. Maybe doesn't hurt to do extra check and give meaningful error message.

lihaosky · 2025-09-04T20:35:02Z

@flinkbot run azure

twalthr

Awesome work @lihaosky. I just added some final minor comments. The PR should be ready to merge in the next iteration.

twalthr · 2025-09-11T13:57:57Z

...main/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java

@@ -117,6 +133,35 @@ public static InputTypeStrategy windowTimeIndicator() {
                                    and(logical(LogicalTypeFamily.CHARACTER_STRING), LITERAL),
                                    JSON_ARGUMENT));

+    /** Input strategy for {@link BuiltInFunctionDefinitions#ML_PREDICT}. */
+    public static final InputTypeStrategy ML_PREDICT_INPUT_TYPE_STRATEGY =
+            new InputTypeStrategy() {


usually we put this strategy into a separate class in same package but with default scope. take ArrayComparableElementArgumentTypeStrategy as an example

twalthr · 2025-09-11T13:58:50Z

.../src/main/java/org/apache/flink/table/types/inference/strategies/SpecificTypeStrategies.java

@@ -194,6 +202,44 @@ public final class SpecificTypeStrategies {
    /** Type strategy specific for {@link BuiltInFunctionDefinitions#OBJECT_UPDATE}. */
    public static final TypeStrategy OBJECT_UPDATE = new ObjectUpdateTypeStrategy();

+    /** Type strategy specific for {@link BuiltInFunctionDefinitions#ML_PREDICT}. */
+    public static final TypeStrategy ML_PREDICT_OUTPUT_TYPE_STRATEGY =
+            callContext -> {


usually we put this strategy into a separate class in same package but with default scope. take ObjectUpdateTypeStrategy as an example.

twalthr · 2025-09-11T14:06:32Z

...e/flink/table/planner/plan/nodes/physical/stream/StreamPhysicalProcessTableFunctionRule.java

    }

    @Override
    public @Nullable RelNode convert(RelNode rel) {
        final FlinkLogicalTableFunctionScan scan = (FlinkLogicalTableFunctionScan) rel;
        final RexCall rexCall = (RexCall) scan.getCall();
+        validateAllowSystemArgs(rexCall);


Validating one time should be enough. So we can drop the call in StreamPhysicalProcessTableFunction.

twalthr · 2025-09-11T14:09:43Z

...r/src/test/java/org/apache/flink/table/planner/plan/stream/sql/ProcessTableFunctionTest.java

+    void testNoSystemArgsAllowedForScalarPtf() {
+        util.addTemporarySystemFunction("f", NoSystemArgsScalarFunction.class);
+        assertThatThrownBy(() -> util.verifyRelPlan("SELECT * FROM f(i => 1);"))
+                .satisfies(anyCauseMatches("Disabling system arguments is not supported for PTF."));


Suggested change

.satisfies(anyCauseMatches("Disabling system arguments is not supported for PTF."));

.satisfies(anyCauseMatches("Disabling system arguments is not supported for user-defined PTF."));

twalthr · 2025-09-11T14:10:29Z

...src/test/java/org/apache/flink/table/planner/plan/stream/sql/MLPredictTableFunctionTest.java

-                        "SQL validation failed. Unsupported expression -1 is in runtime config at position line 2, column 109. Currently, runtime config should be be a MAP of string literals.");
+                .hasCauseInstanceOf(ValidationException.class)
+                .hasStackTraceContaining(
+                        "Config param of ML_PREDICT function should be a MAP of String literals.");


Suggested change

"Config param of ML_PREDICT function should be a MAP of String literals.");

"Config paramameter of ML_PREDICT function should be a MAP data type consisting string literals.");

lihaosky added 4 commits August 18, 2025 20:43

[FLINK-38258][model] support model args in ptf

6b21de3

fix

9f527fd

fix

a024491

fix

ed213d6

airlock-confluentinc bot force-pushed the model-ptf branch from 2eeadb8 to ed213d6 Compare August 19, 2025 03:43

davidradl reviewed Aug 27, 2025

View reviewed changes

github-actions bot added community-reviewed PR has been reviewed by the community. and removed community-reviewed PR has been reviewed by the community. labels Aug 27, 2025

fix

3d56b5e

github-actions bot added community-reviewed PR has been reviewed by the community. and removed community-reviewed PR has been reviewed by the community. labels Sep 2, 2025

twalthr reviewed Sep 3, 2025

View reviewed changes

github-actions bot added community-reviewed PR has been reviewed by the community. and removed community-reviewed PR has been reviewed by the community. labels Sep 3, 2025

address comments

c2a61c5

github-actions bot removed the community-reviewed PR has been reviewed by the community. label Sep 4, 2025

github-actions bot added the community-reviewed PR has been reviewed by the community. label Sep 4, 2025

fix tests

ce04509

lihaosky commented Sep 4, 2025

View reviewed changes

trigger CI

94e4553

github-actions bot added community-reviewed PR has been reviewed by the community. and removed community-reviewed PR has been reviewed by the community. labels Sep 5, 2025

twalthr reviewed Sep 11, 2025

View reviewed changes

	* Output data type produced by the passed model. Extracting type from PTF class definition is
	* Output data type produced by the passed model.

	anyCauseMatches("Disabling uid/time attributes is not supported for PTF."));
	anyCauseMatches("Disabling system arguments is not supported for user-defined PTF yet."));

	.satisfies(anyCauseMatches("Disabling system arguments is not supported for PTF."));
	.satisfies(anyCauseMatches("Disabling system arguments is not supported for user-defined PTF."));

	"Config param of ML_PREDICT function should be a MAP of String literals.");
	"Config paramameter of ML_PREDICT function should be a MAP data type consisting string literals.");

[FLINK-38258][model] support model args in ptf and add ml_predict builtin ptf #26924

Are you sure you want to change the base?

[FLINK-38258][model] support model args in ptf and add ml_predict builtin ptf #26924

Uh oh!

Conversation

lihaosky commented Aug 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What is the purpose of the change

Brief change log

Verifying this change

Does this pull request potentially affect one of the following parts:

Documentation

Uh oh!

flinkbot commented Aug 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CI report:

Uh oh!

davidradl Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

twalthr left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lihaosky commented Sep 4, 2025

Uh oh!

twalthr left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lihaosky commented Aug 18, 2025 •

edited

Loading

flinkbot commented Aug 18, 2025 •

edited

Loading

davidradl Aug 27, 2025 •

edited

Loading