Kafka: Return consumed message results on request

providers/apache/kafka/docs/operators/index.rst

@@ -33,6 +33,10 @@ For parameter definitions take a look at :class:`~airflow.providers.apache.kafka
     If you set the ``commit_cadence`` parameter, ensure that the ``enable.auto.commit`` option in the Kafka connection configuration is explicitly set to ``false``.
     By default, ``enable.auto.commit`` is ``true``, which causes the consumer to auto-commit offsets every 5 seconds, potentially overriding the behavior defined by ``commit_cadence``.
 
+Set ``return_apply_function_results=True`` to return a list containing each non-``None`` value returned by the per-message ``apply_function`` in consume order.
+This option does not apply to ``apply_function_batch``.
+Returned values use normal task return handling and may be stored in XCom, so avoid returning large payloads.
+
 Using the operator
 """"""""""""""""""
 

providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py

@@ -45,6 +45,8 @@ class ConsumeFromTopicOperator(BaseOperator):
     :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
     :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
         defaults to None
+    :param return_apply_function_results: Whether to collect non-None return values from the per-message
+        ``apply_function`` and return them as a list. This option does not apply to ``apply_function_batch``.
     :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
         defaults to "end_of_operator";
         if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
@@ -81,6 +83,7 @@ def __init__(
         apply_function_batch: Callable[..., Any] | str | None = None,
         apply_function_args: Sequence[Any] | None = None,
         apply_function_kwargs: dict[Any, Any] | None = None,
+        return_apply_function_results: bool = False,
         commit_cadence: str = "end_of_operator",
         max_messages: int | None = None,
         max_batch_size: int = 1000,
@@ -94,6 +97,7 @@ def __init__(
         self.apply_function_batch = apply_function_batch
         self.apply_function_args = apply_function_args or ()
         self.apply_function_kwargs = apply_function_kwargs or {}
+        self.return_apply_function_results = return_apply_function_results
         self.kafka_config_id = kafka_config_id
         self.commit_cadence = commit_cadence
         self.max_messages = max_messages
@@ -155,6 +159,7 @@ def execute(self, context) -> Any:
             )
 
         messages_left = self.max_messages or True
+        apply_function_results: list[Any] = []
 
         while self.read_to_end or (
             messages_left > 0
@@ -174,7 +179,9 @@ def execute(self, context) -> Any:
 
             if self.apply_function:
                 for m in msgs:
-                    apply_callable(m)
+                    result = apply_callable(m)
+                    if self.return_apply_function_results and result is not None:
+                        apply_function_results.append(result)
 
             if self.apply_function_batch:
                 apply_callable(msgs)
@@ -189,7 +196,10 @@ def execute(self, context) -> Any:
 
         consumer.close()
 
-        return
+        if self.return_apply_function_results and self.apply_function:
+            return apply_function_results
+
+        return None
 
     def _validate_commit_cadence_on_construct(self):
         """Validate the commit_cadence parameter when the operator is constructed."""

providers/apache/kafka/tests/unit/apache/kafka/operators/test_consume.py

@@ -43,7 +43,7 @@ def _no_op(*args, **kwargs) -> Any:
 
 def create_mock_kafka_consumer(
     message_count: int = 1001, message_content: Any = "test_message", track_consumed_messages: bool = False
-) -> tuple[mock.MagicMock, mock.MagicMock, list[int] | None]:
+) -> tuple[mock.MagicMock, Any, list[int] | None]:
     """
     Creates a mock Kafka consumer with configurable behavior.
 
@@ -81,7 +81,34 @@ def mock_consume(num_messages=0, timeout=-1):
     )
 
     #
-    return mock_consumer, mock_get_consumer, total_consumed_count  # type: ignore[return-value]
+    return mock_consumer, mock_get_consumer, total_consumed_count
+
+
+def create_mock_kafka_consumer_from_messages(
+    messages: list[Any],
+) -> tuple[mock.MagicMock, Any]:
+    mocked_messages = messages.copy()
+
+    def mock_consume(num_messages=0, timeout=-1):
+        nonlocal mocked_messages
+        if num_messages < 0:
+            raise Exception("Number of messages needs to be positive")
+
+        msg_count = min(num_messages, len(mocked_messages))
+        returned_messages = mocked_messages[:msg_count]
+        mocked_messages = mocked_messages[msg_count:]
+
+        return returned_messages
+
+    mock_consumer = mock.MagicMock()
+    mock_consumer.consume = mock_consume
+
+    mock_get_consumer = mock.patch(
+        "airflow.providers.apache.kafka.hooks.consume.KafkaConsumerHook.get_consumer",
+        return_value=mock_consumer,
+    )
+
+    return mock_consumer, mock_get_consumer
 
 
 class TestConsumeFromTopic:
@@ -279,3 +306,61 @@ def test_commit_cadence_behavior(self, commit_cadence, max_messages, expected_co
 
             # Verify consumer was closed
             mock_consumer.close.assert_called_once()
+
+    def test_apply_function_results_return_none_by_default(self):
+        mock_consumer, mock_get_consumer = create_mock_kafka_consumer_from_messages(["one", "two"])
+
+        with mock_get_consumer:
+            operator = ConsumeFromTopicOperator(
+                kafka_config_id="kafka_d",
+                topics=["test"],
+                task_id="test",
+                poll_timeout=0.0001,
+                max_messages=2,
+                max_batch_size=2,
+                apply_function=lambda message: f"processed-{message}",
+            )
+
+            assert operator.execute(context={}) is None
+            mock_consumer.close.assert_called_once()
+
+    def test_return_apply_function_results_filters_none_and_preserves_order(self):
+        mock_consumer, mock_get_consumer = create_mock_kafka_consumer_from_messages(
+            ["first", "skip", "second"]
+        )
+
+        def apply_function(message):
+            return None if message == "skip" else f"processed-{message}"
+
+        with mock_get_consumer:
+            operator = ConsumeFromTopicOperator(
+                kafka_config_id="kafka_d",
+                topics=["test"],
+                task_id="test",
+                poll_timeout=0.0001,
+                max_messages=3,
+                max_batch_size=2,
+                apply_function=apply_function,
+                return_apply_function_results=True,
+            )
+
+            assert operator.execute(context={}) == ["processed-first", "processed-second"]
+            mock_consumer.close.assert_called_once()
+
+    def test_return_apply_function_results_does_not_change_batch_return_behavior(self):
+        mock_consumer, mock_get_consumer = create_mock_kafka_consumer_from_messages(["one", "two"])
+
+        with mock_get_consumer:
+            operator = ConsumeFromTopicOperator(
+                kafka_config_id="kafka_d",
+                topics=["test"],
+                task_id="test",
+                poll_timeout=0.0001,
+                max_messages=2,
+                max_batch_size=2,
+                apply_function_batch=lambda messages: [f"processed-{message}" for message in messages],
+                return_apply_function_results=True,
+            )
+
+            assert operator.execute(context={}) is None
+            mock_consumer.close.assert_called_once()