Python DataFrame API reference

Register scalar functions

The udf decorator registers a Python function as a scalar function for use in DataFrames.

• API name: udf.

  udf(
   func=None,
   *,
   return_dtype=None,
   deterministic=True,
   name=None,
   func_type=None,
   concurrency=None,
   batch_size=None,
   num_gpus=None,
   gpu_type=None
  )

• Description: Registers a Python function as a DataFrame scalar UDF. Supports regular functions, async functions, ScalarFunction / AsyncScalarFunction subclasses, and callable classes.

• Parameters

  Parameter name	Parameter type	Required	Description
  func	Callable / Class	No	The Python function, ScalarFunction instance/subclass, or callable class to wrap. If omitted, a decorator is returned.
  return_dtype	DataType / str / type	No	Return type. Can be DataType instance (for example, DataType.int64()), Python types (for example, int), or SQL type strings (for example, 'BIGINT'). If omitted, it is auto-inferred from the function type hints.
  deterministic	Boolean	No	Whether the function is deterministic. Default: True.
  name	String	No	Name of the UDF. Defaults to the function name.
  func_type	String	No	Execution format. Options: "general", "pandas", "arrow". If omitted, it is auto-detected from the function type hints.
  concurrency	int	No	Parallelism of the UDF operator. UDFs with different concurrency values are split into separate operators.
  batch_size	int	No	Maximum number of elements per batch. Applies only to batch UDFs (pandas / arrow mode).
  num_gpus	float	No	Number of GPUs to request for the UDF (for example, 0.5, 1). When set, the UDF runs in a separate operator and is not chained with other UDFs (including other GPU UDFs).
  gpu_type	String	No	GPU type. Required when num_gpus is specified.

• Return value

  A DataFrameUDFWrapper object that can be used in with_column, with_columns, map, map_batches, and similar operations.

Synchronous row UDF

from pyflink.dataframe import col, udf

orders = pf.from_records(
 [
 (1001, 1, 99.9, "PAID"),
 (1002, 2, 35.5, "CREATED"),
 (1003, 1, 188.0, "PAID"),
 (1004, 3, 88.8, None),
 ],
 schema=["order_id", "user_id", "amount", "status"],
)

@udf
def normalize_status(status: str) -> str:
 if status is None:
 return "UNKNOWN"
 return status.strip().upper()

@udf(concurrency=32)
def order_tag(amount: float, status: str) -> str:
 """Takes multiple columns as input"""
 if status == "PAID" and amount is not None and amount >= 100:
 return "high_value_paid"
 return "normal"

# with_column adds a single column
orders_with_status = orders.with_column(
 "status_norm", normalize_status(col("status"))
)

# with_columns adds multiple columns at once
orders_with_flags = orders.with_columns(
 status_norm=normalize_status(col("status")),
 tag=order_tag(col("amount"), col("status")),
)

def build_order_feature(row):
 amount = row["amount"] or 0.0
 return {
 "order_id": row["order_id"],
 "user_id": row["user_id"],
 "feature": f"{row['status']}:{'large' if amount >= 100 else 'normal'}",
 }

order_features = orders.map(
 build_order_feature,
 return_dtype=DataType.struct({
 "order_id": DataType.int64(),
 "user_id": DataType.int64(),
 "feature": DataType.string(),
 }),
)

from typing import TypedDict

class OrderFeature(TypedDict):
 order_id: int
 user_id: int
 feature: str

def build_order_feature_typed(row) -> OrderFeature:
 amount = row["amount"] or 0.0
 return {
 "order_id": row["order_id"],
 "user_id": row["user_id"],
 "feature": f"{row['status']}:{'large' if amount >= 100 else 'normal'}",
 }

order_features_typed = orders.map(build_order_feature_typed)

Synchronous vector UDF

import pandas as pd

@udf(return_dtype=DataType.float64(), concurrency=32, batch_size=64)
def scale_amount_pandas(amounts: pd.Series) -> pd.Series:
 return amounts * 100.0

orders_scaled_pandas = orders.with_column(
 "amount_scaled", scale_amount_pandas(col("amount")),
)

import pyarrow as pa
import pyarrow.compute as pc

@udf(return_dtype=DataType.float64(), concurrency=32, batch_size=64)
def scale_amount_arrow(amounts: pa.Array) -> pa.Array:
 return pc.multiply(amounts, 100.0)

orders_scaled_arrow = orders.with_column(
 "amount_scaled", scale_amount_arrow(col("amount")),
)

def score_batch_pandas(batch: dict[str, pd.Series]) -> dict[str, pd.Series]:
 amount = batch["amount"].fillna(0.0)
 return {
 "order_id": batch["order_id"],
 "score": (amount / 100.0).clip(0.0, 1.0),
 }

order_scores_pandas = orders.map_batches(
 score_batch_pandas,
 batch_format="pandas",
 batch_size=1024,
 return_dtype=DataType.struct(
 {
 "order_id": DataType.int64(),
 "score": DataType.float64(),
 }
 ),
)

def score_batch_arrow(batch: dict[str, pa.Array]) -> dict[str, pa.Array]:
 amount = pc.if_else(pc.is_null(batch["amount"]), 0.0, batch["amount"])
 raw_score = pc.divide(pc.cast(amount, pa.float64()), 100.0)
 score = pc.if_else(
 pc.less(raw_score, 0.0),
 0.0,
 pc.if_else(pc.greater(raw_score, 1.0), 1.0, raw_score),
 )
 return {
 "order_id": batch["order_id"],
 "score": score,
 }

order_scores_arrow = orders.map_batches(
 score_batch_arrow,
 batch_format="arrow",
 batch_size=1024,
 return_dtype=DataType.struct(
 {
 "order_id": DataType.int64(),
 "score": DataType.float64(),
 }
 ),
)

Asynchronous row UDF

Asynchronous UDFs let multiple rows perform I/O operations concurrently, which improves throughput. They suit logic that calls external services or that is I/O intensive. You can invoke asynchronous row-based UDFs in with_column or with_columns.

import asyncio

@udf(concurrency=32)
async def query_region(user_id: int) -> str:
 await asyncio.sleep(0.01) # simulate async I/O
 return f"region_for_{user_id}"

orders_with_region = orders.with_column(
 "region", query_region(col("user_id")),
)

Asynchronous vector UDF

Asynchronous vectorized UDFs combine batch processing throughput with asynchronous I/O concurrency. Use batch_size to control the size of each batch. You can invoke asynchronous vectorized UDFs in with_column or with_columns.

@udf(return_dtype=DataType.string(), concurrency=32, batch_size=64)
async def batch_enrich(statuses: pd.Series) -> pd.Series:
 async def enrich_one(s):
 await asyncio.sleep(0.01) # simulate async API call
 return f"enriched_{s}"
 tasks = [enrich_one(s) for s in statuses]
 results = await asyncio.gather(*tasks)
 return pd.Series(results)

orders_enriched = orders.with_column(
 "status_enriched", batch_enrich(col("status")),
)

Register table functions

The udtf decorator registers a Python function as a table function for use in DataFrames.

• API name: udtf.

  udtf(
   func=None,
   *,
   return_dtype=None,
   deterministic=True,
   name=None,
   concurrency=None,
   num_gpus=None,
   gpu_type=None
  )

• Description: Registers a Python function, a TableFunction subclass or instance, or a callable class as a table function for use in DataFrames.

• Parameters

  Parameter name	Parameter type	Required	Description
  func	Callable / Class	No	The Python function, TableFunction instance/subclass, or callable class to wrap. If omitted, a decorator is returned.
  return_dtype	DataType / str / type	No	Type of each output row. For multi-column output, use DataType.struct({...}) or TypedDict.
  deterministic	Boolean	No	Whether the function is deterministic. Default: True.
  name	String	No	Name of the UDTF. Defaults to the function name.
  concurrency	int	No	Parallelism of the UDTF operator.
  num_gpus	float	No	Number of GPUs to request for the UDTF (for example, 0.5, 1). When set, the UDTF runs in a separate operator and is not chained with other UDFs (including other GPU UDFs).
  gpu_type	String	No	GPU type. Required when num_gpus is specified.

• Return value

  A DataFrameUDTFWrapper object that can be used in flat_map or join_lateral.

In join_lateral Used in

join_lateral appends UDTF results after the original input row. By default, input rows that produce no UDTF output are removed; you can set ignore_empty=False to keep input rows with no UDTF output.

from typing import Iterator, Tuple
from pyflink.dataframe import udtf

# Use type hint and .alias() to specify return types and column names
@udtf
def split_words(text: str) -> Iterator[Tuple[str, int]]:
 for word in text.split():
 yield word, len(word)

words_with_source = texts.join_lateral(
 split_words(col("text")).alias("word", "word_length"),
 ignore_empty=False,
)

# Or use return_dtype
@udtf(return_dtype=DataType.struct({
 "word": DataType.string(),
 "word_length": DataType.int32(),
}))
def split_words_with_length(text):
 for word in text.split():
 yield {"word": word, "word_length": len(word)}

words_with_detail = texts.join_lateral(
 split_words_with_length(col("text")),
)

In flat_map Used in

flat_map calls the UDTF on each row. The result contains only the UDTF output columns, and the original input columns are not preserved. You can pass a Python function directly to flat_map without wrapping it with the udtf decorator.

from typing import Any, Dict, Iterator, TypedDict

# Use type hint with TypedDict to specify return types and column names
class Word(TypedDict):
 word: str
 length: int

def split_row(row: Dict[str, Any]) -> Iterator[Word]:
 for word in row["text"].split():
 yield {"word": word, "length": len(word)}

words = texts.flat_map(split_row)

# Or use return_dtype
def split_row_with_dtype(row):
 for word in row["text"].split():
 yield {"word": word}

words_with_dtype = texts.flat_map(
 split_row_with_dtype,
 return_dtype=DataType.struct({"word": DataType.string()}),
)

Text classification

• API name: ai_classify.

  DataFrame.llm.ai_classify(
   input_col,
   labels,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )

• Description: Classifies text into one of the specified labels.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  labels	List[String]	Yes	A list of classification labels, for example, ["positive", "negative", "neutral"].

  See the descriptions of other parameters at Generic invocation.

  Note

  Domain-specific AI functions such as ai_classify use a built-in system prompt. The system prompt configured on the Model Provider does not take effect.

• Return value

  A DataFrame with the category (STRING) and confidence (DOUBLE) columns appended, which contain the classification result and the confidence score.

• Example

  pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
  df = pf.from_dict({"review": ["Great product", "Terrible, do not buy", "It is okay"]})
  df = df.llm.ai_classify("review",
   labels=["positive", "negative", "neutral"],
   model="qwen3.6-plus")

Sentiment analysis

• API name: ai_sentiment.

  DataFrame.llm.ai_sentiment(
   input_col,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )

• Description: Performs sentiment analysis on the input text.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.

  See the descriptions of other parameters at Generic invocation.

• Return value

  A DataFrame with the following columns appended, which contain the sentiment analysis result:

  • score (DOUBLE): The sentiment analysis score, ranging from -1.0 to 1.0.

  • label  (STRING): One of "positive", "negative", or "neutral".

  • confidence  (DOUBLE): Confidence score.

• Example

  pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
  df = pf.from_dict({"comment": ["This feature is amazing!", "Broke after one month"]})
  df = df.llm.ai_sentiment("comment", model="qwen3.6-plus")

Information extraction

• API name: ai_extract.

  DataFrame.llm.ai_extract(
   input_col,
   schema,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )
  

• Description: Extracts structured information from text based on a given JSON Schema.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  schema	String	Yes	A JSON Schema string that describes the fields to extract, for example, '{"name":"STRING", "phone":"STRING"}'.

  See the descriptions of other parameters at Generic invocation.

• Return value

  A DataFrame with the extracted_json (STRING) column appended, which contains the extracted structured information.

• Example

  pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
  df = pf.from_dict({"text": ["John Smith, male, 28 years old, phone ***-****"]})
  schema = '{"name": "STRING", "age": "INTEGER", "phone": "STRING"}'
  df = df.llm.ai_extract("text", schema=schema,
   model="qwen3.6-plus")

Text translation

• API name: ai_translate.

  DataFrame.llm.ai_translate(
   input_col,
   source_lang,
   target_lang,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )

• Description: Translates text from a source language to a target language.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  source_lang	String	Yes	Source language code, for example: "zh", "en", "auto" (auto-detected). Supported: auto, zh, en, ja, ko, fr, de, es, ru, ar, pt.
  target_lang	String	Yes	The target language code. This value cannot be "auto".

  See the descriptions of other parameters at Generic invocation.

• Return value

  A DataFrame with the translated_text (STRING) and detected_language (STRING) columns appended, which contain the translated text and the detected language.

• Example

  pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
  df = pf.from_dict({"text": ["Hello World", "How are you?"]})
  df = df.llm.ai_translate("text", source_lang="en", target_lang="zh",
   model="qwen3.6-plus")

Text summarization

• API name: ai_summarize.

  DataFrame.llm.ai_summarize(
   input_col,
   max_length,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )
  

• Description: Summarizes text to a specified maximum length.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  max_length	int	Yes	Maximum number of characters in the summary. Must be greater than 0.

  See the descriptions of other parameters at Generic invocation.

• Return value

  A DataFrame with the summary (STRING) column appended, which contains the summary text.

• Example

  pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
  df = pf.from_dict({"article": ["This is a long article with lots of content..."]})
  df = df.llm.ai_summarize("article", max_length=100,
   model="qwen3.6-plus")
  

Data masking

• API name: ai_mask.

  DataFrame.llm.ai_mask(
   input_col,
   entities,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )
  

• Description: Masks sensitive information in text.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  entities	List[String]	Yes	A list of entity types to mask, for example, ["name", "phone"].

  See the descriptions of other parameters at Generic invocation.

• Return value

  A DataFrame with the masked_text (STRING) and detected_entities (ARRAY<STRING>) columns appended, which contain the masked text and the detected entities.

• Example

  pf.set_model_provider(pf.OpenAICompatProvider(task="chat/completions"))
  df = pf.from_dict({"text": ["Please contact John Smith at 555-0123"]})
  df = df.llm.ai_mask("text", entities=["name", "phone"],
   model="qwen3.6-plus")

Vector search

• API name: vector_search.

  DataFrame.llm.vector_search(
   search_source,
   column_to_search,
   column_to_query,
   top_k,
   *,
   agg=False,
   output_columns=None,
   config=None,
   ignore_empty=False
  )
  

• Description: Searches the vector search source for the most similar records by using the query vectors in the current DataFrame, and appends the search results to the current DataFrame.

• Parameters

  Parameter name	Parameter type	Required	Description
  search_source	DataFrame	Yes	The vector search source DataFrame, for example, a vector collection read by using read_milvus.
  column_to_search	String	Yes	The name of the vector column in the search source.
  column_to_query	String / Expression	Yes	The query vector column name or column expression in the current DataFrame.
  top_k	int	Yes	Number of most similar records returned per input record.
  agg	Boolean	No	Specifies whether to aggregate the top K results into a single array column. Default value: False.
  output_columns	String / List[String]	No	The output column names. When agg=False, specify a name for each search source output column and for the score column. When agg=True, specify a single array column name.
  config	Dict	No	The runtime vector search configuration.
  ignore_empty	Boolean	No	Specifies whether to drop input rows that have no search results. Default value: False.

• Return value

  A DataFrame that contains the original columns plus the appended vector search result columns. Assume that the search source search_source has N columns:

  • When agg=False, each input row corresponds to at most top_k output rows, and N + 1 columns are appended: the N columns from search_source are appended to the result with their original types preserved, followed by a score column of the DOUBLE type. If you specify output_columns, you must provide N + 1 column names, which are used in order as the names of these appended columns.

  • When agg=True, each input record corresponds to one output record, and one array column is appended. Each array element is a row with N + 1 nested fields: the first N fields come from search_source and preserve their original types, and the last field is a score of the DOUBLE type. If you specify output_columns, you must provide exactly one column name, which is used as the name of this appended array column.

• Example

  from array import array
  
  query_df = pf.from_dict({
   "query_id": [1],
   "query_embedding": [array("f", [0.1, 0.2, 0.3])],
  })
  
  documents = pf.read_milvus(
   endpoint="http://milvus.example.com",
   username="${secret_values.milvus_user}",
   password="${secret_values.milvus_password}",
   database_name="commerce",
   collection_name="support_docs",
   schema={
   "doc_id": DataType.int64(),
   "embedding": DataType.list(DataType.float32()),
   "title": DataType.string(),
   },
   search_metric="COSINE",
  )
  
  matched_docs = query_df.llm.vector_search(
   documents,
   column_to_search="embedding",
   column_to_query="query_embedding",
   top_k=3,
   output_columns=["doc_id", "doc_embedding", "doc_title", "score"],
  )

Text embedding

• API name: ai_embed.

  DataFrame.llm.ai_embed(
   input_col,
   dimension=1024,
   *,
   provider=None,
   model=None,
   cache_table=None,
   cache_key=None,
   config=None,
   **kwargs
  )
  

• Description: Generates vector embeddings for text.

• Parameters

  Parameter name	Parameter type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  dimension	int	No	The vector dimension. Default value: 1024.

  See the descriptions of other parameters at Generic invocation.

• Return value

  A DataFrame with the embedding (ARRAY<FLOAT>) column appended, which contains the generated vector.

• Example

  pf.set_model_provider("embedding", pf.OpenAICompatProvider(task="embeddings"))
  df = pf.from_dict({"text": ["stream processing with Flink", "real-time analytics"]})
  df = df.llm.ai_embed("text", dimension=1024,
   provider="embedding", model="text-embedding-v4")