Python Apache Beam连接Kafka失败：ReadFromKafka返回None求助

我需要开发一个Python Apache Beam Kafka客户端，用于从Kafka读取百万/十亿级的流数据，根据键值过滤后以字典列表的批量形式返回结果。目前卡在Kafka连接环节：使用ReadFromKafka方法时似乎仅返回None，但我用confluent-kafka的测试脚本验证过Kafka连接是正常的，当前使用DirectRunner作为管道运行器。

我的Apache Beam Kafka管道代码 #

import logging
from typing import Any, Dict, Iterator, List, Optional

import apache_beam as beam
from apache_beam.io.kafka import ReadFromKafka
from apache_beam.options.pipeline_options import PipelineOptions

from internal_kafka_client import KafkaConfig


class KafkaProcessData(beam.DoFn):
    # pylint: disable = abstract-method
    def __init__(self, filter_key: Any = None, filter_value: Any = None):
        super().__init__()
        self.filter_key = filter_key
        self.filter_value = filter_value

    def process(self, element: Any, *args, **kwargs) -> Iterator[Any]:
        if self.filter_key and self.filter_value:
            if (
                self.filter_key in element
                and element[self.filter_key] == self.filter_value
            ):
                yield element
        yield element


class BatchElements(beam.DoFn):
    # pylint: disable = abstract-method
    def __init__(self, batch_size: int):
        super().__init__()
        self.batch_size = batch_size
        self.buffer = []

    def process(self, element: Any, *args, **kwargs) -> Iterator[List[Any]]:
        self.buffer.append(element)
        if len(self.buffer) >= self.batch_size:
            yield self.buffer
            self.buffer = []


class KafkaClient:
    # pylint: disable = too-few-public-methods
    """
    pipeline_runner choices:
        - DataflowRunner (streaming mode)
        - FlinkRunner (
                streaming mode, locally, not on cluster, haven't tested with cluster
            )
        - DirectRunner (streaming mode)
    """
    _PIPELINE_TYPES = ["DataflowRunner", "FlinkRunner", "DirectRunner"]

    def __init__(
        self,
        kafka_config: KafkaConfig,
        logger: logging.Logger,
        pipeline_runner: str,
        apache_beam_pipeline_options: Optional[Dict] = None,
    ) -> None:
        if not apache_beam_pipeline_options:
            apache_beam_pipeline_options = {}
        if pipeline_runner not in self._PIPELINE_TYPES:
            raise ValueError(
                f"Given pipeline type {pipeline_runner} is not in "
                f"available pipeline list {self._PIPELINE_TYPES}"
            )
        self.pipeline_runner = pipeline_runner
        self.kafka_config = kafka_config
        self.logger = logger
        self.beam_options = PipelineOptions(
            **apache_beam_pipeline_options, save_main_session=True
        )

    def process(self, filter_key: Any = None, filter_value: Any = None) -> List[Dict]:
        kafka_consumer_config = {
            "bootstrap.servers": self.kafka_config.server_address,
            "auto.offset.reset": self.kafka_config.offset_reset,
        }
        if self.kafka_config.group_id:
            kafka_consumer_config["group.id"] = self.kafka_config.group_id

        with beam.Pipeline(self.pipeline_runner, options=self.beam_options) as pipeline:
            # pylint: disable = unsupported-binary-operation
            kafka_data = pipeline | "Read from Kafka" >> ReadFromKafka(
                topics=self.kafka_config.topics,
                consumer_config=kafka_consumer_config,
            )
            processed_data = kafka_data | "Process Data" >> beam.ParDo(
                KafkaProcessData(filter_key=filter_key, filter_value=filter_value)
            )
            processed_data_batches = processed_data | "Batch Data" >> beam.ParDo(
                BatchElements(batch_size=self.kafka_config.batch_size)
            )
            results = pipeline.run()
            results.wait_until_finish()
            return processed_data_batches[0] if processed_data_batches else []

Kafka连接测试脚本（可正常运行） #

from confluent_kafka import Consumer, KafkaError

bootstrap_servers = '<<kafka_address>>:<<kafka_port>>'
topic = '<<topic_name>>'

consumer_config = {
    'bootstrap.servers': bootstrap_servers,
    'group.id': 'my-group',
    'auto.offset.reset': 'earliest'
}

def test_kafka_connection():
    consumer = Consumer(consumer_config)
    consumer.subscribe([topic])
    try:
        while True:
            msg = consumer.poll(1.0)

            if msg is None:
                continue

            if msg.error():
                if msg.error().code() == KafkaError._PARTITION_EOF:
                    continue
                else:
                    print('Kafka error: {}'.format(msg.error()))
                    break
            print('Received message: {}'.format(msg.value().decode('utf-8')))

    except KeyboardInterrupt:
        consumer.close()
        print('Consumer closed')

if __name__ == '__main__':
    test_kafka_connection()

你的代码存在几个核心问题，导致你误以为ReadFromKafka返回None，逐一修复如下：

1. 错误访问PCollection元素 #

Beam的PCollection是分布式数据集，不能通过索引（如processed_data_batches[0]）直接访问元素，这完全不符合Beam的编程模型，会直接返回None或报错。

本地测试（DirectRunner）时，可通过自定义DoFn将结果收集到线程安全容器中：

from typing import List
import threading

class CollectBatches(beam.DoFn):
    def __init__(self, results_list: List[List[Dict]]):
        self.results_list = results_list
        self.lock = threading.Lock()

    def process(self, batch: List[Dict]):
        with self.lock:
            self.results_list.append(batch)

管道中使用方式：

def process(self, filter_key: Any = None, filter_value: Any = None) -> List[Dict]:
    # ... 其他配置代码不变 ...
    collected_batches = []
    with beam.Pipeline(self.pipeline_runner, options=self.beam_options) as pipeline:
        kafka_data = pipeline | "Read from Kafka" >> ReadFromKafka(...)
        # ... 中间处理步骤 ...
        processed_data_batches | "Collect Results" >> beam.ParDo(CollectBatches(collected_batches))
        results = pipeline.run()
        results.wait_until_finish(duration=30000)  # 设置30秒超时结束流式任务
    # 合并所有批次为一维列表
    return [item for batch in collected_batches for item in batch]

2. 未正确解析Kafka消息 #

ReadFromKafka返回的是**(key, value)元组**，而非直接的字典。你的KafkaProcessData直接把element当字典处理，会导致过滤逻辑完全失效，甚至抛出KeyError。

添加消息解析步骤（假设Kafka消息为JSON格式）：

import json

class ParseKafkaMessage(beam.DoFn):
    def process(self, element):
        _, value = element
        try:
            yield json.loads(value.decode('utf-8'))
        except json.JSONDecodeError:
            # 可根据需求处理解析失败的消息，比如丢弃或记录日志
            pass

管道中插入解析步骤：

kafka_data = pipeline | "Read from Kafka" >> ReadFromKafka(...)
parsed_data = kafka_data | "Parse Messages" >> beam.ParDo(ParseKafkaMessage())
processed_data = parsed_data | "Process Data" >> beam.ParDo(KafkaProcessData(...))

3. KafkaProcessData过滤逻辑错误 #

当前逻辑会无条件返回所有元素：符合过滤条件的元素会被yield两次，不符合的也会被yield一次。修正过滤逻辑：

class KafkaProcessData(beam.DoFn):
    def __init__(self, filter_key: Any = None, filter_value: Any = None):
        super().__init__()
        self.filter_key = filter_key
        self.filter_value = filter_value

    def process(self, element: Dict):
        if not self.filter_key or not self.filter_value:
            yield element
            return
        if self.filter_key in element and element[self.filter_key] == self.filter_value:
            yield element

4. BatchElements批量逻辑不完整 #

当前实现不会在管道结束时输出剩余的缓冲元素，需要重写finish_bundle方法：

class BatchElements(beam.DoFn):
    def __init__(self, batch_size: int):
        super().__init__()
        self.batch_size = batch_size
        self.buffer = []

    def process(self, element: Dict):
        self.buffer.append(element)
        if len(self.buffer) >= self.batch_size:
            yield self.buffer.copy()  # 复制避免后续修改影响已输出批次
            self.buffer.clear()

    def finish_bundle(self):
        if self.buffer:
            yield self.buffer.copy()
            self.buffer.clear()

5. 流式管道运行问题 #

DirectRunner的流式模式默认会持续运行，wait_until_finish不会自动结束。测试场景下可选择：

• 改为批处理模式（ReadFromKafka支持批处理）：

  self.beam_options = PipelineOptions(
      **apache_beam_pipeline_options,
      save_main_session=True,
      streaming=False
  )
  

• 或给wait_until_finish设置超时时间：

  results.wait_until_finish(duration=30000)  # 30秒后强制结束任务
  

内容的提问来源于stack exchange，提问作者muertto