如何通过DataflowRunner将DeferredDataFrame保存为feather格式？

核心问题分析 #

你的代码存在三个关键问题导致文件无法生成：

1. Pipeline上下文外的延迟执行问题
   Beam采用延迟执行模型，所有PCollection的操作必须在with beam.Pipeline(...)的上下文内触发。你的df = to_dataframe(embedding_dataframe)和df.to_feather语句写在Pipeline上下文之外，此时Pipeline已经结束，PCollection的计算从未被调度执行，自然不会生成文件。

2. 字段不匹配错误
   BigQuery查询的字段是text_to_embed和identifier，但GetFeatures DoFn中却尝试获取overview和iid字段，导致提取的特征为空字符串，ID为None，后续逻辑即使执行也无法得到有效数据。

3. ModelHandler返回类型不兼容
   EmbeddingModelHandler继承自ModelHandler[str, PredictionResult, SentenceTransformer]，要求run_inference返回Iterable[PredictionResult]，但你直接返回了numpy数组，类型不匹配，会导致后续处理出现隐性错误。

修复后的完整代码 #

import apache_beam as beam
from apache_beam.ml.inference.base import (
    ModelHandler,
    PredictionResult,
    RunInference,
)
from sentence_transformers import SentenceTransformer
import argparse
from apache_beam.options.pipeline_options import PipelineOptions
from typing import Sequence, Optional, Any, Dict, Iterable
from apache_beam.ml.inference.base import KeyedModelHandler
from apache_beam.dataframe.convert import to_dataframe
from apache_beam.dataframe.io import write

ENCODING_MODEL_NAME = "distiluse-base-multilingual-cased-v1"


class EmbeddingModelHandler(
    ModelHandler[str, PredictionResult, SentenceTransformer]
):
    def __init__(self, model_name: str = ENCODING_MODEL_NAME):
        self._model_name = model_name

    def load_model(self) -> SentenceTransformer:
        from sentence_transformers import SentenceTransformer
        return SentenceTransformer(self._model_name)

    def run_inference(
        self,
        batch: Sequence[str],
        model: SentenceTransformer,
        inference_args: Optional[Dict[str, Any]] = None,
    ) -> Iterable[PredictionResult]:
        embedding_matrix = model.encode(
            batch, show_progress_bar=True, normalize_embeddings=True
        )
        # 将numpy数组包装为PredictionResult对象，符合ModelHandler的类型要求
        return [PredictionResult(input_text, embedding) for input_text, embedding in zip(batch, embedding_matrix)]


class GetFeatures(beam.DoFn):
    def process(self, element):
        # 修正字段名，与BigQuery查询结果匹配
        feature = element.get("text_to_embed", "")
        iid = element.get("identifier")
        return [(iid, feature)]


def run(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output",
        dest="output",
        required=True,
        help="Output file to write results to.",
    )
    known_args, pipeline_args = parser.parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)

    with beam.Pipeline(options=pipeline_options) as pipeline:
        embedding_data = (
            pipeline
            | "Read BigQuery"
            >> beam.io.ReadFromBigQuery(
                query="""SELECT text_to_embed,
                                identifier
                                FROM  `gcp-project.gcp-dataset.gcp-table`
                                LIMIT 20
                                """,
                project="gcp-project",
                gcs_location="gs://ml-apache-beam/tmp/",
            )
            | "Get features" >> beam.ParDo(GetFeatures())
            | "Run inference"
            >> RunInference(
                KeyedModelHandler(EmbeddingModelHandler(ENCODING_MODEL_NAME))
            )
            | "Format Rows"
            >> beam.Map(
                lambda element: beam.Row(
                    identifier=int(element[0]), 
                    embedding=element[1].prediction.tolist()
                )
            )
        )

        # 在Pipeline上下文内转换为DataFrame并写入文件
        df = to_dataframe(embedding_data)
        write(df, known_args.output, file_format='feather')


if __name__ == "__main__":
    run()

关键修复点说明 #

1. 调整Pipeline上下文
   将DataFrame转换和写入操作移到with beam.Pipeline(...)内部，确保Beam调度执行所有计算步骤。使用apache_beam.dataframe.io.write替代直接调用df.to_feather，该方法会自动适配Beam的分布式执行模型，支持本地和GCS路径。

2. 修正字段匹配
   将GetFeatures中的overview改为text_to_embed，iid改为identifier，与BigQuery查询的字段一致，确保能正确提取数据。

3. 修复ModelHandler返回类型
   在run_inference中，将每个输入文本和对应的embedding包装为PredictionResult对象，符合ModelHandler的泛型约束，后续处理时通过element[1].prediction获取embedding数组。

4. 修正BigQuery表名格式
   BigQuery表名的正确格式是`project.dataset.table`，而非[project:dataset.table]（旧格式可能引发解析错误）。

验证命令 #

本地运行 #

python beam_pipeline.py --requirements_file requirements.txt --output embedding_output.feather

GCP Dataflow运行 #

python beam_pipeline.py --requirements_file requirements.txt --output "gs://my-bucket/embedding_output.feather" --runner DataflowRunner --project my-gcp-project --region us-central1

内容的提问来源于stack exchange，提问作者Duarte O.C