Python DataFrame API reference

Use in with_column or with_columns

A scalar UDF takes one or more columns and returns a single column. Use Python type hints for automatic return type inference, or specify the type explicitly with return_dtype. You can also set the operator's parallelism with the concurrency parameter.

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

Use in map

The map function takes an entire row and returns a new row. The input row can be accessed by column name in the function. The return type can be explicitly declared with DataType.struct or automatically inferred from a TypedDict return type annotation.

Explicitly declare with return_dtype

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

Automatically infer type with TypedDict

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)

Use in with_column or with_columns

Vectorized UDFs receive and return data in batches, reducing the overhead of row-by-row calls between Python and the Flink runtime. They support both Pandas (pandas.Series) and Arrow (pyarrow.Array) formats. Use batch_size to control the batch size.

Pandas format

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

Arrow format

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

Use in map_batches

map_batches performs vectorized processing on entire rows. When batch_format="pandas", the input and output are dict[str, pandas.Series]. When batch_format="arrow", the input and output are dict[str, pyarrow.Array].

Pandas format

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

Arrow format

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

Text classification

• API name: ai_classify.

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

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

• Parameters

  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, e.g., ["positive", "negative", "neutral"].
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

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

• Example

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

Sentiment analysis

• API name: ai_sentiment.

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

• Description: Performs sentiment analysis on the input text.

• Parameters

  Parameter	Type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

  A DataFrame with the following columns added:

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

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

  • confidence (DOUBLE): The confidence score.

• Example

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

Information extraction

• API name: ai_extract.

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

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

• Parameters

  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 describing the fields to extract, e.g., '{"name":"STRING", "phone":"STRING"}'.
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

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

• Example

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

Text translation

• API name: ai_translate.

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

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

• Parameters

  Parameter	Type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  source_lang	String	Yes	The source language code, such as "zh", "en", or "auto" (for automatic detection). Supported values: auto, zh, en, ja, ko, fr, de, es, ru, ar, pt.
  target_lang	String	Yes	The target language code. This value cannot be "auto".
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

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

• Example

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

Text summarization

• API name: ai_summarize.

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

• Description: Summarizes text to a specified maximum length.

• Parameters

  Parameter	Type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  max_length	int	Yes	The maximum number of characters in the summary. This value must be greater than 0.
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

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

• Example

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

Data masking

• API name: ai_mask.

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

• Description: Masks sensitive information in text.

• Parameters

  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, e.g., ["name", "phone"].
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

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

• Example

  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

• 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_and_null=False
  )
  

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

• Parameters

  Parameter	Type	Required	Description
  search_source	DataFrame	Yes	The vector search source DataFrame, for example, a vector collection loaded via 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	The number of most similar results to return per input record.
  agg	Boolean	No	Whether to aggregate the top-K results into an array column. Defaults to False.
  output_columns	String / List[String]	No	The output column names. When agg=False, provide names for each search-source output column and the score column. When agg=True, provide a single array column name.
  config	Dict	No	Runtime vector search configuration.
  ignore_empty_and_null	Boolean	No	Whether to drop input rows that have no search results. Defaults to False.

• Return value

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

  • When agg=False, each input row corresponds to at most top_k output rows. N + 1 columns are appended: the N columns from search_source retain their original types, plus a DOUBLE-typed score column. If output_columns is specified, provide N + 1 column names in order.

  • When agg=True, each input row corresponds to one output row with an appended array column. Each array element is a row with N + 1 nested fields: the first N fields come from search_source and retain their original types, and the last field is a DOUBLE-typed score. If output_columns is specified, provide a single column name for the array column.

• Example

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

Text embedding

• API name: ai_embed.

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

• Description: Generates a vector embedding for text.

• Parameters

  Parameter	Type	Required	Description
  input_col	String / Expression	Yes	The input text column name or a column expression.
  dimension	int	No	The dimension of the embedding vector. Defaults to 1024.
  provider	String	No	The provider name.
  model	String	No	The model name.
  config	Dict	No	Runtime configuration options.

• Returns

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

• Example

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