################################################################################
# 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.
################################################################################
"""
The AI module provides integration with large language models (LLMs) for tasks
such as classification, sentiment analysis, text extraction, translation,
summarization, PII masking, embedding generation, and vector search. Access
these functions through the df.llm accessor on a DataFrame.
Example::
import pyflink.dataframe as pf
pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
df = pf.from_dict({"text": ["Hello world", "Bonjour le monde"]})
df.llm.ai_translate(
"text", source_lang="auto", target_lang="en", model="qwen3.6-plus").collect()
Caching:
LLM functions accept cache_table and cache_key arguments. cache_table must
identify a Fluss catalog table.
If cache_table is set and cache_key is omitted, Flink uses a synthetic
cache key derived from model options and function input. The cache table
must declare exactly one primary-key column for that generated key, and
that column must be STRING.
If cache_key is specified, pass either a column name or a list of column
names from the input DataFrame. The cache table must declare the same
number of primary-key columns, in the same order. Each DataFrame cache-key
column must have the same type as the corresponding primary-key column.
The non-primary-key cache columns must match the LLM function output
columns.
Implicit cache key example::
# Fluss catalog cache table:
# fluss_catalog.default_database.translate_cache
# +-------------------+--------+-----------------------------+
# | Column | Type | Notes |
# +-------------------+--------+-----------------------------+
# | cache_key | STRING | Primary key (used for |
# | | | generated cache key) |
# | translated_text | STRING | ai_translate output |
# | detected_language | STRING | ai_translate output |
# +-------------------+--------+-----------------------------+
df.llm.ai_translate(
"text", source_lang="auto", target_lang="en",
cache_table="fluss_catalog.default_database.translate_cache")
Explicit cache key example::
# Input DataFrame columns used as cache keys:
# tenant_id STRING
# text STRING
#
# Fluss catalog cache table:
# fluss_catalog.default_database.translate_cache_by_text
# +-------------------+--------+-----------------------------+
# | Column | Type | Notes |
# +-------------------+--------+-----------------------------+
# | tenant_id | STRING | Primary key (used for |
# | | | cache_key[0]) |
# | text | STRING | Primary key (used for |
# | | | cache_key[1]) |
# | translated_text | STRING | ai_translate output |
# | detected_language | STRING | ai_translate output |
# +-------------------+--------+-----------------------------+
df.llm.ai_translate(
"text", source_lang="auto", target_lang="en",
cache_table=(
"fluss_catalog.default_database.translate_cache_by_text"),
cache_key=["tenant_id", "text"])
"""
from functools import lru_cache
import hashlib
import json
from typing import Dict, List, Mapping, Optional, Union, TYPE_CHECKING
from pyflink.dataframe.ai.providers import ModelProvider, create_provider
CacheKey = Union[str, List[str]]
if TYPE_CHECKING:
from pyflink.dataframe.dataframe import DataFrame
from pyflink.dataframe.datatype import DataType
from pyflink.table import Expression
# ---------------------------------------------------------------------------
# Global provider registry
# ---------------------------------------------------------------------------
_provider_registry: Dict[str, ModelProvider] = {}
_default_provider: Optional[str] = None
[docs]def set_model_provider(name_or_provider: Union[str, ModelProvider],
provider: Optional[ModelProvider] = None,
**options) -> None:
"""
Register a global model provider configuration.
The registered provider name is the Python-side lookup key accepted by
``provider=`` arguments, ``set_default_model_provider()``, and returned by
``list_model_providers()``. Registered provider names are unique registry
keys; registering the same name again raises ``ValueError``.
This lookup name is separate from ``ModelProvider.provider_identifier()``,
which returns the Flink Java provider identifier. Multiple registered names
may point to providers with the same Java identifier, such as separate
OpenAI-compatible providers for chat and embeddings.
Can be called in three ways:
1. ``set_model_provider(model_provider)`` — use
``model_provider.provider_identifier()`` as the lookup name.
2. ``set_model_provider("name", model_provider)`` — register under a
custom lookup name. This allows registering multiple instances of the
same provider type (e.g. one for chat, one for embeddings).
3. ``set_model_provider("name", **options)`` — create a built-in provider
when ``name`` matches a built-in provider identifier, otherwise create
a :class:`GenericProvider` with the given options.
Args:
name_or_provider: Either a :class:`ModelProvider` instance (form 1) or a
lookup name string (forms 2 and 3).
provider: A :class:`ModelProvider` instance to register under the custom
lookup name (form 2 only).
**options: Model provider options (form 3 only, dispatches when ``name``
matches a built-in provider identifier, otherwise wraps in
:class:`GenericProvider`).
Example::
>>> import pyflink.dataframe as pf
>>> # Form 1: ModelProvider instance (registered as "openai-compat")
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> # Form 2: Custom name + ModelProvider instance
>>> pf.set_model_provider("chat", pf.OpenAICompatProvider(
... task="chat/completions"))
>>> pf.set_model_provider("embedding", pf.OpenAICompatProvider(
... task="embeddings"))
>>> # Form 3: Generic string API
>>> pf.set_model_provider(
... "openai-compat",
... task="chat/completions")
"""
if isinstance(name_or_provider, ModelProvider):
if provider is not None:
raise ValueError(
"When passing a ModelProvider as the first argument, "
"the second argument must not be a ModelProvider.")
if options:
raise ValueError(
"When passing a ModelProvider as the first argument, "
"**options must not be provided.")
p = name_or_provider
_register_model_provider(p.provider_identifier(), p)
elif isinstance(name_or_provider, str):
if provider is not None:
if not isinstance(provider, ModelProvider):
raise TypeError(
f"The second argument must be a ModelProvider instance, "
f"got {type(provider).__name__}.")
if options:
raise ValueError(
"Cannot pass both a ModelProvider instance and **options.")
_register_model_provider(name_or_provider, provider)
else:
_ensure_model_provider_name_available(name_or_provider)
_register_model_provider(
name_or_provider, create_provider(name_or_provider, **options))
else:
raise TypeError(
f"First argument must be a str or ModelProvider, got "
f"{type(name_or_provider).__name__}.")
def _register_model_provider(name: str, provider: ModelProvider) -> None:
_ensure_model_provider_name_available(name)
_provider_registry[name] = provider
def _ensure_model_provider_name_available(name: str) -> None:
if name in _provider_registry:
raise ValueError(f"Model provider name {name!r} is already registered.")
[docs]def set_default_model_provider(name: str) -> None:
"""
Set the default provider lookup name for AI functions.
When multiple providers are registered, this determines which one
is used when ``provider=`` is not specified in a function call.
Args:
name: Registered provider lookup name (must already be registered via
set_model_provider).
Example::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider("chat", pf.OpenAICompatProvider(
... task="chat/completions"))
>>> pf.set_model_provider("embedding", pf.OpenAICompatProvider(
... task="embeddings"))
>>> pf.set_default_model_provider("chat")
"""
if name not in _provider_registry:
raise ValueError(
f"ModelProvider {name!r} is not registered. "
"Call pf.set_model_provider(...) first.")
global _default_provider
_default_provider = name
[docs]def list_model_providers() -> List[str]:
"""
List all registered provider lookup names.
Returns:
A list of registered provider lookup names. These are the Python-side
registry keys, not necessarily the Flink Java provider identifiers.
Example::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider("chat", pf.OpenAICompatProvider(
... task="chat/completions"))
>>> pf.set_model_provider("embedding", pf.OpenAICompatProvider(
... task="embeddings"))
>>> pf.list_model_providers()
['chat', 'embedding']
"""
return list(_provider_registry.keys())
def _resolve_provider(explicit: Optional[str] = None) -> Optional[str]:
"""Resolve the registered provider lookup name, or None if no provider exists."""
if explicit is not None:
if not explicit:
raise ValueError("Model provider name cannot be empty.")
if explicit not in _provider_registry:
raise ValueError(
f"ModelProvider {explicit!r} is not registered. "
"Call pf.set_model_provider(...) first.")
return explicit
if _default_provider:
return _default_provider
if len(_provider_registry) == 1:
return next(iter(_provider_registry))
if len(_provider_registry) == 0:
return None
raise ValueError(
"Multiple providers registered. Either call "
"pf.set_default_model_provider(...) or pass provider='...' explicitly.")
# ---------------------------------------------------------------------------
# LLMAccessor — accessed via df.llm
# ---------------------------------------------------------------------------
@lru_cache(maxsize=1)
def _ai_function_schemas():
"""Schema specs for AI functions: {func_name: (input_cols, output_cols)}.
Each col is (name, DataType). Built lazily because DataTypes requires JVM.
"""
from pyflink.table import DataTypes
return {
"ai_classify": ([("input", DataTypes.STRING())],
[("content", DataTypes.VARIANT())]),
"ai_sentiment": ([("input", DataTypes.STRING())],
[("content", DataTypes.VARIANT())]),
"ai_extract": ([("input", DataTypes.STRING())],
[("content", DataTypes.VARIANT())]),
"ai_translate": ([("input", DataTypes.STRING())],
[("content", DataTypes.VARIANT())]),
"ai_summarize": ([("input", DataTypes.STRING())],
[("content", DataTypes.VARIANT())]),
"ai_mask": ([("input", DataTypes.STRING())],
[("content", DataTypes.VARIANT())]),
"ai_embed": ([("input", DataTypes.STRING())],
[("embedding", DataTypes.ARRAY(DataTypes.FLOAT()))]),
}
def _build_schema(columns):
"""Build a Schema from a list of (name, DataType) tuples."""
from pyflink.table import Schema
builder = Schema.new_builder()
for col_name, data_type in columns:
builder.column(col_name, data_type)
return builder.build()
def _canonicalize_model_columns(columns):
return [
{"name": str(col_name), "data_type": str(data_type)}
for col_name, data_type in (columns or [])
]
def _model_descriptor_digest(java_provider_identifier, descriptor_options,
model_option_key, model_name, input_columns,
output_columns):
effective_options = {"provider": str(java_provider_identifier)}
effective_options.update({
str(key): str(value) for key, value in descriptor_options.items()
})
if model_name:
effective_options[str(model_option_key)] = str(model_name)
payload = {
"comment": None,
"input_schema": _canonicalize_model_columns(input_columns),
"options": effective_options,
"output_schema": _canonicalize_model_columns(output_columns),
}
canonical_payload = json.dumps(
payload, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(canonical_payload.encode("utf-8")).hexdigest()[:16]
def _create_model_descriptor_and_digest(t_env, provider_name, model_name,
input_columns, output_columns,
provider_options=None):
from pyflink.dataframe.util.artifacts import add_built_in_model
from pyflink.table import ModelDescriptor
provider = _provider_registry.get(provider_name)
if provider is None:
raise ValueError(
f"ModelProvider {provider_name!r} is not registered. "
"Call pf.set_model_provider(...) first.")
descriptor_options = provider.to_options(
input_columns, overrides=provider_options)
java_provider_identifier = provider.provider_identifier()
add_built_in_model(t_env, java_provider_identifier)
builder = ModelDescriptor.for_provider(java_provider_identifier)
for k, v in descriptor_options.items():
builder = builder.option(k, v)
model_option_key = provider.model_option_key()
if model_name:
builder = builder.option(model_option_key, model_name)
builder = builder.input_schema(_build_schema(input_columns))
builder = builder.output_schema(_build_schema(output_columns))
descriptor_digest = _model_descriptor_digest(
java_provider_identifier, descriptor_options, model_option_key,
model_name, input_columns, output_columns)
return builder.build(), descriptor_digest
def _create_model(t_env, provider_name, model_name, input_columns,
output_columns, provider_options=None):
"""Create a temporary model and return the Model object.
Args:
t_env: The TableEnvironment.
provider_name: Registered provider lookup name.
model_name: Model name.
input_columns: List of (name, DataType) tuples for input schema.
output_columns: List of (name, DataType) tuples for output schema.
provider_options: Optional per-call provider option overrides.
"""
descriptor, descriptor_digest = _create_model_descriptor_and_digest(
t_env, provider_name, model_name, input_columns, output_columns,
provider_options)
# AI function cache's default synthetic cache key includes the model name,
# so derive the temporary model name from canonical descriptor components
# instead of a random UUID.
temp_name = f"_pf_llm_{descriptor_digest}"
t_env.create_temporary_model(temp_name, descriptor, True)
return t_env.from_model(temp_name)
def _resolve_model(t_env, model, provider=None, input_columns=None,
output_columns=None, provider_options=None):
"""Resolve a Model, either from provider registry or catalog.
If a registered provider is available (explicitly given or resolved from
registry), creates a temporary model. Otherwise, treats ``model`` as a
catalog model name.
Args:
t_env: The TableEnvironment.
model: Model name — either a provider model id (e.g. "qwen3.6-plus")
or a catalog model name (e.g. "my_catalog_model").
provider: Explicit registered provider lookup name, or None to
auto-resolve.
input_columns: List of (name, DataType) tuples for input schema.
output_columns: List of (name, DataType) tuples for output schema.
provider_options: Optional per-call provider option overrides.
Returns:
A Model object.
"""
provider_name = _resolve_provider(provider)
if provider_name is not None:
# Provider path: lazily create a temporary model descriptor from the
# registered DataFrame model provider and per-call overrides.
return _create_model(t_env, provider_name, model,
input_columns, output_columns, provider_options)
# Catalog path: no DataFrame model provider is configured, so resolve
# ``model`` as an existing catalog model and reject provider overrides.
if (provider_options
and any(value is not None for value in provider_options.values())):
raise ValueError(
"Provider option overrides require a configured model provider.")
if model is None:
raise ValueError(
"Either configure a provider via pf.set_model_provider(...) "
"or pass a catalog model name via model='...'.")
return t_env.from_model(model)
def _normalize_cache_key(cache_key):
from pyflink.table.ai_functions import _validate_cache_key_list
if cache_key is None:
return None
if isinstance(cache_key, str):
return _validate_cache_key_list("cache_key", [cache_key])
if isinstance(cache_key, list):
if not cache_key:
raise ValueError("cache_key must not be empty")
return _validate_cache_key_list("cache_key", cache_key)
raise TypeError("cache_key must be a string, a list of strings, or None")
def _validate_cache_args(cache_table, cache_key):
from pyflink.table.ai_functions import _validate_cache_table
if cache_key is not None and cache_table is None:
raise ValueError("cache_key requires cache_table to be specified")
_validate_cache_table(cache_table)
return _normalize_cache_key(cache_key)
class LLMAccessor:
"""
Accessor for LLM / AI functions on a DataFrame.
Accessed via ``df.llm``. Groups all AI-related operations together.
Example::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>>
>>> df = pf.from_dict({"text": ["hello", "world"]})
>>> df = df.llm.predict("text", model="qwen3.6-plus")
>>> df = df.llm.ai_classify("text", ["pos", "neg"], model="qwen3.6-plus")
"""
def __init__(self, df: "DataFrame"):
self._df = df
@staticmethod
def _to_col_expr(input_col: Union[str, "Expression"]) -> "Expression":
"""Convert a column name string to an Expression if needed."""
if isinstance(input_col, str):
from pyflink.dataframe.dataframe import col
return col(input_col)
return input_col
def _get_column_type(self, col_name: str):
"""Get the DataType of a column from the DataFrame's schema."""
schema = self._df._table.get_resolved_schema()
for column in schema.get_columns():
if column.get_name() == col_name:
return column.get_data_type()
available = [c.get_name() for c in schema.get_columns()]
raise ValueError(
f"Column {col_name!r} not found. Available columns: {available}")
[docs] def vector_search(
self,
search_source: "DataFrame",
column_to_search: str,
column_to_query: Union[str, "Expression"],
top_k: int,
*,
agg: bool = False,
output_columns: Optional[Union[str, List[str]]] = None,
config: Optional[Mapping[str, str]] = None,
ignore_empty: bool = False,
) -> "DataFrame":
"""
Join this DataFrame with vector search results from another source.
Args:
search_source: Search source as a DataFrame.
column_to_search: Vector column in the search source.
column_to_query: Query vector column or direct column reference from
this DataFrame.
top_k: Number of nearest matches to return per input row.
agg: Whether to aggregate matches into one array column.
output_columns: Optional output column name or names.
config: Optional dict of runtime vector search options.
ignore_empty: Whether to drop rows with no search output.
Returns:
A new DataFrame with vector search outputs appended.
Note:
Let ``N`` be the number of columns in ``search_source``.
In non-aggregation mode (``agg=False``), the result may contain up
to ``top_k`` rows for every input row and appends ``N + 1`` columns:
the ``N`` columns from ``search_source`` with their original types,
followed by one score column of type ``DOUBLE``. If
``output_columns`` is specified, it must contain exactly ``N + 1``
names and is used as the names of these appended columns in order.
In aggregation mode (``agg=True``), the result contains one row for
every input row and appends one array column. The array element is a
row with ``N + 1`` nested columns: the ``N`` columns from
``search_source`` with their original types, followed by one score
field of type ``DOUBLE``. If ``output_columns`` is specified, it
must contain exactly one name and is used as the name of the
appended array column.
Example:
::
>>> import pyflink.dataframe as pf
>>> from pyflink.table import DataTypes
>>> queries = pf.from_records(
... [(1, [0.1, 0.2]), (2, [0.3, 0.4])],
... schema=["query_id", "query_embedding"],
... )
>>> documents = pf.read_milvus(
... endpoint="http://milvus.example.com",
... username="${secret_values.milvus_user}",
... password="${secret_values.milvus_password}",
... database_name="my_db",
... collection_name="documents",
... schema={
... "doc_id": DataTypes.BIGINT(),
... "embedding": DataTypes.ARRAY(DataTypes.FLOAT()),
... "title": DataTypes.STRING(),
... },
... columns=["doc_id", "embedding", "title"],
... )
>>> result = queries.llm.vector_search(
... documents,
... column_to_search="embedding",
... column_to_query="query_embedding",
... top_k=5,
... output_columns=[
... "match_doc_id",
... "match_embedding",
... "match_title",
... "score",
... ],
... )
"""
from pyflink.dataframe.dataframe import DataFrame as DF
from pyflink.table.expressions import (
col as table_col,
vector_search as table_vector_search,
vector_search_agg as table_vector_search_agg,
)
def resolve_search_source(source: "DataFrame"):
if isinstance(source, DF):
return source.to_table(), list(source.columns)
raise TypeError(
f"search_source must be a DataFrame, got {type(source).__name__}"
)
def make_unique_score_name(search_columns: List[str]) -> str:
existing_columns = set(search_columns)
output_name = "score"
candidate = output_name
count = 0
while candidate in existing_columns:
candidate = f"{output_name}{count}"
count += 1
return candidate
def resolve_default_outputs(search_columns: List[str]) -> List[str]:
if agg:
return ["search_results"]
return search_columns + [make_unique_score_name(search_columns)]
def validate_no_output_conflicts(resolved_outputs: List[str]) -> None:
conflicting_columns = set(self._df.columns).intersection(resolved_outputs)
if conflicting_columns:
raise ValueError(
"output_columns conflicts with existing DataFrame columns; "
"use unique output_columns: "
f"{sorted(conflicting_columns)}"
)
def resolve_output_columns(default_outputs: List[str]) -> List[str]:
if output_columns is None:
resolved_outputs = list(default_outputs)
elif isinstance(output_columns, str):
resolved_outputs = [output_columns]
elif isinstance(output_columns, list):
if not output_columns:
raise ValueError("output_columns must not be empty")
if not all(isinstance(name, str) for name in output_columns):
raise TypeError(
"output_columns list must contain only strings")
resolved_outputs = list(output_columns)
else:
raise TypeError(
"output_columns must be a string, a list of strings, "
f"or None, got {type(output_columns).__name__}"
)
if any(not name.strip() for name in resolved_outputs):
raise ValueError(
"output_columns must contain only non-empty strings")
expected_count = 1 if agg else len(default_outputs)
if len(resolved_outputs) != expected_count:
raise ValueError(
"output_columns count must match vector_search output "
f"column count: expected {expected_count}, "
f"got {len(resolved_outputs)}"
)
if len(set(resolved_outputs)) != len(resolved_outputs):
raise ValueError(
"output_columns contains duplicate column names")
validate_no_output_conflicts(resolved_outputs)
return resolved_outputs
if not isinstance(column_to_search, str):
raise TypeError(
"column_to_search must be a string, "
f"got {type(column_to_search).__name__}"
)
if not column_to_search.strip():
raise ValueError("column_to_search must not be empty")
if not isinstance(top_k, int) or isinstance(top_k, bool):
raise TypeError("top_k must be an integer")
if top_k <= 0:
raise ValueError("top_k must be greater than 0")
if not isinstance(agg, bool):
raise TypeError("agg must be a boolean")
if not isinstance(ignore_empty, bool):
raise TypeError("ignore_empty must be a boolean")
search_table, search_columns = resolve_search_source(search_source)
default_outputs = resolve_default_outputs(search_columns)
resolved_outputs = resolve_output_columns(default_outputs)
vector_search_function = (
table_vector_search_agg if agg else table_vector_search
)
table_function = vector_search_function(
search_table,
column_to_search,
column_to_query,
top_k,
config,
).alias(*resolved_outputs)
if ignore_empty:
joined = self._df._table.join_lateral(table_function)
else:
joined = self._df._table.left_outer_join_lateral(table_function)
return DF(
joined.select(
*[table_col(c) for c in self._df.columns + resolved_outputs]))
[docs] def predict(self, *input_cols: str, provider: Optional[str] = None,
model: Optional[str] = None,
output_type: Optional[Union[str, "DataType"]] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Perform prediction using a model.
This is the general-purpose prediction method.
When a provider is configured, a temporary model is created with
the given input/output schema. When using a catalog model
(no provider), the model's registered schema is used and
``output_type`` is ignored.
Args:
*input_cols: Column names to use as input.
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name (e.g. "qwen3.6-plus") or catalog model
name.
output_type: Output type declaration. ``None`` defaults to a
single ``STRING`` column named ``output``. A SQL type string
such as ``"VARIANT"`` appends a single column named
``output`` with that parsed type. A scalar or container
``DataType`` also appends a single column named ``output``.
A ``DataType.struct(...)`` appends one column per top-level
struct field.
cache_table: Optional cache table identifier.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Returns:
A new DataFrame with the model output columns appended.
Examples:
Runtime config keys depend on the provider.
Set up a provider and input DataFrame:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> df = pf.from_dict({
... "tenant_id": ["tenant-a"],
... "question": ["Which products need review?"],
... "image_url": ["https://example.com/product.jpg"],
... })
Default output: appends one ``output`` column:
::
>>> df.llm.predict("question", model="qwen3.6-plus")
Multiple input columns with per-column content types. ``content_type``
is passed through to
:class:`pyflink.dataframe.OpenAICompatProvider` as a per-call provider
option:
::
>>> df.llm.predict("question", "image_url",
... model="qwen3.6-plus",
... content_type=["TEXT", "IMAGE_URL"])
Free-form JSON in one VARIANT column named ``output``.
``response_format`` is passed through to
:class:`pyflink.dataframe.OpenAICompatProvider` as a per-call provider
option:
::
>>> df.llm.predict("question", model="qwen3.6-plus",
... output_type="VARIANT",
... response_format="json_object")
Named output column:
::
>>> df.llm.predict("question", model="qwen3.6-plus",
... output_type=pf.DataType.struct({
... "answer": pf.DataType.string()
... }))
Caching prediction results:
::
>>> df.llm.predict(
... "question", model="qwen3.6-plus",
... cache_table="`fluss-catalog`.default_database.predict_cache",
... cache_key=["tenant_id", "question"])
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.predict(
... "question", model="qwen3.6-plus",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.dataframe.dataframe import DataFrame as DF
from pyflink.dataframe.datatype import DataType
from pyflink.table.types import RowType
def validate_output_column_names(output_names: List[str]) -> None:
existing_columns = set(self._df.columns)
conflicting_columns = sorted(existing_columns.intersection(output_names))
if conflicting_columns:
raise ValueError(
"predict output columns conflicts with existing DataFrame "
"columns; use unique output field names: "
f"{conflicting_columns}")
def resolve_output_columns(
output_type: Optional[Union[str, "DataType"]]):
if output_type is None:
dtype = DataType.string()
output_columns = [("output", dtype._table_type)]
elif isinstance(output_type, str):
dtype = DataType._from_sql(output_type)
output_columns = [("output", dtype._table_type)]
elif isinstance(output_type, DataType):
table_type = output_type._table_type
if isinstance(table_type, RowType):
if len(table_type.fields) == 0:
raise ValueError(
"Struct output_type must contain at least one field.")
seen = set()
duplicate_names = []
for field_name in table_type.field_names():
if field_name in seen and field_name not in duplicate_names:
duplicate_names.append(field_name)
seen.add(field_name)
if duplicate_names:
raise ValueError(
"Struct output_type contains duplicate field names: "
f"{duplicate_names}")
output_columns = [
(field.name, field.data_type) for field in table_type.fields]
else:
output_columns = [("output", table_type)]
elif isinstance(output_type, Mapping):
raise TypeError(
"output_type mapping form is no longer supported; use a SQL "
"type string or DataType for a single 'output' column, or "
"DataType.struct({...}) for named output columns.")
else:
raise TypeError(
"output_type must be None, a SQL type string, or a DataType; "
f"got {type(output_type).__name__}.")
validate_output_column_names([name for name, _ in output_columns])
return output_columns
cache_keys = _validate_cache_args(cache_table, cache_key)
output_columns = resolve_output_columns(output_type)
# Build input schema from actual column types
input_columns = [
(col_name, self._get_column_type(col_name))
for col_name in input_cols
]
flink_model = _resolve_model(
self._df._table._t_env, model, provider,
input_columns, output_columns, provider_options=kwargs)
result = flink_model.predict(
self._df._table, list(input_cols),
cache_table=cache_table, cache_keys=cache_keys,
options=config)
return DF(result)
def _ai_func(self, func_name: str, input_col, model, provider,
cache_table: Optional[str], cache_key: Optional[CacheKey],
config, ai_func, *extra_args, provider_options=None):
"""Common logic for ai_* functions."""
from pyflink.dataframe.dataframe import DataFrame as DF
cache_keys = _validate_cache_args(cache_table, cache_key)
input_cols, output_cols = _ai_function_schemas()[func_name]
flink_model = _resolve_model(
self._df._table._t_env, model, provider,
input_cols, output_cols, provider_options=provider_options)
expr = ai_func(
flink_model, self._to_col_expr(input_col), *extra_args,
cache_table=cache_table, cache_keys=cache_keys, config=config)
return DF(self._df._table.join_lateral(expr))
[docs] def ai_classify(self, input_col: Union[str, "Expression"],
labels: List[str], *,
provider: Optional[str] = None,
model: Optional[str] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Classify text into one of the provided labels.
Args:
input_col: Column name (str) or Expression for the input text.
labels: List of label strings.
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name.
cache_table: Optional cache table identifier.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Note: This function uses its own task-specific prompt.
Provider configurations for the system prompt, such as
``system_prompt`` on ``OpenAICompatProvider`` and
``DashScopeProvider``, do not take effect for this call.
Returns:
A new DataFrame with columns appended:
- ``category`` (STRING): the predicted label.
- ``confidence`` (DOUBLE): confidence score.
Examples:
Simple invocation:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> df = pf.from_dict({"text": ["The delivery was fast."]})
>>> df.llm.ai_classify("text", ["positive", "negative"],
... model="qwen3.6-plus")
With cache:
::
>>> df.llm.ai_classify(
... "text", ["positive", "negative"], model="qwen3.6-plus",
... cache_table="`fluss-catalog`.default_database.classify_cache",
... cache_key="text")
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.ai_classify(
... "text", ["positive", "negative"], model="qwen3.6-plus",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.table.ai_functions import ai_classify as _ai_classify
return self._ai_func("ai_classify", input_col, model, provider,
cache_table, cache_key, config, _ai_classify,
labels, provider_options=kwargs)
[docs] def ai_sentiment(self, input_col: Union[str, "Expression"], *,
provider: Optional[str] = None,
model: Optional[str] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Analyze the sentiment of input text.
Args:
input_col: Column name (str) or Expression for the input text.
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name.
cache_table: Optional cache table identifier.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Note: This function uses its own task-specific prompt.
Provider configurations for the system prompt, such as
``system_prompt`` on ``OpenAICompatProvider`` and
``DashScopeProvider``, do not take effect for this call.
Returns:
A new DataFrame with columns appended:
- ``score`` (DOUBLE): sentiment score from -1.0 to 1.0.
- ``label`` (STRING): one of "positive", "negative", "neutral".
- ``confidence`` (DOUBLE): confidence score.
Examples:
Simple invocation:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> df = pf.from_dict({"review": ["The food was excellent."]})
>>> df.llm.ai_sentiment("review", model="qwen3.6-plus")
With cache:
::
>>> df.llm.ai_sentiment(
... "review", model="qwen3.6-plus",
... cache_table="`fluss-catalog`.default_database.sentiment_cache",
... cache_key="review")
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.ai_sentiment(
... "review", model="qwen3.6-plus",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.table.ai_functions import ai_sentiment as _ai_sentiment
return self._ai_func("ai_sentiment", input_col, model, provider,
cache_table, cache_key, config, _ai_sentiment,
provider_options=kwargs)
[docs] def ai_translate(self, input_col: Union[str, "Expression"],
source_lang: str, target_lang: str, *,
provider: Optional[str] = None,
model: Optional[str] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Translate text from one language to another.
Args:
input_col: Column name (str) or Expression for the input text.
source_lang: Source language code (e.g. ``"zh"``, ``"en"``,
``"auto"`` for auto-detection). Supported: ``auto``, ``zh``,
``en``, ``ja``, ``ko``, ``fr``, ``de``, ``es``, ``ru``,
``ar``, ``pt``.
target_lang: Target language code (cannot be ``"auto"``).
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name.
cache_table: Optional cache table identifier.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Note: This function uses its own task-specific prompt.
Provider configurations for the system prompt, such as
``system_prompt`` on ``OpenAICompatProvider`` and
``DashScopeProvider``, do not take effect for this call.
Returns:
A new DataFrame with columns appended:
- ``translated_text`` (STRING): the translated text.
- ``detected_language`` (STRING): detected source language code.
Examples:
Simple invocation:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> df = pf.from_dict({"text": ["你好"]})
>>> df.llm.ai_translate("text", "zh", "en", model="qwen3.6-plus")
With cache:
::
>>> df.llm.ai_translate(
... "text", "zh", "en", model="qwen3.6-plus",
... cache_table="`fluss-catalog`.default_database.translate_cache",
... cache_key="text")
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.ai_translate(
... "text", "zh", "en", model="qwen3.6-plus",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.table.ai_functions import ai_translate as _ai_translate
return self._ai_func("ai_translate", input_col, model, provider,
cache_table, cache_key, config, _ai_translate,
source_lang, target_lang, provider_options=kwargs)
[docs] def ai_summarize(self, input_col: Union[str, "Expression"],
max_length: int, *, provider: Optional[str] = None,
model: Optional[str] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Summarize text to a maximum length.
Args:
input_col: Column name (str) or Expression for the input text.
max_length: Maximum length of the summary in characters.
Must be greater than 0.
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name.
cache_table: Optional cache table identifier.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Note: This function uses its own task-specific prompt.
Provider configurations for the system prompt, such as
``system_prompt`` on ``OpenAICompatProvider`` and
``DashScopeProvider``, do not take effect for this call.
Returns:
A new DataFrame with a column appended:
- ``summary`` (STRING): the summarized text.
Examples:
Simple invocation:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> df = pf.from_dict({"article": ["Long article text..."]})
>>> df.llm.ai_summarize("article", 200, model="qwen3.6-plus")
With cache:
::
>>> df.llm.ai_summarize(
... "article", 200, model="qwen3.6-plus",
... cache_table="`fluss-catalog`.default_database.summarize_cache",
... cache_key="article")
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.ai_summarize(
... "article", 200, model="qwen3.6-plus",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.table.ai_functions import ai_summarize as _ai_summarize
return self._ai_func("ai_summarize", input_col, model, provider,
cache_table, cache_key, config, _ai_summarize,
max_length, provider_options=kwargs)
[docs] def ai_mask(self, input_col: Union[str, "Expression"],
entities: List[str], *, provider: Optional[str] = None,
model: Optional[str] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Mask sensitive information in text.
Args:
input_col: Column name (str) or Expression for the input text.
entities: List of entity types to mask (e.g. ``["name", "phone"]``).
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name.
cache_table: Optional cache table identifier.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Note: This function uses its own task-specific prompt.
Provider configurations for the system prompt, such as
``system_prompt`` on ``OpenAICompatProvider`` and
``DashScopeProvider``, do not take effect for this call.
Returns:
A new DataFrame with columns appended:
- ``masked_text`` (STRING): text with sensitive info replaced.
- ``detected_entities`` (ARRAY<STRING>): list of detected entity
types.
Examples:
Simple invocation:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="chat/completions"))
>>> df = pf.from_dict({"text": ["Alice called 555-0100."]})
>>> df.llm.ai_mask("text", ["name", "phone"], model="qwen3.6-plus")
With cache:
::
>>> df.llm.ai_mask(
... "text", ["name", "phone"], model="qwen3.6-plus",
... cache_table="`fluss-catalog`.default_database.mask_cache",
... cache_key="text")
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.ai_mask(
... "text", ["name", "phone"], model="qwen3.6-plus",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.table.ai_functions import ai_mask as _ai_mask
return self._ai_func("ai_mask", input_col, model, provider,
cache_table, cache_key, config, _ai_mask,
entities, provider_options=kwargs)
[docs] def ai_embed(self, input_col: Union[str, "Expression"],
dimension: int = 1024, *, provider: Optional[str] = None,
model: Optional[str] = None,
cache_table: Optional[str] = None,
cache_key: Optional[CacheKey] = None,
config: Optional[Dict[str, str]] = None,
**kwargs) -> "DataFrame":
"""
Generate embedding vectors for text.
Args:
input_col: Column name (str) or Expression for the input text.
dimension: Dimension of the embedding vector (default 1024).
provider: Registered model provider lookup name. Uses default if
not specified. If no provider is configured, ``model`` is
treated as a catalog model name.
model: Model name.
cache_table: Pre-registered catalog table identifier for the
embedding cache.
cache_key: Optional cache key column name or list of column names.
config: Optional dict of runtime prediction options.
**kwargs: Per-call overrides for the selected DataFrame model
provider's constructor parameters. Overrides are applied only
when the call uses a configured provider.
Returns:
A new DataFrame with a column appended:
- ``embedding`` (ARRAY<FLOAT>): the embedding vector.
Examples:
Simple invocation:
::
>>> import pyflink.dataframe as pf
>>> pf.set_model_provider(pf.OpenAICompatProvider(
... task="embeddings"))
>>> df = pf.from_dict({
... "tenant_id": ["tenant-a"],
... "text": ["Hello world"],
... })
>>> df.llm.ai_embed("text", 512, model="text-embedding-v4")
With cache:
::
>>> df.llm.ai_embed(
... "text", 512, model="text-embedding-v4",
... cache_table="`fluss-catalog`.default_database.embed_cache",
... cache_key="text")
Using runtime config supported by
:class:`pyflink.dataframe.OpenAICompatProvider`:
::
>>> df.llm.ai_embed(
... "text", 512, model="text-embedding-v4",
... config={"max-concurrent-operations": "20"})
"""
from pyflink.dataframe.dataframe import DataFrame as DF
from pyflink.table.ai_functions import ai_embed as _ai_embed
cache_keys = _validate_cache_args(cache_table, cache_key)
input_cols, output_cols = _ai_function_schemas()["ai_embed"]
flink_model = _resolve_model(
self._df._table._t_env, model, provider, input_cols, output_cols,
provider_options=kwargs)
expr = _ai_embed(
flink_model, self._to_col_expr(input_col),
dimension, cache_table=cache_table, cache_keys=cache_keys,
config=config)
return DF(self._df._table.join_lateral(expr))