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 6144aa7367d [FLINK-40440][python] Add model prediction support to the
DataFrame API (#29478)
6144aa7367d is described below
commit 6144aa7367d9c56f600fe86afa08d576e6b0e8c0
Author: Liu Liu <[email protected]>
AuthorDate: Thu Oct 15 17:44:27 2026 +0800
[FLINK-40440][python] Add model prediction support to the DataFrame API
(#29478)
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 1 +
.../docs/reference/pyflink.dataframe/llm.rst | 39 +++
flink-python/pyflink/dataframe/dataframe.py | 17 ++
flink-python/pyflink/dataframe/llm.py | 274 ++++++++++++++++-
flink-python/pyflink/dataframe/tests/test_llm.py | 327 +++++++++++++++++++++
.../dataframe/tests/test_llm_integration.py | 88 ++++++
flink-python/pyflink/table/model.py | 6 +-
flink-python/pyflink/table/table_environment.py | 27 +-
flink-python/pyflink/table/tests/test_model.py | 32 +-
9 files changed, 787 insertions(+), 24 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index 1fc52491ec5..e31a2bf176d 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -161,6 +161,7 @@ Properties
DataFrame.schema
DataFrame.columns
+ DataFrame.llm
Results
-------
diff --git a/flink-python/docs/reference/pyflink.dataframe/llm.rst
b/flink-python/docs/reference/pyflink.dataframe/llm.rst
index b51f9c45eff..ba96681a3ac 100644
--- a/flink-python/docs/reference/pyflink.dataframe/llm.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/llm.rst
@@ -48,6 +48,43 @@ For other providers available in the application's
classpath, use ``GenericProvi
with provider-specific option names and string values, or subclass
``ModelProvider``
to create a reusable configuration.
+Prediction
+----------
+
+Use ``df.llm.predict`` to infer a result for each row and append the output
columns.
+Supply one or more column names in model input order. The provider determines
which
+input types and number of columns it supports.
+
+.. code-block:: python
+
+ pf.set_model_provider("chat", pf.OpenAIProvider(
+ endpoint="https://api.openai.com/v1/chat/completions";, api_key="key",
retry_num=3))
+ questions = pf.from_records([{"question": "What is Apache Flink?"}])
+ answers = questions.llm.predict(
+ "question", provider="chat", model="gpt-6-luna")
+
+The default output is one ``STRING`` column named ``output``. Set
``output_type``
+to a SQL type string or a ``DataType`` to declare the provider's native output
schema.
+
+A ``DataType.struct(...)`` expands its top-level fields into named output
columns.
+Output field names must be unique and must not already occur in the input
DataFrame.
+The provider validates the output schema it supports.
+
+.. code-block:: python
+
+ # Name the prediction column "answer" instead of the default "output".
+ answers = questions.llm.predict(
+ "question", provider="chat", model="gpt-6-luna",
+ output_type=pf.DataType.struct({"answer": pf.DataType.string()}))
+
+You can pass in extra arguments to override the selected provider's
constructor settings:
+
+.. code-block:: python
+
+ # Override the provider's retry_num parameter for this call.
+ answers = questions.llm.predict(
+ "question", provider="chat", model="gpt-6-luna", retry_num=0)
+
.. currentmodule:: pyflink.dataframe
.. autosummary::
@@ -60,3 +97,5 @@ to create a reusable configuration.
set_model_provider
set_default_model_provider
list_model_providers
+
+.. automethod:: pyflink.dataframe.llm.LLMAccessor.predict
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 1d3eedd0973..d54f3538a5f 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -41,6 +41,7 @@ if TYPE_CHECKING:
import pyarrow
from pyflink.dataframe.udf import _DataTypeLike
from pyflink.dataframe.udtf import _DataFrameUDTFCall,
_DataFrameUDTFWrapper
+ from pyflink.dataframe.llm import LLMAccessor
from pyflink.table.table_environment import TableEnvironment
from pyflink.table.table_schema import TableSchema
@@ -2639,6 +2640,22 @@ class DataFrame:
# ======================== Properties ========================
+ @property
+ @PublicEvolving()
+ def llm(self) -> "LLMAccessor":
+ """
+ Access model inference operations for this DataFrame.
+
+ Example::
+
+ >>> result = df.llm.predict("question", provider="chat",
model="gpt-6-luna")
+
+ .. versionadded:: 2.4.0
+ """
+ from pyflink.dataframe.llm import LLMAccessor
+
+ return LLMAccessor(self)
+
@property
@PublicEvolving()
def schema(self) -> "TableSchema":
diff --git a/flink-python/pyflink/dataframe/llm.py
b/flink-python/pyflink/dataframe/llm.py
index 122707dafb4..d9e404eb796 100644
--- a/flink-python/pyflink/dataframe/llm.py
+++ b/flink-python/pyflink/dataframe/llm.py
@@ -19,9 +19,16 @@
from abc import ABC, abstractmethod
from collections.abc import Mapping as MappingABC
import json
-from typing import Dict, List, Literal, Mapping, Optional, Tuple, overload
+from typing import Callable, Dict, List, Literal, Mapping, Optional, Tuple,
Union, cast, overload
+from pyflink.dataframe.dataframe import DataFrame
+from pyflink.dataframe.datatype import DataType
from pyflink.dataframe.validation import _require_int, _require_number
+from pyflink.table.model_descriptor import ModelDescriptor
+from pyflink.table.schema import Schema
+from pyflink.table.types import (
+ DataType as TableDataType, DataTypes, RowType, _from_java_data_type,
+)
from pyflink.util.api_stability_decorators import PublicEvolving
__all__ = [
@@ -140,12 +147,12 @@ class ModelProvider(ABC):
@PublicEvolving()
class OpenAIProvider(ModelProvider):
"""
- Configure an OpenAI-compatible service for chat completions or embeddings.
+ Configure an OpenAI-compatible service.
- The full endpoint URL selects the task: use a URL ending in
``/chat/completions``
- for chat or ``/embeddings`` for embedding vectors. Chat options such as
- ``system_prompt`` and ``temperature`` apply to chat requests;
``dimension`` applies
- to embeddings. Supply a model name here or when calling an AI function.
+ For chat completions, use a full endpoint URL ending in
``/chat/completions``.
+ Supply a model name here or when calling ``df.llm.predict``. Chat accepts
one
+ string input column and returns one string output column, matching the
default
+ ``output_type``.
All optional parameters default to ``None``. Omitted options use the Flink
defaults
described below, or the service's defaults where Flink defines none.
@@ -175,8 +182,8 @@ class OpenAIProvider(ModelProvider):
supplied).
:param response_format: Chat response format: ``text`` or ``json_object``.
Default: ``None`` (service default).
- :param dimension: Number of elements in each embedding vector. Default:
``None``
- (the embedding model's default dimension).
+ :param dimension: Requested output vector dimension, when supported by the
service.
+ Default: ``None`` (service default).
:param max_context_size: Maximum number of input context tokens before
applying
``context_overflow_action``. Default: ``None``
(no Flink
context limit).
@@ -201,13 +208,23 @@ class OpenAIProvider(ModelProvider):
>>> chat = pf.OpenAIProvider(
... endpoint="https://api.openai.com/v1/chat/completions";,
api_key="key")
>>> pf.set_model_provider("chat", chat)
- >>> embeddings = pf.OpenAIProvider(
- ... endpoint="https://api.openai.com/v1/embeddings";, api_key="key")
- >>> pf.set_model_provider("embed", embeddings)
+ >>> questions = pf.from_records([{"question": "What is Apache
Flink?"}])
+ >>> answers = questions.llm.predict("question", provider="chat",
model="gpt-6-luna")
.. versionadded:: 2.4.0
"""
+ _OPTION_MAP = {
+ "endpoint": "endpoint", "api_key": "api-key", "model": "model",
+ "system_prompt": "system-prompt", "temperature": "temperature",
"top_p": "top-p",
+ "max_tokens": "max-tokens", "stop": "stop", "presence_penalty":
"presence-penalty",
+ "n": "n", "seed": "seed", "response_format": "response-format",
"dimension": "dimension",
+ "max_context_size": "max-context-size",
+ "context_overflow_action": "context-overflow-action",
+ "error_handling_strategy": "error-handling-strategy", "retry_num":
"retry-num",
+ "retry_fallback_strategy": "retry-fallback-strategy",
+ }
+
def __init__(
self, endpoint: str, api_key: str, *, model: Optional[str] = None,
system_prompt: Optional[str] = None,
@@ -230,8 +247,7 @@ class OpenAIProvider(ModelProvider):
retry_fallback_strategy: Optional[Literal["FAILOVER", "IGNORE"]] =
None,
options: Optional[Mapping[str, str]] = None,
):
- _validate_name(endpoint, "endpoint")
- _validate_name(api_key, "api_key")
+ self._validate_required_options({"endpoint": endpoint, "api-key":
api_key})
typed_options: Dict[str, object] = {
"endpoint": endpoint, "api-key": api_key, "model": model,
"system-prompt": system_prompt, "temperature": temperature,
"top-p": top_p,
@@ -265,7 +281,14 @@ class OpenAIProvider(ModelProvider):
):
if number is not None:
_require_number(number, name, lower_bound, maximum=upper_bound)
- self._options = _merge_options(options,
_stringify_options(typed_options))
+ self._parameters = {name: typed_options[key] for name, key in
self._OPTION_MAP.items()}
+ self._typed_options = _stringify_options(typed_options)
+ self._options = _merge_options(options, self._typed_options)
+
+ @staticmethod
+ def _validate_required_options(options: Mapping[str, object]) -> None:
+ _validate_name(options.get("endpoint"), "endpoint")
+ _validate_name(options.get("api-key"), "api_key")
def provider_identifier(self) -> str:
"""Return ``openai``.
@@ -292,6 +315,9 @@ class TritonProvider(ModelProvider):
fallback options control how failed requests are handled; health checks
and the
circuit breaker can reduce requests to an unavailable server.
+ Set ``output_type`` in ``df.llm.predict`` to match the deployed model's
output
+ type. The default is ``STRING``; other output types must be declared
explicitly.
+
All optional parameters default to ``None``. Omitted options use the Flink
defaults
described below. Specify durations as strings such as ``"30 s"`` or ``"100
ms"``.
@@ -352,10 +378,29 @@ class TritonProvider(ModelProvider):
... endpoint="http://localhost:8000";, model_name="classifier",
... flatten_batch_dim=True, max_retries=2, default_value="-1")
>>> pf.set_model_provider("classifier", provider)
+ >>> features = pf.from_records([{"features": [0.1, 0.9]}])
+ >>> predictions = features.llm.predict(
+ ... "features", provider="classifier",
output_type=pf.DataType.int32())
.. versionadded:: 2.4.0
"""
+ _OPTION_MAP = {
+ "endpoint": "endpoint", "model_name": "model-name", "model_version":
"model-version",
+ "timeout": "timeout", "flatten_batch_dim": "flatten-batch-dim",
"priority": "priority",
+ "sequence_id": "sequence-id", "sequence_start": "sequence-start",
+ "sequence_end": "sequence-end", "compression": "compression",
"auth_token": "auth-token",
+ "custom_headers": "custom-headers",
+ "max_retries": "max-retries", "retry_initial_backoff":
"retry-initial-backoff",
+ "retry_max_backoff": "retry-max-backoff", "default_value":
"default-value",
+ "health_check_enabled": "health-check-enabled",
+ "health_check_interval": "health-check-interval",
+ "circuit_breaker_enabled": "circuit-breaker-enabled",
+ "circuit_breaker_failure_threshold":
"circuit-breaker-failure-threshold",
+ "circuit_breaker_timeout": "circuit-breaker-timeout",
+ "circuit_breaker_half_open_requests":
"circuit-breaker-half-open-requests",
+ }
+
def __init__(
self, endpoint: str, *, model_name: Optional[str] = None,
model_version: Optional[str] = None,
@@ -380,7 +425,7 @@ class TritonProvider(ModelProvider):
circuit_breaker_half_open_requests: Optional[int] = None,
options: Optional[Mapping[str, str]] = None,
):
- _validate_name(endpoint, "endpoint")
+ self._validate_required_options({"endpoint": endpoint})
typed_options: Dict[str, object] = {
"endpoint": endpoint, "model-name": model_name, "model-version":
model_version,
"timeout": timeout, "flatten-batch-dim": flatten_batch_dim,
"priority": priority,
@@ -413,7 +458,15 @@ class TritonProvider(ModelProvider):
if circuit_breaker_failure_threshold is not None:
_require_number(circuit_breaker_failure_threshold,
"circuit_breaker_failure_threshold",
0, maximum=1, include_minimum=False)
- self._options = _merge_options(options,
_stringify_options(typed_options))
+ self._parameters = {name: typed_options[key] for name, key in
self._OPTION_MAP.items()}
+ self._parameters["custom_headers"] = (
+ dict(custom_headers) if custom_headers is not None else None)
+ self._typed_options = _stringify_options(typed_options)
+ self._options = _merge_options(options, self._typed_options)
+
+ @staticmethod
+ def _validate_required_options(options: Mapping[str, object]) -> None:
+ _validate_name(options.get("endpoint"), "endpoint")
def provider_identifier(self) -> str:
"""Return ``triton``.
@@ -437,6 +490,9 @@ class TritonProvider(ModelProvider):
return dict(self._options)
+TYPED_PROVIDERS = (OpenAIProvider, TritonProvider)
+
+
@PublicEvolving()
class GenericProvider(ModelProvider):
"""
@@ -610,3 +666,189 @@ def _resolve_provider(name: Optional[str] = None) ->
ModelProvider:
"Multiple model providers are registered; call
set_default_model_provider() "
"or select a provider explicitly"
)
+
+
+def _override_options(
+ options: Mapping[str, str], overrides: Mapping[str, object],
+) -> Dict[str, str]:
+ result = dict(options)
+ for key, value in overrides.items():
+ if not isinstance(key, str):
+ raise TypeError("Option keys must be strings")
+ if value is None:
+ result.pop(key, None)
+ elif isinstance(value, str):
+ result[key] = value
+ else:
+ raise TypeError(f"Option {key!r} must have a string value or None")
+ return result
+
+
+def _provider_options(
+ provider: ModelProvider, overrides: Mapping[str, object],
+) -> Dict[str, str]:
+ configured = provider.to_options()
+ if not isinstance(configured, MappingABC):
+ raise TypeError("ModelProvider.to_options() must return a mapping of
strings to strings")
+ options = _merge_options(configured, {})
+ if type(provider) not in TYPED_PROVIDERS:
+ return _override_options(options, overrides)
+ if not overrides:
+ return options
+
+ typed_provider = cast(Union[OpenAIProvider, TritonProvider], provider)
+ parameters = dict(overrides)
+ raw_supplied = "options" in parameters
+ raw = parameters.pop("options", None)
+ if raw is not None and not isinstance(raw, MappingABC):
+ raise TypeError("options must be a mapping of strings to strings or
None")
+
+ for name in parameters:
+ if name not in typed_provider._OPTION_MAP:
+ raise TypeError(f"Unexpected provider parameter {name!r}")
+ constructor: Callable[..., ModelProvider] = type(typed_provider)
+ validated = constructor(**{**typed_provider._parameters,
**parameters}).to_options()
+ typed = {typed_provider._OPTION_MAP[name]:
validated.get(typed_provider._OPTION_MAP[name])
+ for name in parameters}
+ if raw_supplied and raw is None:
+ options = dict(typed_provider._typed_options)
+ elif raw is not None:
+ # Validate raw values even if an equal typed override supplies the
same key.
+ _override_options({}, raw)
+ for key in typed.keys() & raw.keys():
+ if typed[key] != raw[key]:
+ raise ValueError(f"conflicting values for option {key!r}")
+
+ options = _override_options(options, typed)
+ if raw is not None:
+ options = _override_options(options, raw)
+ typed_provider._validate_required_options(options)
+ return options
+
+
+def _prediction_output_columns(
+ output_type: Optional[Union[str, DataType]],
+) -> List[Tuple[str, TableDataType]]:
+ if output_type is None:
+ return [("output", DataTypes.STRING())]
+ if isinstance(output_type, str):
+ return [("output",
DataType._from_sql(output_type)._to_table_data_type())]
+ if not isinstance(output_type, DataType):
+ raise TypeError("output_type must be a SQL type string or a
pyflink.dataframe.DataType")
+ table_type = output_type._to_table_data_type()
+ if not isinstance(table_type, RowType):
+ return [("output", table_type)]
+ if not table_type.fields:
+ raise ValueError("Struct output_type must contain at least one field")
+ return [(field.name, field.data_type) for field in table_type.fields]
+
+
+@PublicEvolving()
+class LLMAccessor:
+ """Model inference operations accessed through :attr:`DataFrame.llm`.
+
+ .. versionadded:: 2.4.0
+ """
+
+ def __init__(self, dataframe: DataFrame):
+ self._df = dataframe
+
+ @PublicEvolving()
+ def predict(
+ self, *input_cols: str, provider: Optional[str] = None, model:
Optional[str] = None,
+ output_type: Optional[Union[str, DataType]] = None,
+ config: Optional[Mapping[str, str]] = None, **kwargs: object,
+ ) -> DataFrame:
+ """
+ Predict each row and append the model's output columns.
+
+ The selected provider validates input and output types and the number
of columns.
+
+ :param input_cols: One or more existing column names, in model input
order.
+ :param provider: Registered provider name. If omitted, use the
selected default
+ or the sole registered provider.
+ :param model: Override the provider's model name for this call.
+ :param output_type: SQL type string or :class:`DataType` declaring the
provider's
+ native output schema. Defaults to ``STRING``. A
scalar or container
+ type appends one column named ``output``.
+ ``DataType.struct(...)`` appends its top-level
fields as
+ named columns, preserving nested types. Output
names must
+ be unique and must not occur in the input
DataFrame.
+ :param config: Runtime configuration as string keys and values.
+ :param kwargs: Provider constructor overrides for this call. Built-ins
accept
+ their typed Python parameter names and merge
``options`` per key.
+ Custom providers accept raw option names and string
values.
+ ``None`` removes an inherited option; ``options=None``
clears
+ inherited raw options. The registered provider stays
unchanged.
+ :return: A new DataFrame containing the input and prediction columns.
+
+ Examples::
+
+ >>> import pyflink.dataframe as pf
+ >>> pf.set_model_provider(pf.OpenAIProvider(
+ ... "https://api.openai.com/v1/chat/completions";,
api_key="key", retry_num=3))
+ >>> df = pf.from_records([{"question": "What is Apache Flink?"}])
+ >>> answer = df.llm.predict("question", model="gpt-6-luna")
+
+ >>> # Override the provider's retry_num parameter for this call.
+ >>> answer = df.llm.predict(
+ ... "question", model="gpt-6-luna", retry_num=0)
+
+ >>> # Name the prediction column "answer" instead of the default
"output".
+ >>> answer = df.llm.predict(
+ ... "question", model="gpt-6-luna",
+ ... output_type=pf.DataType.struct({"answer":
pf.DataType.string()}))
+
+ .. versionadded:: 2.4.0
+ """
+ if not input_cols:
+ raise ValueError("predict() requires at least one input column")
+ for name in input_cols:
+ _validate_name(name, "input column")
+ if len(set(input_cols)) != len(input_cols):
+ raise ValueError("predict() input columns must be unique")
+ if model is not None:
+ _validate_name(model, "model")
+ if config is not None and not isinstance(config, MappingABC):
+ raise TypeError("config must be a mapping of strings to strings")
+ runtime_options = _merge_options(config, {})
+ selected = _resolve_provider(provider)
+ identifier = _validate_name(selected.provider_identifier(), "provider
identifier")
+ options = _provider_options(selected, kwargs)
+ if model is not None:
+ key = _validate_name(selected.model_option_key(), "model option
key")
+ options[key] = model
+ output_columns = _prediction_output_columns(output_type)
+ output_names = [name for name, _ in output_columns]
+ for name in output_names:
+ _validate_name(name, "output field name")
+ if len(set(output_names)) != len(output_names):
+ raise ValueError("Struct output_type contains duplicate field
names")
+
+ resolved = self._df._table.get_resolved_schema()
+ names = list(resolved.get_column_names())
+ missing = [name for name in input_cols if name not in names]
+ if missing:
+ raise ValueError(f"Input columns not found in DataFrame:
{missing}")
+ collisions = sorted(set(names).intersection(output_names))
+ if collisions:
+ raise ValueError(
+ f"Prediction output columns conflict with existing columns:
{collisions}")
+
+ input_schema = Schema.new_builder()
+ column_indexes = {name: index for index, name in enumerate(names)}
+ j_column_types = resolved._j_resolved_schema.getColumnDataTypes()
+ for name in input_cols:
+ input_schema.column(
+ name,
_from_java_data_type(j_column_types.get(column_indexes[name])))
+ output_schema = Schema.new_builder()
+ for name, dtype in output_columns:
+ output_schema.column(name, dtype)
+ descriptor = (ModelDescriptor.for_provider(identifier)
+ .input_schema(input_schema.build())
+ .output_schema(output_schema.build()))
+ for key, value in options.items():
+ if key != "provider":
+ descriptor.option(key, value)
+ flink_model = self._df._table._t_env.from_model(descriptor.build())
+ return DataFrame(flink_model.predict(self._df._table,
list(input_cols), runtime_options))
diff --git a/flink-python/pyflink/dataframe/tests/test_llm.py
b/flink-python/pyflink/dataframe/tests/test_llm.py
new file mode 100644
index 00000000000..6bbe3973a01
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_llm.py
@@ -0,0 +1,327 @@
+################################################################################
+# 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.
+################################################################################
+
+import json
+import unittest
+from types import MappingProxyType
+from unittest.mock import Mock, call, patch, sentinel
+
+import pyflink.dataframe as pf
+from pyflink.dataframe import llm
+from pyflink.table import DataTypes
+from pyflink.testing.test_case_utils import PyFlinkDataFrameUTTestCase
+
+
+class ProviderOverrideTests(unittest.TestCase):
+ def test_builtin_overrides_merge_without_mutating_provider(self):
+ provider = pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key", temperature=0.7,
+ options={"top-p": "0.9", "seed": "42", "future-option": "old"},
+ )
+ configured = provider.to_options()
+ overrides = {
+ "temperature": 0.2,
+ "options": MappingProxyType({"top-p": "0.3", "seed": None}),
+ }
+
+ self.assertEqual(llm._provider_options(provider, overrides), {
+ "endpoint": "https://example.test/chat/completions";, "api-key":
"key",
+ "temperature": "0.2", "top-p": "0.3", "future-option": "old",
+ })
+ self.assertEqual(provider.to_options(), configured)
+ self.assertEqual(overrides["options"], {"top-p": "0.3", "seed": None})
+ self.assertNotIn("top-p", llm._provider_options(provider, {"top_p":
None}))
+
+ def test_none_removes_inherited_parameters_or_raw_options(self):
+ provider = pf.TritonProvider(
+ "http://localhost:8000";, timeout="10 s", flatten_batch_dim=False,
+ options={"priority": "2", "future-option": "value"},
+ )
+ self.assertEqual(llm._provider_options(provider, {"timeout": None}), {
+ "endpoint": "http://localhost:8000";, "flatten-batch-dim": "false",
+ "priority": "2", "future-option": "value",
+ })
+ self.assertEqual(llm._provider_options(provider, {"options": None}), {
+ "endpoint": "http://localhost:8000";, "timeout": "10 s",
+ "flatten-batch-dim": "false",
+ })
+ self.assertEqual(llm._provider_options(provider, {"options": {}}),
provider.to_options())
+
+ def test_builtin_overrides_validate_typed_values_and_raw_conflicts(self):
+ provider = pf.OpenAIProvider("https://example.test/chat/completions";,
"key")
+ for overrides, error in [
+ ({"unknown": "value"}, TypeError),
+ ({"temperature": True}, TypeError),
+ ({"top_p": 1.1}, ValueError),
+ ({"endpoint": None}, TypeError),
+ ({"api_key": None}, TypeError),
+ ({"options": []}, TypeError),
+ ({"options": {1: "value"}}, TypeError),
+ ({"options": {"seed": 42}}, TypeError),
+ ({"options": {"api-key": None}}, TypeError),
+ ({"options": {"api-key": " "}}, ValueError),
+ ({"options": {"endpoint": None}}, TypeError),
+ ({"temperature": 0.2, "options": {"temperature": "0.3"}},
ValueError),
+ ]:
+ with self.subTest(overrides=overrides):
+ with self.assertRaises(error):
+ llm._provider_options(provider, overrides)
+ self.assertEqual(llm._provider_options(provider, {
+ "temperature": 0.2, "options": {"temperature": "0.2"},
+ })["temperature"], "0.2")
+
+ def test_triton_overrides_preserve_typed_parameters_and_copy_headers(self):
+ headers = {"X-Request": "original"}
+ provider = pf.TritonProvider(
+ "http://localhost:8000";, custom_headers=headers,
flatten_batch_dim=False,
+ max_retries=2, compression="gzip")
+ headers["X-Request"] = "changed"
+
+ options = llm._provider_options(provider, {"max_retries": 0})
+
+ self.assertEqual(options["max-retries"], "0")
+ self.assertEqual(options["flatten-batch-dim"], "false")
+ self.assertEqual(json.loads(options["custom-headers"]), {"X-Request":
"original"})
+ self.assertEqual(provider.to_options()["max-retries"], "2")
+ overridden = llm._provider_options(provider, {"custom_headers":
{"X-Request": "new"}})
+ self.assertEqual(json.loads(overridden["custom-headers"]),
{"X-Request": "new"})
+ self.assertNotIn("custom-headers", llm._provider_options(
+ provider, {"custom_headers": None}))
+ for value in (None, " "):
+ with self.subTest(endpoint=value):
+ with self.assertRaises((TypeError, ValueError)):
+ llm._provider_options(provider, {"options": {"endpoint":
value}})
+
+ def test_custom_provider_keeps_no_argument_serialization_contract(self):
+ class CustomProvider(pf.ModelProvider):
+ def provider_identifier(self):
+ return "custom"
+
+ def to_options(self):
+ return {"endpoint": "original", "seed": "42"}
+
+ for provider in (CustomProvider(), pf.GenericProvider(
+ "custom", endpoint="original", seed="42")):
+ with self.subTest(provider=type(provider).__name__):
+ self.assertEqual(llm._provider_options(provider, {
+ "endpoint": "override", "seed": None, "new-option":
"value",
+ }), {"endpoint": "override", "new-option": "value"})
+ self.assertEqual(provider.to_options(), {"endpoint":
"original", "seed": "42"})
+ with self.assertRaises(TypeError):
+ llm._provider_options(provider, {"seed": 1})
+
+
+class DataFramePredictTests(unittest.TestCase):
+ def setUp(self):
+ previous = dict(llm._provider_registry)
+ self.addCleanup(llm._provider_registry.update, previous)
+ self.addCleanup(llm._provider_registry.clear)
+ self.addCleanup(setattr, llm, "_default_provider",
llm._default_provider)
+ llm._provider_registry.clear()
+ llm._default_provider = None
+ self.provider = pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key", model="configured",
+ response_format="text",
+ )
+ pf.set_model_provider("chat", self.provider)
+ self.table = Mock()
+ resolved = self.table.get_resolved_schema.return_value
+ resolved.get_column_names.return_value = ["id", "prompt"]
+ resolved.get_column_data_types.return_value = [DataTypes.BIGINT(),
DataTypes.STRING()]
+ j_column_types =
resolved._j_resolved_schema.getColumnDataTypes.return_value
+ j_column_types.get.side_effect =
resolved.get_column_data_types.return_value.__getitem__
+ conversion_patch = patch.object(
+ llm, "_from_java_data_type", side_effect=lambda data_type:
data_type)
+ self.addCleanup(conversion_patch.stop)
+ conversion_patch.start()
+ self.dataframe = pf.DataFrame(self.table)
+ self.model = self.table._t_env.from_model.return_value
+ self.model.predict.return_value = sentinel.prediction
+ self.builders = []
+
+ def schema_builder():
+ builder = Mock()
+ self.builders.append(builder)
+ return builder
+
+ builder_patch = patch.object(llm.Schema, "new_builder",
side_effect=schema_builder)
+ self.addCleanup(builder_patch.stop)
+ builder_patch.start()
+ self.descriptor = Mock()
+ self.descriptor.input_schema.return_value = self.descriptor
+ self.descriptor.output_schema.return_value = self.descriptor
+ descriptor_patch = patch.object(
+ llm.ModelDescriptor, "for_provider", return_value=self.descriptor)
+ self.addCleanup(descriptor_patch.stop)
+ self.for_provider = descriptor_patch.start()
+
+ def
test_predict_forwards_input_order_runtime_config_and_model_precedence(self):
+ self.assertIsInstance(self.dataframe.llm, llm.LLMAccessor)
+ resolved = self.table.get_resolved_schema.return_value
+ resolved.get_column_names.return_value += [f"metadata_{i}" for i in
range(1000)]
+ resolved.get_column_data_types.return_value += [DataTypes.STRING()] *
1000
+ provider = pf.TritonProvider(
+ "http://localhost:8000";, model_name="configured",
options={"provider": "other"})
+ pf.set_model_provider("triton", provider)
+ pf.set_default_model_provider("chat")
+ config = MappingProxyType({"timeout": "30 s", "custom-runtime-option":
"value"})
+
+ result = self.dataframe.llm.predict(
+ "prompt", "id", provider="triton", model="selected", config=config,
+ model_name="overridden", timeout="10 s",
+ )
+
+ self.assertIs(result._table, sentinel.prediction)
+ self.for_provider.assert_called_once_with("triton")
+ self.assertEqual(self.builders[0].column.call_args_list, [
+ call("prompt", DataTypes.STRING()), call("id", DataTypes.BIGINT()),
+ ])
+ self.assertEqual(self.builders[1].column.call_args_list,
+ [call("output", DataTypes.STRING())])
+
self.table._t_env.from_model.assert_called_once_with(self.descriptor.build.return_value)
+ self.model.predict.assert_called_once_with(self.table, ["prompt",
"id"], dict(config))
+ self.assertEqual(dict(args.args for args in
self.descriptor.option.call_args_list), {
+ "endpoint": "http://localhost:8000";, "model-name": "selected",
"timeout": "10 s",
+ })
+ self.assertEqual(provider.to_options()["model-name"], "configured")
+ self.assertIs(llm._resolve_provider(), self.provider)
+
+ def
test_predict_rejects_invalid_inputs_models_and_config_before_planning(self):
+ for columns, kwargs, error in [
+ ((), {}, ValueError),
+ ((1,), {}, TypeError),
+ ((" ",), {}, ValueError),
+ (("prompt", "prompt"), {}, ValueError),
+ (("missing",), {}, ValueError),
+ (("prompt",), {"model": ""}, ValueError),
+ (("prompt",), {"provider": "missing"}, ValueError),
+ (("prompt",), {"config": []}, TypeError),
+ (("prompt",), {"config": {"async": True}}, TypeError),
+ (("prompt",), {"config": {1: "value"}}, TypeError),
+ (("prompt",), {"output_type": DataTypes.STRING()}, TypeError),
+ ]:
+ with self.subTest(columns=columns, kwargs=kwargs):
+ with self.assertRaises(error):
+ self.dataframe.llm.predict(*columns, **kwargs)
+ self.assertEqual(self.builders, [])
+ self.table._t_env.from_model.assert_not_called()
+
+ def test_predict_rejects_empty_duplicate_and_colliding_output_fields(self):
+ for fields in ({}, [("x", pf.DataType.string()), ("x",
pf.DataType.int32())],
+ {" ": pf.DataType.string()}, {"prompt":
pf.DataType.string()}):
+ with self.subTest(fields=fields):
+ with self.assertRaises(ValueError):
+ self.dataframe.llm.predict("prompt",
output_type=pf.DataType.struct(fields))
+
self.table.get_resolved_schema.return_value.get_column_names.return_value = [
+ "prompt", "output",
+ ]
+ with self.assertRaisesRegex(ValueError, "conflict with existing
columns"):
+ self.dataframe.llm.predict("prompt")
+ self.assertEqual(self.builders, [])
+
+ def test_default_output_is_string_for_all_providers(self):
+ pf.set_model_provider("other-task", pf.OpenAIProvider(
+ "https://example.test/embeddings";, "key"))
+ pf.set_model_provider("triton",
pf.TritonProvider("http://localhost:8000";))
+ pf.set_model_provider("generic", pf.GenericProvider("custom"))
+ for provider in ("chat", "other-task", "triton", "generic"):
+ with self.subTest(provider=provider):
+ self.dataframe.llm.predict("prompt", provider=provider)
+ self.assertEqual(self.builders[-1].column.call_args_list,
+ [call("output", DataTypes.STRING())])
+
+ def
test_prediction_declares_native_output_types_and_preserves_nested_fields(self):
+ pf.set_model_provider("custom", pf.GenericProvider("custom"))
+ pf.set_model_provider("embeddings", pf.OpenAIProvider(
+ "https://example.test/embeddings";, "key", model="embedding"))
+ nested = pf.DataType.struct({"label": pf.DataType.string()})
+ for provider, output_type, fields in [
+ ("custom", None, [("output", DataTypes.STRING())]),
+ ("custom", pf.DataType.int32(), [("output", DataTypes.INT())]),
+ ("custom", pf.DataType(DataTypes.VARCHAR(5)), [("output",
DataTypes.VARCHAR(5))]),
+ ("custom", pf.DataType.string().not_null(),
+ [("output", DataTypes.STRING().not_null())]),
+ ("custom", pf.DataType.struct({"answer": nested, "score":
pf.DataType.int32()}),
+ [("answer", nested._to_table_data_type()), ("score",
DataTypes.INT())]),
+ ("embeddings", pf.DataType.list(pf.DataType.float32()),
+ [("output", DataTypes.ARRAY(DataTypes.FLOAT()))]),
+ ]:
+ with self.subTest(output_type=output_type):
+ self.descriptor.reset_mock()
+ self.dataframe.llm.predict("prompt", provider=provider,
output_type=output_type)
+ self.assertEqual(self.builders[-1].column.call_args_list,
+ [call(name, dtype) for name, dtype in fields])
+ options = dict(args.args for args in
self.descriptor.option.call_args_list)
+ self.assertNotIn("response-format", options)
+
+ def
test_prediction_preserves_inherited_and_explicit_response_formats(self):
+ for overrides, response_format in [
+ ({}, "text"), ({"response_format": "json_object"}, "json_object"),
+ ]:
+ with self.subTest(overrides=overrides):
+ self.descriptor.reset_mock()
+ self.dataframe.llm.predict(
+ "prompt", output_type=pf.DataType.int32(), **overrides)
+ options = dict(args.args for args in
self.descriptor.option.call_args_list)
+ self.assertEqual(options["response-format"], response_format)
+ self.assertEqual(self.provider.to_options()["response-format"], "text")
+
+
+class DataFramePredictSchemaTests(PyFlinkDataFrameUTTestCase):
+ def setUp(self):
+ super().setUp()
+ previous = dict(llm._provider_registry)
+ self.addCleanup(llm._provider_registry.update, previous)
+ self.addCleanup(llm._provider_registry.clear)
+ self.addCleanup(setattr, llm, "_default_provider",
llm._default_provider)
+ llm._provider_registry.clear()
+ llm._default_provider = None
+ pf.set_model_provider(pf.GenericProvider("python-values"))
+ self.dataframe = pf.from_records([(1, "question")], schema=["id",
"prompt"])
+
+ def test_prediction_preserves_input_columns_and_is_not_registered(self):
+ registered_models = self.t_env.list_models()
+ prediction = self.dataframe.llm.predict("id", "prompt")
+
+ self.assert_dataframe_schema(prediction, ["id", "prompt", "output"], [
+ DataTypes.BIGINT(), DataTypes.STRING(), DataTypes.STRING(),
+ ])
+ self.assertEqual(self.t_env.list_models(), registered_models)
+
+ def test_sql_row_is_one_column_and_struct_expands_only_top_level(self):
+ nested = DataTypes.ROW([DataTypes.FIELD("label", DataTypes.STRING())])
+ for output_type, names, types in [
+ ("ROW<label STRING>", ["output"], [nested]),
+ (pf.DataType.struct({"details": pf.DataType(nested), "score":
pf.DataType.int32()}),
+ ["details", "score"], [nested, DataTypes.INT()]),
+ ]:
+ with self.subTest(output_type=output_type):
+ prediction = self.dataframe.llm.predict("id",
output_type=output_type)
+ self.assert_dataframe_schema(prediction, ["id", "prompt"] +
names, [
+ DataTypes.BIGINT(), DataTypes.STRING(),
+ ] + types)
+
+
+if __name__ == '__main__':
+ 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/dataframe/tests/test_llm_integration.py
b/flink-python/pyflink/dataframe/tests/test_llm_integration.py
new file mode 100644
index 00000000000..56d4c829ec1
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_llm_integration.py
@@ -0,0 +1,88 @@
+################################################################################
+# 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.
+################################################################################
+
+import unittest
+
+import pyflink.dataframe as pf
+from pyflink.common import Row
+from pyflink.dataframe import llm
+from pyflink.testing.test_case_utils import PyFlinkStreamDataFrameTestCase
+
+
+class DataFrameLlmITCase(PyFlinkStreamDataFrameTestCase):
+
+ def setUp(self):
+ previous = dict(llm._provider_registry)
+ self.addCleanup(llm._provider_registry.update, previous)
+ self.addCleanup(llm._provider_registry.clear)
+ self.addCleanup(setattr, llm, "_default_provider",
llm._default_provider)
+ llm._provider_registry.clear()
+ llm._default_provider = None
+ self.provider = pf.GenericProvider("python-values", {"async": "false",
"latency": "0"})
+ pf.set_model_provider("values", self.provider)
+ self.output_type = pf.DataType.struct({
+ "x": pf.DataType.string(),
+ "y": pf.DataType.int32(),
+ "z": pf.DataType.string(),
+ })
+
+ def test_sync_prediction_appends_typed_columns(self):
+ registered_models = self.t_env.list_models()
+ dataframe = pf.from_records([(1, "first"), (3, "third")],
schema=["id", "content"])
+
+ prediction = dataframe.llm.predict(
+ "id", output_type=self.output_type, config={"async": "false"},
+ )
+
+ self.assertEqual(prediction.columns, ["id", "content", "x", "y", "z"])
+ self.assertCountEqual(prediction.collect(), [
+ Row(1, "first", "x1", 1, "z1"),
+ Row(3, "third", "x3", 3, "z3"),
+ ])
+ self.assertEqual(self.t_env.list_models(), registered_models)
+
+ def test_async_prediction_with_provider_override_and_runtime_config(self):
+ dataframe = pf.from_records([(2, "second"), (4, "fourth")],
schema=["id", "content"])
+
+ prediction = dataframe.llm.predict(
+ "id", provider="values", output_type=self.output_type,
+ config={
+ "async": "true",
+ "output-mode": "allow_unordered",
+ "max-concurrent-operations": "2",
+ "timeout": "30 s",
+ },
+ **{"async": "true"},
+ )
+ selected = prediction.select("content", "x", (pf.col("y") +
1).alias("next_id"), "z")
+
+ self.assertCountEqual(selected.collect(), [
+ Row("second", "x2", 3, "z2"),
+ Row("fourth", "x4", 5, "z4"),
+ ])
+ self.assertEqual(self.provider.to_options(), {"async": "false",
"latency": "0"})
+
+
+if __name__ == '__main__':
+ 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/model.py
b/flink-python/pyflink/table/model.py
index 485bf3a64b8..db57f19f640 100644
--- a/flink-python/pyflink/table/model.py
+++ b/flink-python/pyflink/table/model.py
@@ -28,8 +28,10 @@ __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.
+ A model object represents a machine learning model retrieved from the
catalog or
+ created from a model descriptor. Use :meth:`~TableEnvironment.from_model`
to retrieve
+ a named model or create an anonymous model, then :meth:`predict` to perform
+ predictions on input tables.
.. versionadded:: 2.4.0
"""
diff --git a/flink-python/pyflink/table/table_environment.py
b/flink-python/pyflink/table/table_environment.py
index 6d84c39ec0a..ebab2034a38 100644
--- a/flink-python/pyflink/table/table_environment.py
+++ b/flink-python/pyflink/table/table_environment.py
@@ -1250,9 +1250,12 @@ 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':
+ def from_model(self, model_path: Union[str, ModelDescriptor]) -> 'Model':
"""
- Gets a model registered in the catalog.
+ Gets a registered model or creates an anonymous model from a
descriptor.
+
+ A descriptor defines the provider options and input/output schemas
directly,
+ without registering a model in the catalog.
Example:
::
@@ -1262,12 +1265,28 @@ class TableEnvironment(object):
>>> 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.
+ Create an anonymous model:
+ ::
+
+ >>> from pyflink.table import DataTypes
+ >>> descriptor = (ModelDescriptor.for_provider("my-provider")
+ ... .input_schema(Schema.new_builder()
+ ... .column("text", DataTypes.STRING()).build())
+ ... .output_schema(Schema.new_builder()
+ ... .column("label",
DataTypes.STRING()).build())
+ ... .build())
+ >>> model = t_env.from_model(descriptor)
+
+ :param model_path: The path of the registered model, or a descriptor
defining
+ an anonymous model.
:return: The requested model instance.
.. versionadded:: 2.4.0
"""
- j_model = self._j_tenv.fromModel(model_path)
+ if isinstance(model_path, ModelDescriptor):
+ j_model = self._j_tenv.fromModel(model_path._j_model_descriptor)
+ else:
+ j_model = self._j_tenv.fromModel(model_path)
return Model(j_model)
def create_temporary_model(self,
diff --git a/flink-python/pyflink/table/tests/test_model.py
b/flink-python/pyflink/table/tests/test_model.py
index 5e1a64d6458..7ceb3a795cc 100644
--- a/flink-python/pyflink/table/tests/test_model.py
+++ b/flink-python/pyflink/table/tests/test_model.py
@@ -18,7 +18,35 @@
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
+from pyflink.testing.test_case_utils import PyFlinkStreamTableTestCase,
PyFlinkUTTestCase
+
+
+class AnonymousModelTests(PyFlinkUTTestCase):
+
+ def test_anonymous_model_descriptor(self):
+ 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()
+ registered_models = self.t_env.list_models()
+
+ model = self.t_env.from_model(descriptor)
+
+
self.assertEqual(list(model.get_resolved_input_schema().get_column_names()),
["id"])
+
self.assertEqual(model.get_resolved_input_schema().get_column_data_types(),
+ [DataTypes.BIGINT()])
+
self.assertEqual(list(model.get_resolved_output_schema().get_column_names()),
+ ["x", "y", "z"])
+
self.assertEqual(model.get_resolved_output_schema().get_column_data_types(),
+ [DataTypes.STRING(), DataTypes.INT(),
DataTypes.STRING()])
+
+ self.assertEqual(self.t_env.list_models(), registered_models)
class ModelTest(PyFlinkStreamTableTestCase):
@@ -125,7 +153,7 @@ class ModelTest(PyFlinkStreamTableTestCase):
.build()
self.t_env.create_temporary_model("opt_model", model_descriptor, True)
- model = self.t_env.from_model("opt_model")
+ model = self.t_env.from_model(model_path="opt_model")
data = [
(1, 12, "Julian"),
