自定义离线算子（UDF）

接入方式

将函数或可调用类传给 Ray Data 的 map_batches，即可在数据集流水线中分布式执行自定义逻辑。

本文示例使用 Ray Data 公开接口。基础示例需要 ray[data] 和 numpy，串联内置算子时还需要 mmops。请确保 Driver 与 Worker 节点使用相同的依赖环境。mmops 的获取和安装方法，请参见mmops SDK 使用指南。mmops 的环境要求和调用约定，请参见离线算子概述。

示例算子：文本标准化

本文以文本标准化算子为例：将连续空白合并为单个空格、去掉首尾空白，并可选转为小写。算子保留原字段，每条输入对应一条输出，并新增清洗后的文本和字符数。

编写并执行函数 UDF

当 batch_format="numpy" 时，每个批次是“列名 → 数组”的字典。函数需返回同样的列式结构，且各列长度必须一致；本例的输出行数与输入相同。函数创建新数组，不修改输入数组。

import os
import re
import numpy as np
import ray


def normalize_text_batch(batch, *, lowercase=False):
    if "text" not in batch or "source_id" not in batch:
        raise ValueError("需要 source_id 和 text 两列")
    if "clean_text" in batch or "char_count" in batch:
        raise ValueError("输出列 clean_text / char_count 已存在")
    if any(not isinstance(text, str) for text in batch["text"]):
        raise TypeError("text 必须是字符串；空字符串允许")

    cleaned = [re.sub(r"\s+", " ", text).strip()
               for text in batch["text"]]
    if lowercase:
        cleaned = [text.lower() for text in cleaned]
    # 保留输入列；新建数组，不修改输入缓冲区。
    out = dict(batch)
    out["clean_text"] = np.asarray(cleaned, dtype=str)
    out["char_count"] = np.asarray([len(t) for t in cleaned], dtype=np.int64)
    return out


# 设置 RAY_ADDRESS=auto 连接现有集群；未设置时启动本地 Ray。
ray.init(address=os.environ.get("RAY_ADDRESS"))
source = ray.data.from_items([
    {"source_id": "doc-001", "text": "  产品\n  使用说明  "},
    {"source_id": "doc-002", "text": "Ray   Data"},
])
result = source.map_batches(
    normalize_text_batch, batch_format="numpy", batch_size=32,
    fn_kwargs={"lowercase": False}, num_cpus=1,
)
# 小样本读取触发执行；大规模任务使用 write_json 等输出接口。
print(result.take_all())

示例输出保留 source_id 和原 text 列，并新增 clean_text 和 char_count 列。部分结果列如下：

使用类 UDF 复用状态

类的构造方法在执行实例内初始化，__call__ 方法处理后续批次。将模型权重或客户端放在构造方法中创建，可以避免每个批次重复加载。

import os
import re
import numpy as np
import ray


class NormalizeTextUDF:
    def __init__(self, lowercase=False):
        self.lowercase = lowercase
        # 模型或客户端可在此初始化，由当前执行实例复用。

    def __call__(self, batch):
        if "text" not in batch or "source_id" not in batch:
            raise ValueError("需要 source_id 和 text 两列")
        if "clean_text" in batch or "char_count" in batch:
            raise ValueError("输出列 clean_text / char_count 已存在")
        if any(not isinstance(text, str) for text in batch["text"]):
            raise TypeError("text 必须是字符串；空字符串允许")

        cleaned = [re.sub(r"\s+", " ", text).strip()
                   for text in batch["text"]]
        if self.lowercase:
            cleaned = [text.lower() for text in cleaned]
        out = dict(batch)
        out["clean_text"] = np.asarray(cleaned, dtype=str)
        out["char_count"] = np.asarray([len(t) for t in cleaned], dtype=np.int64)
        return out


# 设置 RAY_ADDRESS=auto 连接现有集群；未设置时启动本地 Ray。
ray.init(address=os.environ.get("RAY_ADDRESS"))
source = ray.data.from_items([
    {"source_id": "doc-001", "text": "  产品\n  使用说明  "},
    {"source_id": "doc-002", "text": "Ray   Data"},
])
result = source.map_batches(
    NormalizeTextUDF,
    fn_constructor_kwargs={"lowercase": False},
    compute=ray.data.ActorPoolStrategy(size=2),
    batch_format="numpy", batch_size=32, num_cpus=1,
)
print(result.take_all())

ActorPoolStrategy(size=2) 表示使用两个执行实例，每个实例分别持有自己的状态。实际并行度还受输入分块和可用资源的限制，因此不能把实例内的计数当作全局计数。使用 GPU 模型时，请为执行实例配置 num_gpus，并在构造方法中加载对应模型。

与内置算子串联

以下示例的处理链为：原始文本 → 自定义文本标准化 UDF → split_text → 展开片段 → 输出数据集。自定义逻辑与内置算子通过数据列衔接，并保留来源标识。

import re
from uuid import uuid4
import numpy as np
import ray
from mmops.ray import enable_accessor

ray.init()
enable_accessor()

def normalize_text_batch(batch):
    if any(not isinstance(t, str) for t in batch["text"]):
        raise TypeError("text 必须是字符串")
    return {
        "source_id": batch["source_id"],
        "clean_text": np.asarray([
            re.sub(r"\s+", " ", t).strip() for t in batch["text"]
        ], dtype=str),
    }

source = ray.data.from_items([
    {"source_id": "doc-001", "text": "  产品  使用说明  "},
    {"source_id": "doc-002", "text": "Ray   Data"},
])
cleaned = source.map_batches(
    normalize_text_batch, batch_format="numpy", batch_size=32, num_cpus=1,
)
chunks = cleaned.mmops.split_text(
    text_column="clean_text", keep_columns=["source_id"],
    config={"strategy": "length", "max_chars": 800, "overlap_chars": 80},
    batch_size=1, concurrency=4, num_cpus=1,
)
result = chunks.flat_map(lambda row: [
    {"source_id": row["source_id"],
     "chunk_index": chunk["chunk_index"], "text": chunk["content"]}
    for chunk in row["chunks"] if chunk["content"].strip()
])
result.select_columns(["source_id", "chunk_index", "text"]).write_parquet(
    "/shared/output/custom-udf/" + uuid4().hex
)

结果为包含 source_id、chunk_index 和 text 列的片段数据集。后续可以继续接入向量化等算子，具体输入列与参数以对应接口文档为准。

import json

rows = result.take(2)
print(json.dumps(rows, ensure_ascii=False, indent=2))

[
  {"source_id": "doc-001", "chunk_index": 0, "text": "产品 使用说明"},
  {"source_id": "doc-002", "chunk_index": 0, "text": "Ray Data"}
]

参数与使用约定

相关文档

• Ray Data map_batches

• Ray Data 数据转换与有状态 UDF