Python DataFrame API 参考

注册标量函数

udf 装饰器支持将 Python 函数注册为 DataFrame 可用的标量函数。

• 接口名称：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
  )

• 功能描述：将 Python 函数注册为 DataFrame 可用的标量用户自定义函数。支持注册普通函数、async 函数、ScalarFunction / AsyncScalarFunction 子类及可调用类。

• 入参

  参数名称	参数类型	是否必填	描述
  func	Callable / Class	否	要包装的 Python 函数、ScalarFunction 实例/子类或可调用类。若省略则返回装饰器。
  return_dtype	DataType / str / type	否	返回类型。可以是 DataType 实例（如 DataType.int64()）、Python 类型（如 int）或 SQL 类型字符串（如 'BIGINT'）。若省略则从函数返回类型提示自动推断。
  deterministic	Boolean	否	函数是否是确定性的。默认值为 True。
  name	String	否	UDF 的名称。若不指定则使用函数名。
  func_type	String	否	执行格式，可选 "general"、"pandas"、"arrow"。若不指定则根据参数类型提示自动检测。
  concurrency	int	否	UDF 算子的并行度。不同 concurrency 值的 UDF 会被拆分到独立算子。
  batch_size	int	否	每批最大元素数，仅适用于批处理 UDF（pandas / arrow 模式）。
  num_gpus	float	否	为该 UDF 申请的 GPU 数量（如 0.5、1）。设置后，该 UDF 会在独立算子中运行，不与其他 UDF 融合（包括其它 GPU UDF）。
  gpu_type	String	否	GPU 类型。当指定 num_gpus 时必填。

• 返回值

  DataFrameUDFWrapper 对象，可在 with_column、with_columns、map、map_batches等操作中使用。

同步行式 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"

@udf
class ModelPredictor:
    def __init__(self):
        # Initialization is executed on the task manager
        self._model = load_model()

    def __call__(self, amount: float) -> str:
        return self._model.predict(amount)

# 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")),
)

# Use stateful UDFs
orders_with_predictions = orders.with_column(
    "prediction", ModelPredictor(col("amount"))
)

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)

同步向量 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(),
        }
    ),
)

异步行式 UDF

异步 UDF 允许多行同时执行 I/O 操作，从而提升吞吐，适合调用外部服务或执行 I/O 密集型逻辑。您可在with_column、with_columns 和 map 中调用异步行式 UDF。

import asyncio

# Use in `with_column`
@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")),
)

# Use in `map`
class OrderWithRegion(TypedDict):
    order_id: int
    region: str

async def enrich_order_region(row) -> OrderWithRegion:
    await asyncio.sleep(0.01)  # Mock async I/O
    return {
        "order_id": row["order_id"],
        "region": f"region_for_{row['user_id']}",
    }

orders_with_region = orders.map(enrich_order_region)

异步向量 UDF

异步向量 UDF 兼顾批处理吞吐与异步 I/O 并发。可以通过 batch_size 控制每批大小。您可在with_column、with_columns 和 map_batches 中调用异步向量 UDF。

# Use in `with_column`
@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")),
)

# Use in `map_batches`
async def enrich_order_batch(
    batch: dict[str, pd.Series],
) -> dict[str, pd.Series]:
    async def query_region(user_id):
        await asyncio.sleep(0.01)  # 模拟异步 I/O
        return f"region_for_{user_id}"

    regions = await asyncio.gather(
        *(query_region(user_id) for user_id in batch["user_id"])
    )

    return {
        "order_id": batch["order_id"],
        "region": pd.Series(
            regions,
            index=batch["user_id"].index,
        ),
    }

orders_with_region = orders.map_batches(
    enrich_order_batch,
    batch_format="pandas",
    batch_size=64,
    concurrency=32,
    return_dtype=DataType.struct({
        "order_id": DataType.int64(),
        "region": DataType.string(),
    }),
)

注册表值函数

udtf 装饰器支持将 Python 函数注册为 DataFrame 可用的表值函数。

• 接口名称：udtf。

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

• 功能描述：将 Python 函数、TableFunction 子类/实例或可调用类注册为 DataFrame 可用的表值函数。

• 入参

  参数名称	参数类型	是否必填	描述
  func	Callable / Class	否	要包装的 Python 函数、TableFunction 实例/子类或可调用类。若省略则返回装饰器。
  return_dtype	DataType / str / type	否	每个输出行的类型。多列输出可使用 DataType.struct({...}) 或 TypedDict。
  deterministic	Boolean	否	函数是否是确定性的。默认值为 True。
  name	String	否	UDTF 的名称。若不指定则使用函数名。
  concurrency	int	否	UDTF 算子的并行度。
  num_gpus	float	否	为该 UDTF 申请的 GPU 数量（如 0.5、1）。设置后，该 UDTF 会在独立算子中运行，不与其他 UDF 融合（包括其它 GPU UDF）。
  gpu_type	String	否	GPU 类型。当指定 num_gpus 时必填。

• 返回值

  DataFrameUDTFWrapper 对象，可在 flat_map 或 join_lateral 中使用。

在join_lateral中使用

join_lateral会把 UDTF 结果追加到原始输入行后面。默认删除 UDTF 无输出的输入行，您可设置 ignore_empty=False 保留 UDTF 无输出的输入行。

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")),
)

在flat_map中使用

flat_map对每行调用 UDTF，结果只包含 UDTF 输出列，不保留原始输入列。您可直接将 Python 函数传入 flat_map，而无需使用 udtf 装饰器包装。

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()}),
)

文本分类

• 接口名称：ai_classify。

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

• 功能描述：将文本分类到指定标签列表中的某一个。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  labels	List[String]	是	分类标签列表，如 ["positive", "negative", "neutral"]。

  其余参数的含义参见通用调用。

  说明

  ai_classify等垂类 AI 函数会使用内置 system prompt，Model Provider 配置的 system prompt 不生效。

• 返回值

  DataFrame 对象，追加 category (STRING) 和 confidence (DOUBLE) 列，包含分类结果和置信分数。

• 示例

  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")

情感分析

• 接口名称：ai_sentiment。

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

• 功能描述：对输入文本进行情感分析。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。

  其余参数的含义参见通用调用。

• 返回值

  DataFrame 对象，追加以下列，包含情感分析结果：

  • score (DOUBLE): 情感分析分数，在 -1.0 到 1.0 之间。

  • label (STRING): “positive”, “negative”, “neutral” 之一。

  • confidence (DOUBLE): 置信分数。

• 示例

  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")

信息提取

• 接口名称：ai_extract。

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

• 功能描述：按照给定的 JSON Schema 从文本中提取结构化信息。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  schema	String	是	JSON Schema 字符串，描述要提取的字段，如 '{"name":"STRING", "phone":"STRING"}'。

  其余参数的含义参见通用调用。

• 返回值

  DataFrame 对象，追加 extracted_json (STRING) 列，包含提取的结构化信息。

• 示例

  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")

文本翻译

• 接口名称：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
  )

• 功能描述：将文本从源语言翻译为目标语言。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  source_lang	String	是	源语言代码，如 "zh"、"en"、"auto"（自动检测）。支持：auto、zh、en、ja、ko、fr、de、es、ru、ar、pt。
  target_lang	String	是	目标语言代码，不可为 "auto"。

  其余参数的含义参见通用调用。

• 返回值

  DataFrame 对象，追加 translated_text (STRING) 和 detected_language (STRING) 列，包含翻译结果和识别到的语言。

• 示例

  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")

文本摘要

• 接口名称：ai_summarize。

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

• 功能描述：将文本摘要到指定最大长度。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  max_length	int	是	摘要的最大字符数，必须大于 0。

  其余参数的含义参见通用调用。

• 返回值

  DataFrame 对象，追加 summary (STRING) 列，包含摘要文本。

• 示例

  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")
  

数据脱敏

• 接口名称：ai_mask。

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

• 功能描述：对文本中的敏感信息进行脱敏处理。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  entities	List[String]	是	需要脱敏的实体类型列表，如 ["name", "phone"]。

  其余参数的含义参见通用调用。

• 返回值

  DataFrame 对象，追加 masked_text (STRING) 和 detected_entities (ARRAY<STRING>) 列，包含脱敏后的文字和检测到的实体。

• 示例

  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。

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

• 功能描述：使用当前 DataFrame 中的查询向量，到向量检索源中查找最相似的记录，并把检索结果追加到当前 DataFrame。

• 入参

  参数名称	参数类型	是否必填	描述
  search_source	DataFrame	是	向量检索源 DataFrame，例如通过 read_milvus 读取的向量集合。
  column_to_search	String	是	检索源中的向量列名。
  column_to_query	String / Expression	是	当前 DataFrame 中的查询向量列名或列表达式。
  top_k	int	是	每条输入记录返回的最相似结果数量。
  agg	Boolean	否	是否把 Top K 结果聚合到一个数组列中。默认值为 False。
  output_columns	List[String] / Mapping[String, String]	否	可选的向量搜索结果列。列表用于选择列并保留原列名；agg=False 时还可以使用映射，其键用于选择列、值作为输出别名。可选择 search_source 中的顶层列以及引擎生成的 score 列。默认值为 None，表示选择全部结果列。
  config	Dict	否	运行时向量搜索配置。
  ignore_empty	Boolean	否	是否丢弃没有检索结果的输入行。默认值为 False。

• 返回值

  DataFrame 对象，包含原始列和追加的向量检索结果列。设 output_columns 最终选择 M 个结果列：

  • 未指定 `output_columns` 时，默认返回 search_source 的全部列，以及一个 DOUBLE 类型的检索分值列，默认名为 score。分值的含义和排序方向由 Connector 及其距离度量决定；使用 Milvus 时请结合 search_metric 解释该值。

  • 当 `agg=False` 时，每个输入行最多会对应 `top_k` 条输出行，并追加 M 列。列表按照声明顺序选择列并保留原名；映射按照键的声明顺序选择列，并使用对应的值作为输出别名。

  • 当 `agg=True` 时，每条有检索结果的输入最多产生一条输出，并追加固定名为 search_results 的 ARRAY<ROW<...>> 列。默认 ignore_empty=False 时会保留所有输入记录；没有检索结果时 search_results 为 NULL。设置 ignore_empty=True 后，没有检索结果的输入记录不会出现在结果中。数组元素为包含 M 个所选字段的 ROW。此模式仅支持使用列表投影内部字段，不支持映射。可以在调用完成后使用 `rename_columns("search_results", "新列名")` 重命名外层数组列。

  • `output_columns` 不能为空，列名和别名均不能为空或重复。所选列必须是 `search_source` 的顶层列或者生成的相似度列。返回顺序与列表或映射键的声明顺序一致。

  • 如果 `search_source` 已存在 `score`、`score0` 等列，引擎会为生成的相似度列选择唯一名称，例如 `score2`；此时需要在 `output_columns` 中使用实际生成的列名。

  • 非聚合模式下，可以使用映射为结果列指定不冲突的别名。如果列表或默认输出与当前 DataFrame 的已有列重名，请先使用 `rename_columns` 重命名当前 DataFrame 中的冲突列，再调用 `vector_search`。聚合模式下，如果当前 DataFrame 已有 `search_results` 列，也需要先重命名该列。

  • 如果向量检索源使用的连接器支持列裁剪下推，引擎会尝试将 `output_columns` 对应的列裁剪下推到连接器；无论连接器是否支持下推，最终结果都只包含所选的向量搜索结果列。

• VVR 11.8升级说明

  VVR 11.8中的 output_columns 表示“全部输出列的别名”，VVR 11.9及以上版本中则表示“需要选择的结果列”。例如，VVR 11.8中的以下调用：

  output_columns=["match_doc_id", "match_embedding", "match_title", "match_score"]

  升级后应改为映射形式：

  output_columns={
      "doc_id": "match_doc_id",
      "embedding": "match_embedding",
      "title": "match_title",
      "score": "match_score",
  }

• 示例

  from array import array
  import pyflink.dataframe as pf
  from pyflink.dataframe import DataType
  
  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",
  )
  
  selected_docs = query_df.llm.vector_search(
      documents,
      column_to_search="embedding",
      column_to_query="query_embedding",
      top_k=3,
      output_columns=["doc_id", "title", "score"],
  )

  列表形式会投影结果列，并保留原列名。需要同时指定别名时，可以使用映射：

  matched_docs = query_df.llm.vector_search(
      documents,
      column_to_search="embedding",
      column_to_query="query_embedding",
      top_k=3,
      output_columns={
          "doc_id": "match_doc_id",
          "title": "match_title",
          "score": "match_score",
      },
  )

  聚合模式下，列表选择数组元素中的字段，外层数组列固定名为 search_results。如需重命名外层数组列，可以调用 rename_columns：

  matched_docs_agg = query_df.llm.vector_search(
      documents,
      column_to_search="embedding",
      column_to_query="query_embedding",
      top_k=3,
      agg=True,
      output_columns=["doc_id", "title", "score"],
  ).rename_columns("search_results", "matches")

文本向量化

• 接口名称：ai_embed。

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

• 功能描述：为文本生成向量嵌入（Embedding）。

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  dimension	int	否	向量维度。默认值为 1024。

  其余参数的含义参见通用调用。

• 返回值

  DataFrame 对象，追加 embedding (ARRAY<FLOAT>) 列，包含生成的向量。

• 示例

  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")