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 类型。

• 返回值

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

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

同步向量 UDF

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

import pyarrow as pa
import pyarrow.compute as pc

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中调用异步行式 UDF。

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

异步向量 UDF

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

import asyncio

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

注册表值函数

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 类型。

• 返回值

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

在join_lateral中使用

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

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

# 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
from pyflink.dataframe import DataType

# 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,
      config=None
  )

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

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  labels	List[String]	是	分类标签列表，如 ["positive", "negative", "neutral"]。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

• 示例

  df = pf.from_dict({"review": ["Great product", "Terrible, do not buy", "It is okay"]})
  df = df.llm.ai_classify("review",
                          labels=["positive", "negative", "neutral"],
                          provider="chat", model="qwen-plus")

情感分析

• 接口名称：ai_sentiment。

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

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

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

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

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

  • confidence (DOUBLE): 置信分数。

• 示例

  df = pf.from_dict({"comment": ["This feature is amazing!", "Broke after one month"]})
  df = df.llm.ai_sentiment("comment", provider="chat", model="qwen-plus")

信息提取

• 接口名称：ai_extract。

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

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

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  schema	String	是	JSON Schema 字符串，描述要提取的字段，如 '{"name":"STRING", "phone":"STRING"}'。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

• 示例

  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,
                         provider="chat", model="qwen-plus")

文本翻译

• 接口名称：ai_translate。

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

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

• 入参

  参数名称	参数类型	是否必填	描述
  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"。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

• 示例

  df = pf.from_dict({"text": ["Hello World", "How are you?"]})
  df = df.llm.ai_translate("text", source_lang="en", target_lang="zh",
                           provider="chat", model="qwen-plus")

文本摘要

• 接口名称：ai_summarize。

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

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

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  max_length	int	是	摘要的最大字符数，必须大于 0。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

• 示例

  df = pf.from_dict({"article": ["This is a long article with lots of content..."]})
  df = df.llm.ai_summarize("article", max_length=100,
                           provider="chat", model="qwen-plus")
  

数据脱敏

• 接口名称：ai_mask。

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

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

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  entities	List[String]	是	需要脱敏的实体类型列表，如 ["name", "phone"]。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

• 示例

  df = pf.from_dict({"text": ["Please contact John Smith at 555-0123"]})
  df = df.llm.ai_mask("text", entities=["name", "phone"],
                      provider="chat", model="qwen-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_and_null=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	String / List[String]	否	输出列名。agg=False 时需为检索源输出列和 score 分别指定列名；agg=True 时指定一个数组列名。
  config	Dict	否	运行时向量搜索配置。
  ignore_empty_and_null	Boolean	否	是否丢弃没有检索结果的输入行。默认值为 False。

• 返回值

  DataFrame 对象，包含原始列和追加的向量检索结果列。设检索源 search_source 包含 N 列：

  • 当 agg=False 时，每个输入行最多会对应 top_k 条输出行，并追加 N + 1 列：search_source 的 N 列会保留原类型追加到结果中，最后再追加一个 DOUBLE 类型的 score 列。如果指定 output_columns，必须提供 N + 1 个列名，并按顺序作为这些追加列的名称。

  • 当 agg=True 时，每条输入记录对应一条输出记录，并追加一个数组列。数组元素是包含 N + 1 个嵌套字段的行：前 N 个字段来自 search_source 并保留原类型，最后一个字段为 DOUBLE 类型的 score。如果指定 output_columns，必须只提供一个列名，并作为这个追加数组列的名称。

• 示例

  query_df = pf.from_dict({
      "query_id": [1],
      "query_embedding": [[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(),
      },
      columns=["doc_id", "embedding", "title"],
      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"],
  )

文本向量化

• 接口名称：ai_embed。

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

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

• 入参

  参数名称	参数类型	是否必填	描述
  input_col	String / Expression	是	输入文本列名或列表达式。
  dimension	int	否	向量维度。默认值为 1024。
  provider	String	否	Provider 名称。
  model	String	否	模型名称。
  config	Dict	否	运行时配置选项。

• 返回值

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

• 示例

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