[data] Adding Kafka datasink.

[justinrmiller]

Description

This PR introduces a Kafka Datasink for Ray Data, allowing users to write Ray Datasets directly to Apache Kafka topics. The implementation leverages Confluent's Kafka python library and supports distributed writes with configurable serialization, key extraction, and performance tuning (e.g., periodic flushing to manage memory).

---

Key Changes

ray.data._internal.datasource.kafka_datasink.py

• New KafkaDatasink Class: Implements the Datasink interface to handle parallel writes across Ray workers.

• Serialization Support: Provides built-in support for json, string, and bytes for both message keys and values.

• Smart Row Conversion: Includes a logic handler to convert various Ray row types (Dict, ArrowRow, PandasRow, NamedTuple) into serializable formats.

• Memory Management:

• • Implements a _FLUSH_INTERVAL (set to 10,000) to periodically flush the producer and wait for acknowledgments, preventing memory exhaustion from too many un-flushed futures.
• • Handles BufferError by flushing and retrying the send operation.

• Error Handling: Captures and reports the number of failed messages, ensuring that the first encountered exception is chained to the final RuntimeError for easier debugging.

ray.data.dataset.py

• write_kafka() Method: Added as a top-level convenience method on the Dataset class.
• API Exposure: Decorated with @ConsumptionAPI, allowing users to call ds.write_kafka(...) with standard Ray remote arguments and concurrency controls.

ray.data.tests.datasource.test_kafka.py

• Unit Tests: Validates initialization, row-to-dict conversion, and serialization logic in isolation.
• Integration Tests: Comprehensive test suite (requires a Kafka broker) covering:
• Basic writes and multi-block writes.
• Key extraction and custom serializers.
• Producer configurations and delivery callbacks.
• Edge cases like empty datasets, null values, and invalid connection strings.

---

Basic Usage

import ray

ds = ray.data.range(100)
ds.write_kafka(
    topic="my-topic", 
    bootstrap_servers="localhost:9092",
    key_field="id",
    key_serializer="string",
    value_serializer="json"
)

Related issues

Closes #58725

[gemini-code-assist]

Code Review

This pull request introduces a KafkaDatasink for Ray Data, which is a valuable addition. The implementation is generally well-structured. I've identified a few areas for improvement to enhance robustness and maintainability. My feedback includes suggestions to refactor duplicated code, correct potentially buggy logic in object-to-dictionary conversion, add parameter validation, and fix an incorrect docstring example. Addressing these points will strengthen the new datasink implementation.

[owenowenisme]

Thanks for the contribution!

[justinrmiller]

@owenowenisme Do you think we should make the flush interval a configurable parameter? I went with 10k as I thought that would be a reasonable amount (guidance is generally to keep message sizes under 1.5 KB which is 15 MB max (10k * 1.5 KB).

[gemini-code-assist]

Code Review

This pull request introduces a Kafka Datasink for Ray Data, which is a great addition. The implementation is well-structured, with good error handling for producer buffer errors and delivery failures. The periodic polling to manage memory is also a thoughtful touch. The accompanying test suite is comprehensive, covering unit tests, integration tests, and various edge cases. I have one suggestion to improve code maintainability by refactoring duplicated serialization logic.

_{Note: Security Review is unavailable for this PR.}

[iamjustinhsu]

nice work!! I'm not too too familiar with kafka, so i had some questions. But overall looks good

[iamjustinhsu]

can you add a comment/docstring on what None signifies, and if/how it affects message distribution?

[iamjustinhsu]

if there are a lot of backlog messages that need to be delivered, will this line try to free up all space, and if so, do u know how long that will take? I'm wondering if this should be less blocking in the worst, like 5 seconds, but not sure

[iamjustinhsu]

lgtm!

[owenowenisme]

Overall LGTM. One thing to note: we should document that the current implementation provides best-effort
delivery only — partial writes can occur if a task fails midway, and duplicates are possible on system-error
task retries (e.g., node failures).

For future improvements, we could explore:

1. Per-task Kafka transactions — to guarantee atomicity within each write task (no partial messages from failed
   tasks, no duplicates on retry)
2. Checkpoint-based delivery tracking — to enable resumable writes and stronger end-to-end guarantees

[cursor]

Test never polls Kafka when expected count is zero

Low Severity

The consume_messages helper has a while len(messages) < expected_count loop condition. When expected_count=0, 0 < 0 is False, so the loop body never executes — the function returns an empty list immediately without ever polling Kafka. In test_write_kafka_empty_dataset, the assertion assert len(messages) == 0 is therefore trivially true regardless of whether messages were actually written, making the verification meaningless.

Additional Locations (1)

• python/ray/data/tests/datasource/test_kafka.py#L1088-L1091

 

[cursor]

Dead exception handler catches nothing in try block

Low Severity

The except KafkaException handler in write() is unreachable. Within the try block, producer.produce() is asynchronous and reports errors via the on_delivery callback (not by raising). producer.poll() and producer.flush() also don't raise KafkaException. All other operations are Ray/Python-internal and can't produce a KafkaException. This creates a false sense of error handling — a real KafkaException from the constructor on line 236 occurs before the try block and would propagate uncaught.

 

[justinrmiller]

@owenowenisme Are we good to go on this one? Thanks!

[cursor]

Stale consumer assignment when topic metadata missing

Medium Severity

The consume_messages helper only calls consumer.assign() when topic metadata is found (if topic_meta and topic_meta.partitions). Since the kafka_consumer fixture is session-scoped and shared across all integration tests, if topic metadata is not found for a given test, assign() is never called, and the consumer retains the partition assignment from a previous test. This means it would poll from the wrong topic, returning stale messages and causing misleading test results instead of a clean failure.

 

[owenowenisme]

LGTM

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 4 total unresolved issues (including 3 from previous reviews).

[cursor]

Missing producer flush on non-KafkaException error paths

Medium Severity

The try/except KafkaException block in write() only catches KafkaException. If _produce_with_retry raises RuntimeError (persistent BufferError) or serialization raises TypeError (e.g., numpy types from Pandas blocks not being JSON-serializable), producer.flush() is never called. Messages already successfully enqueued in the producer's internal buffer before the error are silently dropped without any delivery attempt. A try/finally wrapping the producer.flush() call would ensure buffered messages get a chance to be delivered even on unexpected error paths.

 

[owenowenisme]

Not a bug. Flushing on error paths would be worse — the task has failed and Ray Data may retry it from scratch. Flushing a partial batch means the retry produces duplicates for the already-delivered prefix. Skipping flush minimizes the partial write window, which is the right behavior. The class docstring already documents the at-most-once-per-attempt delivery semantics.