This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 700cd617fd3 [FLINK-39163][python] Support ML_PREDICT in Python API
(#27710)
700cd617fd3 is described below
commit 700cd617fd35b44683af7ed57f495158ac13990c
Author: Biao Geng <[email protected]>
AuthorDate: Tue Oct 13 23:22:30 2026 +0800
[FLINK-39163][python] Support ML_PREDICT in Python API (#27710)
---
.../docs/reference/pyflink.table/index.rst | 1 +
.../pyflink.table/{index.rst => model.rst} | 31 ++-
.../reference/pyflink.table/table_environment.rst | 2 +
flink-python/pyflink/table/__init__.py | 4 +
flink-python/pyflink/table/model.py | 113 ++++++++++
flink-python/pyflink/table/table_environment.py | 22 ++
flink-python/pyflink/table/tests/test_model.py | 187 ++++++++++++++++
.../tests/test_table_environment_completeness.py | 1 -
.../flink/table/utils/TestValuesModelFactory.java | 235 +++++++++++++++++++++
.../org.apache.flink.table.factories.Factory | 1 +
10 files changed, 580 insertions(+), 17 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.table/index.rst
b/flink-python/docs/reference/pyflink.table/index.rst
index 326fb86327b..3aef76aeb81 100644
--- a/flink-python/docs/reference/pyflink.table/index.rst
+++ b/flink-python/docs/reference/pyflink.table/index.rst
@@ -32,6 +32,7 @@ This page gives an overview of all public PyFlink Table API.
expressions
udf
descriptors
+ model
statement_set
catalog
compiled_plans
diff --git a/flink-python/docs/reference/pyflink.table/index.rst
b/flink-python/docs/reference/pyflink.table/model.rst
similarity index 69%
copy from flink-python/docs/reference/pyflink.table/index.rst
copy to flink-python/docs/reference/pyflink.table/model.rst
index 326fb86327b..ed10bbf01af 100644
--- a/flink-python/docs/reference/pyflink.table/index.rst
+++ b/flink-python/docs/reference/pyflink.table/model.rst
@@ -16,22 +16,21 @@
limitations under the License.
################################################################################
-=============
-PyFlink Table
-=============
+=====
+Model
+=====
-This page gives an overview of all public PyFlink Table API.
+A model registered in the catalog can be retrieved with
+:meth:`~pyflink.table.TableEnvironment.from_model`. Use
+:meth:`~pyflink.table.Model.predict` to apply it to a table, optionally passing
+runtime options such as asynchronous execution settings.
-.. toctree::
- :maxdepth: 2
+.. currentmodule:: pyflink.table
- table_environment
- table
- data_types
- window
- expressions
- udf
- descriptors
- statement_set
- catalog
- compiled_plans
+.. autosummary::
+ :toctree: api/
+
+ Model
+ Model.predict
+ Model.get_resolved_input_schema
+ Model.get_resolved_output_schema
diff --git a/flink-python/docs/reference/pyflink.table/table_environment.rst
b/flink-python/docs/reference/pyflink.table/table_environment.rst
index 367a4323cc1..759e5825542 100644
--- a/flink-python/docs/reference/pyflink.table/table_environment.rst
+++ b/flink-python/docs/reference/pyflink.table/table_environment.rst
@@ -181,6 +181,7 @@ keyword, thus must be escaped) in a catalog named 'cat.1'
and database named 'db
TableEnvironment.from_elements
TableEnvironment.from_pandas
TableEnvironment.from_path
+ TableEnvironment.from_model
TableEnvironment.get_catalog
TableEnvironment.get_config
TableEnvironment.get_current_catalog
@@ -239,6 +240,7 @@ StreamTableEnvironment
StreamTableEnvironment.from_elements
StreamTableEnvironment.from_pandas
StreamTableEnvironment.from_path
+ StreamTableEnvironment.from_model
StreamTableEnvironment.from_data_stream
StreamTableEnvironment.from_changelog_stream
StreamTableEnvironment.get_catalog
diff --git a/flink-python/pyflink/table/__init__.py
b/flink-python/pyflink/table/__init__.py
index 6fc10c9393b..53ad751fcd8 100644
--- a/flink-python/pyflink/table/__init__.py
+++ b/flink-python/pyflink/table/__init__.py
@@ -118,6 +118,8 @@ from pyflink.table.data_view import DataView, ListView,
MapView
from pyflink.table.environment_settings import EnvironmentSettings
from pyflink.table.explain_detail import ExplainDetail
from pyflink.table.expression import Expression
+from pyflink.table.model import Model
+from pyflink.table.model_descriptor import ModelDescriptor
from pyflink.table.module import Module, ModuleEntry
from pyflink.table.plan_reference import PlanReference
from pyflink.table.result_kind import ResultKind
@@ -159,6 +161,8 @@ __all__ = [
'TableDescriptor',
'FormatDescriptor',
'Schema',
+ 'Model',
+ 'ModelDescriptor',
'Module',
'ModuleEntry',
'SqlDialect',
diff --git a/flink-python/pyflink/table/model.py
b/flink-python/pyflink/table/model.py
new file mode 100644
index 00000000000..485bf3a64b8
--- /dev/null
+++ b/flink-python/pyflink/table/model.py
@@ -0,0 +1,113 @@
+################################################################################
+# 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.
+################################################################################
+from typing import List, Dict, Optional
+
+from pyflink.java_gateway import get_gateway
+from pyflink.table.catalog import ResolvedSchema
+from pyflink.table.table import Table
+from pyflink.util.api_stability_decorators import PublicEvolving
+
+__all__ = ['Model']
+
+
+@PublicEvolving()
+class Model(object):
+ """
+ A model object represents a machine learning model registered in the
catalog.
+ It provides a way to perform predictions on input tables using the predict
method.
+
+ .. versionadded:: 2.4.0
+ """
+
+ def __init__(self, j_model):
+ self._j_model = j_model
+
+ @PublicEvolving()
+ def predict(self,
+ input_table: Table,
+ input_columns: List[str],
+ options: Optional[Dict[str, str]] = None) -> Table:
+ """
+ Performs prediction on the given input table using this model.
+
+ Example:
+ ::
+
+ >>> model = t_env.from_model("my_model")
+ >>> input_table = t_env.from_path("src")
+ >>> result_table = model.predict(input_table, ["user_comment"])
+ >>> result_table.execute_insert("prediction_result").wait()
+
+ :param input_table: The input table to perform prediction on.
+ :param input_columns: The column names from the input table to use as
features.
+ :param options: Optional runtime configuration for the prediction.
+ :return: The resulting table with prediction results appended.
+
+ .. versionadded:: 2.4.0
+ """
+ gateway = get_gateway()
+
+ # Convert input columns to Java ColumnList
+ j_input_columns = gateway.jvm.java.util.ArrayList()
+ for col_item in input_columns:
+ if isinstance(col_item, str):
+ j_input_columns.add(col_item)
+ else:
+ j_input_columns.add(str(col_item))
+ j_column_list =
gateway.jvm.org.apache.flink.types.ColumnList.of(j_input_columns)
+
+ # Convert config to Java map if provided
+ if options is not None:
+ j_config = gateway.jvm.java.util.HashMap()
+ for key, value in options.items():
+ j_config.put(str(key), str(value))
+ else:
+ j_config = gateway.jvm.java.util.Collections.emptyMap()
+
+ j_result_table = self._j_model.predict(input_table._j_table,
j_column_list, j_config)
+ return Table(j_result_table, input_table._t_env)
+
+ @PublicEvolving()
+ def get_resolved_input_schema(self) -> 'ResolvedSchema':
+ """
+ Returns the resolved input schema of this model.
+
+ The input schema describes the structure and data types of the input
columns that the
+ model expects for inference operations.
+
+ :return: the resolved input schema.
+
+ .. versionadded:: 2.4.0
+ """
+ j_resolved_schema = self._j_model.getResolvedInputSchema()
+ return ResolvedSchema(j_resolved_schema=j_resolved_schema)
+
+ @PublicEvolving()
+ def get_resolved_output_schema(self) -> 'ResolvedSchema':
+ """
+ Returns the resolved output schema of this model.
+
+ The output schema describes the structure and data types of the output
columns that the
+ model produces during inference operations.
+
+ :return: the resolved output schema.
+
+ .. versionadded:: 2.4.0
+ """
+ j_resolved_schema = self._j_model.getResolvedOutputSchema()
+ return ResolvedSchema(j_resolved_schema=j_resolved_schema)
diff --git a/flink-python/pyflink/table/table_environment.py
b/flink-python/pyflink/table/table_environment.py
index a1d90c41db1..6d84c39ec0a 100644
--- a/flink-python/pyflink/table/table_environment.py
+++ b/flink-python/pyflink/table/table_environment.py
@@ -37,6 +37,7 @@ from pyflink.serializers import BatchedSerializer,
PickleSerializer
from pyflink.table import Table, EnvironmentSettings, Expression,
ExplainDetail, \
Module, ModuleEntry, Schema, ChangelogMode
from pyflink.table.catalog import Catalog, CatalogDescriptor
+from pyflink.table.model import Model
from pyflink.table.model_descriptor import ModelDescriptor
from pyflink.table.compiled_plan import CompiledPlan
from pyflink.table.plan_reference import PlanReference
@@ -1248,6 +1249,27 @@ class TableEnvironment(object):
"""
self._j_tenv.createModel(model_path,
model_descriptor._j_model_descriptor, ignore_if_exists)
+ @PublicEvolving()
+ def from_model(self, model_path: str) -> 'Model':
+ """
+ Gets a model registered in the catalog.
+
+ Example:
+ ::
+
+ >>> from pyflink.table import ModelDescriptor, Schema
+ >>> model = t_env.from_model("my_model")
+ >>> input_table = t_env.from_path("src")
+ >>> result_table = model.predict(input_table, ["user_comment"])
+
+ :param model_path: The path of the model to retrieve.
+ :return: The requested model instance.
+
+ .. versionadded:: 2.4.0
+ """
+ j_model = self._j_tenv.fromModel(model_path)
+ return Model(j_model)
+
def create_temporary_model(self,
model_path: str,
model_descriptor: ModelDescriptor,
diff --git a/flink-python/pyflink/table/tests/test_model.py
b/flink-python/pyflink/table/tests/test_model.py
new file mode 100644
index 00000000000..5e1a64d6458
--- /dev/null
+++ b/flink-python/pyflink/table/tests/test_model.py
@@ -0,0 +1,187 @@
+################################################################################
+# 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.
+################################################################################
+from pyflink.common import Row
+from pyflink.table import ModelDescriptor, Schema, DataTypes
+from pyflink.table.catalog import ResolvedSchema
+from pyflink.testing.test_case_utils import PyFlinkStreamTableTestCase
+
+
+class ModelTest(PyFlinkStreamTableTestCase):
+
+ def test_model_predict(self):
+ # Model input: (id BIGINT), output: (x STRING, y INT, z STRING)
+ # Factory generates output ("x" + id, id, "z" + id) based on input id
+ model_descriptor = ModelDescriptor.for_provider("python-values") \
+ .input_schema(Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .build()) \
+ .output_schema(Schema.new_builder()
+ .column("x", DataTypes.STRING())
+ .column("y", DataTypes.INT())
+ .column("z", DataTypes.STRING())
+ .build()) \
+ .build()
+
+ self.t_env.create_temporary_model("my_model", model_descriptor, True)
+
+ model = self.t_env.from_model("my_model")
+ self.assertIsNotNone(model)
+
+ data = [
+ (1, 12, "Julian"),
+ (2, 15, "Hello"),
+ (3, 15, "Fabian"),
+ (8, 11, "Hello world"),
+ (9, 12, "Hello world!")
+ ]
+
+ schema = DataTypes.ROW([
+ DataTypes.FIELD("id", DataTypes.BIGINT()),
+ DataTypes.FIELD("len", DataTypes.INT()),
+ DataTypes.FIELD("content", DataTypes.STRING())
+ ])
+
+ input_table = self.t_env.from_elements(data, schema)
+ prediction_table = model.predict(input_table, ["id"])
+ self.assertIsNotNone(prediction_table)
+ result = [i for i in prediction_table.execute().collect()]
+ expected = [Row(1, 12, 'Julian', 'x1', 1, 'z1'),
+ Row(2, 15, 'Hello', 'x2', 2, 'z2'),
+ Row(3, 15, 'Fabian', 'x3', 3, 'z3'),
+ Row(8, 11, 'Hello world', 'x8', 8, 'z8'),
+ Row(9, 12, 'Hello world!', 'x9', 9, 'z9')]
+ self.assertEqual(expected, result)
+
+ def test_model_predict_async(self):
+ model_descriptor = ModelDescriptor.for_provider("python-values") \
+ .input_schema(Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .build()) \
+ .output_schema(Schema.new_builder()
+ .column("x", DataTypes.STRING())
+ .column("y", DataTypes.INT())
+ .column("z", DataTypes.STRING())
+ .build()) \
+ .option("async", "true") \
+ .build()
+
+ self.t_env.create_temporary_model("my_async_model", model_descriptor,
True)
+
+ model = self.t_env.from_model("my_async_model")
+ self.assertIsNotNone(model)
+
+ data = [
+ (1, 12, "Julian"),
+ (2, 15, "Hello"),
+ (3, 15, "Fabian"),
+ (8, 11, "Hello world"),
+ (9, 12, "Hello world!")
+ ]
+
+ schema = DataTypes.ROW([
+ DataTypes.FIELD("id", DataTypes.BIGINT()),
+ DataTypes.FIELD("len", DataTypes.INT()),
+ DataTypes.FIELD("content", DataTypes.STRING())
+ ])
+
+ input_table = self.t_env.from_elements(data, schema)
+ prediction_table = model.predict(input_table, ["id"])
+ self.assertIsNotNone(prediction_table)
+ result = [i for i in prediction_table.execute().collect()]
+ expected = [Row(1, 12, 'Julian', 'x1', 1, 'z1'),
+ Row(2, 15, 'Hello', 'x2', 2, 'z2'),
+ Row(3, 15, 'Fabian', 'x3', 3, 'z3'),
+ Row(8, 11, 'Hello world', 'x8', 8, 'z8'),
+ Row(9, 12, 'Hello world!', 'x9', 9, 'z9')]
+ # Async may return results in different order
+ self.assertEqual(sorted(expected, key=lambda r: r[0]),
+ sorted(result, key=lambda r: r[0]))
+
+ def test_model_predict_with_options(self):
+ model_descriptor = ModelDescriptor.for_provider("python-values") \
+ .input_schema(Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .build()) \
+ .output_schema(Schema.new_builder()
+ .column("x", DataTypes.STRING())
+ .column("y", DataTypes.INT())
+ .column("z", DataTypes.STRING())
+ .build()) \
+ .build()
+
+ self.t_env.create_temporary_model("opt_model", model_descriptor, True)
+ model = self.t_env.from_model("opt_model")
+
+ data = [
+ (1, 12, "Julian"),
+ (2, 15, "Hello"),
+ (3, 15, "Fabian"),
+ ]
+
+ schema = DataTypes.ROW([
+ DataTypes.FIELD("id", DataTypes.BIGINT()),
+ DataTypes.FIELD("len", DataTypes.INT()),
+ DataTypes.FIELD("content", DataTypes.STRING())
+ ])
+
+ input_table = self.t_env.from_elements(data, schema)
+ # Pass options (async=false is a no-op for sync model, just testing
the API)
+ prediction_table = model.predict(input_table, ["id"], {"async":
"false"})
+ self.assertIsNotNone(prediction_table)
+ result = [i for i in prediction_table.execute().collect()]
+ expected = [Row(1, 12, 'Julian', 'x1', 1, 'z1'),
+ Row(2, 15, 'Hello', 'x2', 2, 'z2'),
+ Row(3, 15, 'Fabian', 'x3', 3, 'z3')]
+ self.assertEqual(expected, result)
+
+ def test_model_schema(self):
+ model_descriptor = ModelDescriptor.for_provider("python-values") \
+ .input_schema(Schema.new_builder()
+ .column("id", DataTypes.BIGINT())
+ .build()) \
+ .output_schema(Schema.new_builder()
+ .column("x", DataTypes.STRING())
+ .column("y", DataTypes.INT())
+ .column("z", DataTypes.STRING())
+ .build()) \
+ .build()
+
+ self.t_env.create_temporary_model("schema_model", model_descriptor,
True)
+ model = self.t_env.from_model("schema_model")
+
+ # Test input schema
+ input_schema = model.get_resolved_input_schema()
+ self.assertIsNotNone(input_schema)
+ self.assertIsInstance(input_schema, ResolvedSchema)
+
+ # Test output schema
+ output_schema = model.get_resolved_output_schema()
+ self.assertIsNotNone(output_schema)
+ self.assertIsInstance(output_schema, ResolvedSchema)
+
+
+if __name__ == '__main__':
+ import unittest
+
+ try:
+ import xmlrunner
+
+ testRunner = xmlrunner.XMLTestRunner(output='target/test-reports')
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git
a/flink-python/pyflink/table/tests/test_table_environment_completeness.py
b/flink-python/pyflink/table/tests/test_table_environment_completeness.py
index 2e0cc5c2b39..48c6369bae6 100644
--- a/flink-python/pyflink/table/tests/test_table_environment_completeness.py
+++ b/flink-python/pyflink/table/tests/test_table_environment_completeness.py
@@ -44,7 +44,6 @@ class
TableEnvironmentAPICompletenessTests(PythonAPICompletenessTestCase, PyFlin
"from",
"registerFunction",
"fromCall",
- "fromModel",
}
diff --git
a/flink-python/src/test/java/org/apache/flink/table/utils/TestValuesModelFactory.java
b/flink-python/src/test/java/org/apache/flink/table/utils/TestValuesModelFactory.java
new file mode 100644
index 00000000000..aeb751846c7
--- /dev/null
+++
b/flink-python/src/test/java/org/apache/flink/table/utils/TestValuesModelFactory.java
@@ -0,0 +1,235 @@
+/*
+ * 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.utils;
+
+import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.ConfigOptions;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.StringData;
+import org.apache.flink.table.data.conversion.RowRowConverter;
+import org.apache.flink.table.factories.FactoryUtil;
+import org.apache.flink.table.factories.ModelProviderFactory;
+import org.apache.flink.table.functions.AsyncPredictFunction;
+import org.apache.flink.table.functions.FunctionContext;
+import org.apache.flink.table.functions.PredictFunction;
+import org.apache.flink.table.ml.AsyncPredictRuntimeProvider;
+import org.apache.flink.table.ml.ModelProvider;
+import org.apache.flink.table.ml.PredictRuntimeProvider;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+/**
+ * Test values model factory for Python integration tests.
+ *
+ * <p>This factory generates predictable output based on input: for input with
first column as id
+ * (BIGINT), it generates output ("x" + id, id, "z" + id).
+ */
+public class TestValuesModelFactory implements ModelProviderFactory {
+
+ public static final ConfigOption<Boolean> ASYNC =
+ ConfigOptions.key("async").booleanType().defaultValue(false);
+
+ public static final ConfigOption<Integer> LATENCY =
+ ConfigOptions.key("latency")
+ .intType()
+ .noDefaultValue()
+ .withDescription(
+ "Latency in milliseconds for async prediction. "
+ + "If not set, defaults to 0ms (no
delay).");
+
+ @Override
+ public ModelProvider createModelProvider(Context context) {
+ FactoryUtil.ModelProviderFactoryHelper helper =
+ FactoryUtil.createModelProviderFactoryHelper(this, context);
+ helper.validate();
+
+ RowRowConverter outputConverter =
+ RowRowConverter.create(
+ context.getCatalogModel()
+ .getResolvedOutputSchema()
+ .toPhysicalRowDataType());
+
+ if (helper.getOptions().get(ASYNC)) {
+ return new AsyncValuesModelProvider(
+ new GeneratingPredictFunction(outputConverter),
+ new AsyncGeneratingPredictFunction(
+ outputConverter,
helper.getOptions().getOptional(LATENCY).orElse(0)));
+ } else {
+ return new ValuesModelProvider(new
GeneratingPredictFunction(outputConverter));
+ }
+ }
+
+ @Override
+ public String factoryIdentifier() {
+ return "python-values";
+ }
+
+ @Override
+ public Set<ConfigOption<?>> requiredOptions() {
+ return Collections.emptySet();
+ }
+
+ @Override
+ public Set<ConfigOption<?>> optionalOptions() {
+ return new HashSet<>(Arrays.asList(ASYNC, LATENCY));
+ }
+
+ /** Values model provider implementing sync predict. */
+ public static class ValuesModelProvider implements PredictRuntimeProvider {
+
+ private final PredictFunction function;
+
+ public ValuesModelProvider(PredictFunction function) {
+ this.function = function;
+ }
+
+ @Override
+ public PredictFunction createPredictFunction(Context context) {
+ return function;
+ }
+
+ @Override
+ public ModelProvider copy() {
+ return new ValuesModelProvider(function);
+ }
+ }
+
+ /** Async values model provider implementing both sync and async predict.
*/
+ public static class AsyncValuesModelProvider
+ implements AsyncPredictRuntimeProvider, PredictRuntimeProvider {
+
+ private final PredictFunction predictFunction;
+ private final AsyncPredictFunction asyncPredictFunction;
+
+ public AsyncValuesModelProvider(
+ PredictFunction predictFunction, AsyncPredictFunction
asyncPredictFunction) {
+ this.predictFunction = predictFunction;
+ this.asyncPredictFunction = asyncPredictFunction;
+ }
+
+ @Override
+ public AsyncPredictFunction createAsyncPredictFunction(Context
context) {
+ return asyncPredictFunction;
+ }
+
+ @Override
+ public PredictFunction createPredictFunction(Context context) {
+ return predictFunction;
+ }
+
+ @Override
+ public ModelProvider copy() {
+ return new AsyncValuesModelProvider(predictFunction,
asyncPredictFunction);
+ }
+ }
+
+ /**
+ * Predict function that generates output based on input.
+ *
+ * <p>For input row with first column value as id (BIGINT), generates:
("x" + id, id, "z" + id).
+ */
+ public static class GeneratingPredictFunction extends PredictFunction {
+
+ private final RowRowConverter outputConverter;
+
+ public GeneratingPredictFunction(RowRowConverter outputConverter) {
+ this.outputConverter = outputConverter;
+ }
+
+ @Override
+ public void open(FunctionContext context) throws Exception {
+ super.open(context);
+ outputConverter.open(context.getUserCodeClassLoader());
+ }
+
+ @Override
+ public Collection<RowData> predict(RowData features) {
+ // Extract the first field as id
+ long id = features.getLong(0);
+ // Generate output: ("x" + id, id, "z" + id)
+ GenericRowData output =
+ GenericRowData.of(
+ StringData.fromString("x" + id),
+ (int) id,
+ StringData.fromString("z" + id));
+ return Collections.singletonList(output);
+ }
+ }
+
+ /** Async predict function that generates output with configurable
latency. */
+ public static class AsyncGeneratingPredictFunction extends
AsyncPredictFunction {
+
+ private final RowRowConverter outputConverter;
+ private final int latency;
+
+ private transient ExecutorService executors;
+
+ public AsyncGeneratingPredictFunction(RowRowConverter outputConverter,
int latency) {
+ this.outputConverter = outputConverter;
+ this.latency = latency;
+ }
+
+ @Override
+ public void open(FunctionContext context) throws Exception {
+ super.open(context);
+ outputConverter.open(context.getUserCodeClassLoader());
+ executors = Executors.newFixedThreadPool(5);
+ }
+
+ @Override
+ public CompletableFuture<Collection<RowData>> asyncPredict(RowData
features) {
+ return CompletableFuture.supplyAsync(
+ () -> {
+ try {
+ if (latency > 0) {
+ Thread.sleep(latency);
+ }
+ // Extract the first field as id
+ long id = features.getLong(0);
+ // Generate output: ("x" + id, id, "z" + id)
+ GenericRowData output =
+ GenericRowData.of(
+ StringData.fromString("x" + id),
+ (int) id,
+ StringData.fromString("z" + id));
+ return Collections.singletonList(output);
+ } catch (Exception e) {
+ throw new RuntimeException("Prediction
interrupted.", e);
+ }
+ },
+ executors);
+ }
+
+ @Override
+ public void close() throws Exception {
+ super.close();
+ if (executors != null) {
+ executors.shutdownNow();
+ }
+ }
+ }
+}
diff --git
a/flink-python/src/test/resources/META-INF/services/org.apache.flink.table.factories.Factory
b/flink-python/src/test/resources/META-INF/services/org.apache.flink.table.factories.Factory
index 679e7f4a8c6..45deae20af8 100644
---
a/flink-python/src/test/resources/META-INF/services/org.apache.flink.table.factories.Factory
+++
b/flink-python/src/test/resources/META-INF/services/org.apache.flink.table.factories.Factory
@@ -15,3 +15,4 @@
org.apache.flink.table.utils.TestCollectionTableFactory
org.apache.flink.table.utils.TestingSinkTableFactory
+org.apache.flink.table.utils.TestValuesModelFactory
