如何将ConsumeFromTopicOperator数据推送至XCom并传递给PythonOperator？

核心问题 #

• ConsumeFromTopicOperator消费的Kafka数据无法通过XCom传递给下游PythonOperator，导致process_message中message变量无输出
• AwaitFromKafkaTopicOperator仅捕获单条Topic记录，无法满足批量消费需求

问题原因及修复 #

1. XCom无输出的原因 #

你的consumer_function没有返回任何数据，而ConsumeFromTopicOperator的XCom推送完全依赖apply_function_batch的返回值。当前函数仅做日志打印后return（无返回内容），导致XCom中没有可拉取的数据；同时你在process_message中指定拉取key='message'，但Operator默认推送的XCom key是return_value，两者不匹配。

修复步骤：

修改consumer_function，返回处理后的消息列表：

def consumer_function(messages, prefix=None):
    consumer_logger.info("starting consumer function")
    consumer_logger.info(f"Messages count: {len(messages)}")
    processed_messages = []
    for message in messages:
        decoded_msg = message.value().decode('utf-8')
        consumer_logger.info(
            f"{prefix} {message.topic()} @ {message.offset()}; Message {decoded_msg}"
        )
        # 整理结构化数据以便下游处理
        processed_messages.append({
            "topic": message.topic(),
            "offset": message.offset(),
            "value": decoded_msg
        })
    # 返回数据用于XCom推送
    return processed_messages

修改process_message，拉取默认的return_value：

def process_message(ti, **context):
    # 拉取consume_records任务的返回值，默认key为return_value
    messages = ti.xcom_pull(task_ids='consume_records', key='return_value')
    print("Received messages from Kafka:", messages)
    
    # 示例处理逻辑
    processed_data = [msg['value'] for msg in messages]
    ti.xcom_push(key='processed_data', value=processed_data)

2. AwaitFromKafkaTopicOperator仅捕获单条记录的问题 #

AwaitFromKafkaTopicOperator的设计定位就是等待并获取单条Kafka记录，如果需要批量消费，直接使用ConsumeFromTopicOperator即可——你当前配置的max_messages=10（最多消费10条）、max_batch_size=2（每次批量处理2条）参数是合理的，修复XCom问题后就能实现批量数据传递。

完整修复后代码示例 #

import functools
from airflow import DAG
from airflow.operators.python import PythonOperator
from astronomer.providers.kafka.operators.consume import ConsumeFromTopicOperator
import logging

consumer_logger = logging.getLogger(__name__)

# 补充你的默认参数
default_args = {
    'owner': 'airflow',
    'start_date': ...
}

dag = DAG(dag_id='Stream_Employee_Data_Loader',
          default_args=default_args,
          catchup=False,
          schedule_interval=None)

def process_message(ti, **context):
    messages = ti.xcom_pull(task_ids='consume_records', key='return_value')
    print("Received messages:", messages)
    
    processed_data = [msg['value'] for msg in messages]
    ti.xcom_push(key='processed_data', value=processed_data)

retrieve_data_task = PythonOperator(
         task_id='retrieve_data',
         python_callable=process_message,
         provide_context=True,
         do_xcom_push=True,
         dag=dag
     )

def consumer_function(messages, prefix=None):
    consumer_logger.info("starting consumer function")
    consumer_logger.info(f"Messages count: {len(messages)}")
    processed_messages = []
    for message in messages:
        decoded_msg = message.value().decode('utf-8')
        consumer_logger.info(
            f"{prefix} {message.topic()} @ {message.offset()}; Message {decoded_msg}"
        )
        processed_messages.append({
            "topic": message.topic(),
            "offset": message.offset(),
            "value": decoded_msg
        })
    return processed_messages

consume_task = ConsumeFromTopicOperator(
        task_id='consume_records',
        topics=[kafka_topic],
        apply_function_batch=functools.partial(
            consumer_function, prefix="consumed:::"
        ),
        consumer_config={
            "bootstrap.servers": "kafka:9092",
            "group.id": "kafka_connection",
            "enable.auto.commit": False,
            "auto.offset.reset": "beginning",
        },
        commit_cadence="end_of_batch",
        max_messages=10,
        max_batch_size=2,
        do_xcom_push=True,
        dag=dag
    )

consume_task >> retrieve_data_task 

关键注意点 #

• ConsumeFromTopicOperator的apply_function_batch必须返回数据，才能触发XCom推送
• XCom默认拉取的key是return_value，若需自定义可在Operator中设置xcom_push_key参数
• 批量消费场景优先使用ConsumeFromTopicOperator，AwaitFromKafkaTopicOperator仅适用于单条记录等待场景

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