[Data] Add Kafka as datasource by owenowenisme · Pull Request #58592 · ray-project/ray

owenowenisme · 2025-11-13T09:50:04Z

Description

Added Kafka as data source, we can now use read_kafka for bounded data source

Related issues

Closes #58653

Additional information

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

bveeramani · 2025-11-15T03:50:18Z

python/ray/data/read_api.py

+    start_offset: Optional[Union[int, str]] = "earliest",
+    end_offset: Optional[Union[int, str]] = "latest",


Are earlier and latest the only valid string values? Will we need to add more values later?

Changed from str into Literal["earliest"] & Literal["latest"]

bveeramani · 2025-11-15T03:51:12Z

python/ray/data/read_api.py

+    trigger: Literal["once"] = "once",
+    start_offset: Optional[Union[int, str]] = "earliest",
+    end_offset: Optional[Union[int, str]] = "latest",
+    authentication: Optional[Dict[str, Any]] = None,


What happens if this is None? If it's implicitly required, I think it'd be clearer if this was a required argument.

Also, in the interest of being explicit over implicit, I think we should use a typed dataclass rather than a dict (what are the valid fields and values? You'd only know by looking at the docstring)

Made authentication a dataclass

What happens if this is None? If it's implicitly required, I think it'd be clearer if this was a required argument.

This is not required.

python/ray/data/read_api.py

bveeramani · 2025-11-15T03:55:24Z

python/ray/data/read_api.py

+            By default, the number of output blocks is dynamically decided based on
+            input data size and available resources. You shouldn't manually set this
+            value in most cases.
+        timeout_ms: Timeout in milliseconds to poll to until reaching end_offset (default 10000ms/10s).


What happens if we don't expose this?

I think user should set their timeout?

python/ray/data/_internal/datasource/kafka_datasource.py

bveeramani · 2025-11-15T04:16:10Z

python/ray/data/_internal/datasource/kafka_datasource.py

+
+            # Create metadata for this task
+            metadata = BlockMetadata(
+                num_rows=max_records_per_task,


What happens if we don't set this? This could be inaccurate, right?

I removed max_records_per_task since we current run a task per partition.
Also set the num_rows to None.

bveeramani · 2025-11-15T04:20:12Z

python/ray/data/_internal/datasource/kafka_datasource.py

+                in the remote task without requiring 'self' to be serialized.
+                """
+
+                def kafka_read_fn() -> Iterable[Block]:


The implementation of this function is pretty long. Wonder if there's a way we can make it shorter (and easier to read) by abstracting some of the logic as functions.

If you feel like it'd hurt readability though, okay to keep as-is.

bveeramani · 2025-11-15T04:20:36Z

python/deplocks/ray_img/ray_img_py310.lock

-    --hash=sha256:31f23644fe2602f88ff55e1f5c79ba497e01224ee7737937930c448e4d0e24dc \
-    --hash=sha256:a8a6399716257f45be6a007360200409fce5cda2661e3dec71d23dc15f6189ab
-    # via uvicorn
+python-dotenv==1.2.1 \


For my own understanding, what's the deal with all the lock files? Is this what REEf recommended?

I asked REEf team how do I resolve deps error and they want me to run this command
bazel run //ci/raydepsets:raydepsets -- build --all-configs

Pasting the message

I added some test deps and ran bash ci/ci.sh compile_pip_dependencies to update requirement_compiled.txt
And still getting this error, do you have any idea how we should resolve this?
https://buildkite.com/ray-project/premerge/builds/53872/steps/canvas?sid=019a84cf-060a-43b7-b6ef-39fe56484670

Is this what REEf recommended?

yes.

you can split the python-docenv version upgrade PR into a separate one.

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

gemini-code-assist

Code Review

This pull request introduces a new Kafka datasource for Ray Data, allowing bounded reads from Kafka topics. The implementation is well-structured, using kafka-python and providing a clear API via ray.data.read_kafka. The addition of KafkaAuthConfig makes authentication flexible. The integration tests are comprehensive, covering various scenarios including different offset configurations, multiple partitions/topics, and timeout behavior, which is excellent. I've found a few areas for improvement, mainly around documentation clarity and a couple of small code refinements. Overall, this is a great contribution.

python/ray/data/_internal/datasource/kafka_datasource.py

python/ray/data/read_api.py

python/ray/data/tests/test_kafka.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com>

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

cursor · 2025-11-16T14:31:21Z

python/ray/data/read_api.py

+    trigger: Literal["once"] = "once",
+    start_offset: Union[int, Literal["earliest"]] = "earliest",
+    end_offset: Union[int, Literal["latest"]] = "latest",
+    kafka_auth_config: KafkaAuthConfig = KafkaAuthConfig(),


Bug: Mutable Defaults Break Configuration Isolation

Mutable default argument KafkaAuthConfig() creates a shared instance across all calls to read_kafka. Since KafkaAuthConfig is a non-frozen dataclass, any mutation to this shared instance would affect all subsequent function calls that don't provide kafka_auth_config. The default should be None instead, with instantiation inside the function if needed.

aslonnie

fwiw, @elliot-barn is giving a talk about depset next tue.

aslonnie · 2025-11-16T14:32:55Z

python/deplocks/ray_img/ray_img_py310.lock

-    --hash=sha256:31f23644fe2602f88ff55e1f5c79ba497e01224ee7737937930c448e4d0e24dc \
-    --hash=sha256:a8a6399716257f45be6a007360200409fce5cda2661e3dec71d23dc15f6189ab
-    # via uvicorn
+python-dotenv==1.2.1 \


Is this what REEf recommended?

yes.

you can split the python-docenv version upgrade PR into a separate one.

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

aslonnie · 2025-11-16T14:34:54Z

release/ray_release/byod/ray_base_extra_testdeps_cuda_py3.10.lock

-    --hash=sha256:31f23644fe2602f88ff55e1f5c79ba497e01224ee7737937930c448e4d0e24dc \
-    --hash=sha256:a8a6399716257f45be6a007360200409fce5cda2661e3dec71d23dc15f6189ab
-    # via uvicorn
+python-dotenv==1.2.1 \


@elliot-barn could you change codeowner so that these lock files are not owned by ray-ci ? we do not need to review every single dependency change for release tests.

aslonnie

approval for release test lock file changes.

python/ray/data/tests/test_kafka.py

…-datasource

bveeramani · 2025-11-18T19:11:24Z

python/ray/data/read_api.py

+            By default, the number of output blocks is dynamically decided based on
+            input data size and available resources. You shouldn't manually set this
+            value in most cases.
+        timeout_ms: Timeout in milliseconds to poll to until reaching end_offset (default 10000ms).


Can you elaborate on this? For example, is this the timeout per task, or for the read as a whole? What happens when the timeout is reached? Does the job fail?

bveeramani · 2025-11-18T19:16:45Z

python/ray/data/_internal/datasource/kafka_datasource.py

+                            if elapsed_time >= timeout_seconds:
+                                break


If a read task times out and performs a partial read, I don't think the user would know that from the logs.

Maybe let's add a logging.warning or logging.info so the users aren't surprised (and it's easier to debug)?

Added logging.warning

python/ray/data/tests/test_kafka.py

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

python/ray/data/_internal/datasource/kafka_datasource.py

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

python/ray/data/tests/test_kafka.py

python/ray/data/_internal/datasource/kafka_datasource.py

## Description Added Kafka as data source, we can now use `read_kafka` for bounded data source ## Related issues Closes ray-project#58653 ## Additional information --------- Signed-off-by: You-Cheng Lin <mses010108@gmail.com> Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Aydin Abiar <aydin@anyscale.com>

## Description Added Kafka as data source, we can now use `read_kafka` for bounded data source ## Related issues Closes ray-project#58653 ## Additional information --------- Signed-off-by: You-Cheng Lin <mses010108@gmail.com> Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Pavitra Bhalla <pavitra@rayai.com>

## Description Added Kafka as data source, we can now use `read_kafka` for bounded data source ## Related issues Closes ray-project#58653 ## Additional information --------- Signed-off-by: You-Cheng Lin <mses010108@gmail.com> Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

## Description Added Kafka as data source, we can now use `read_kafka` for bounded data source ## Related issues Closes ray-project#58653 ## Additional information --------- Signed-off-by: You-Cheng Lin <mses010108@gmail.com> Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: YK <1811651+ykdojo@users.noreply.github.com>

## Description Added Kafka as data source, we can now use `read_kafka` for bounded data source ## Related issues Closes ray-project#58653 ## Additional information --------- Signed-off-by: You-Cheng Lin <mses010108@gmail.com> Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

## Description Added Kafka as data source, we can now use `read_kafka` for bounded data source ## Related issues Closes ray-project#58653 ## Additional information --------- Signed-off-by: You-Cheng Lin <mses010108@gmail.com> Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

owenowenisme added 2 commits November 13, 2025 17:49

pick-from-stephen's work

f1eb981

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

update

3221246

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

owenowenisme added the go add ONLY when ready to merge, run all tests label Nov 14, 2025

owenowenisme and others added 11 commits November 14, 2025 21:52

update requirement_compiled.txt

90312c3

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

testcontainer

be7bf8d

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

update test

fad95b2

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

add timeout test

bcc6727

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

remove python-dotenv

3fc991a

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

Merge branch 'master' into data/add-kafka-as-datasource

67040fa

add doc entry

a934dae

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

refine read_kafka

bfa465a

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

update req compile txt

0e9d832

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

update lock

63825e7

Signed-off-by: You-Cheng Lin (Owen) <mses010108@gmail.com>

make value bytes

e08b3a9

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

bveeramani reviewed Nov 15, 2025

View reviewed changes

make key from msg also bytes

6b9ffc7

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

richardliaw mentioned this pull request Nov 15, 2025

Ray Data Q4 Roadmap + Wishlist #58665

Open

owenowenisme and others added 13 commits November 16, 2025 14:06

make offsets only accept int, 'earlist' and 'latest'

f9acc46

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

handle error in _resolve_offsets and make offset non-optional

f526946

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

Merge branch 'master' into data/add-kafka-as-datasource

4120726

nake kafka auth config a type

f28a2e5

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

message header key is string

4d6aaa7

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

make kafkaAurhConfig dataclass

71c96d7

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

remove parallelism param

17421b7

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

remove max_records_per_task and make internal variable underscored

e822674

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

gcm set _target_max_block_size in constructor

37000dd

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

refine yield logic

eb9cdb4

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

move schema to the front

673b4ed

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

make trigger err msg raw string

981a68d

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

remove redundant check

bb56032

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

gemini-code-assist bot reviewed Nov 16, 2025

View reviewed changes

owenowenisme and others added 4 commits November 16, 2025 22:23

Update python/ray/data/read_api.py

134104a

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com>

Update python/ray/data/_internal/datasource/kafka_datasource.py

6e60f8a

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com>

Update python/ray/data/_internal/datasource/kafka_datasource.py

956f063

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com>

update comments

b3639c0

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

cursor bot reviewed Nov 16, 2025

View reviewed changes

aslonnie reviewed Nov 16, 2025

View reviewed changes

update optional KafkaAuthConfig

a5f04b9

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

aslonnie reviewed Nov 16, 2025

View reviewed changes

aslonnie approved these changes Nov 16, 2025

View reviewed changes

cursor bot reviewed Nov 16, 2025

View reviewed changes

python/ray/data/tests/test_kafka.py Show resolved Hide resolved

python/ray/data/tests/test_kafka.py Show resolved Hide resolved

owenowenisme mentioned this pull request Nov 16, 2025

[CI] Add kafka-python&testcontainers[kafka] to deps #58680

Merged

owenowenisme requested a review from bveeramani November 16, 2025 17:24

ray-gardener bot added the data Ray Data-related issues label Nov 16, 2025

Merge remote-tracking branch 'upstream/master' into data/add-kafka-as…

3e628c4

…-datasource

bveeramani approved these changes Nov 18, 2025

View reviewed changes

Merge branch 'master' into data/add-kafka-as-datasource

566f497

cursor bot reviewed Nov 19, 2025

View reviewed changes

python/ray/data/tests/test_kafka.py Show resolved Hide resolved

add log update description

f375dab

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

cursor bot reviewed Nov 19, 2025

View reviewed changes

python/ray/data/_internal/datasource/kafka_datasource.py Show resolved Hide resolved

update log

6d6b71b

Signed-off-by: You-Cheng Lin <mses010108@gmail.com>

bveeramani enabled auto-merge (squash) November 19, 2025 01:46

cursor bot reviewed Nov 19, 2025

View reviewed changes

python/ray/data/tests/test_kafka.py Show resolved Hide resolved

python/ray/data/_internal/datasource/kafka_datasource.py Show resolved Hide resolved

bveeramani merged commit 946e35c into ray-project:master Nov 19, 2025
6 of 7 checks passed

		start_offset: Optional[Union[int, str]] = "earliest",
		end_offset: Optional[Union[int, str]] = "latest",

Conversation

owenowenisme commented Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related issues

Additional information

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

owenowenisme Nov 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

owenowenisme Nov 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor bot Nov 16, 2025

Choose a reason for hiding this comment

Bug: Mutable Defaults Break Configuration Isolation

Uh oh!

aslonnie left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aslonnie left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

owenowenisme commented Nov 13, 2025 •

edited

Loading

owenowenisme Nov 16, 2025 •

edited

Loading

owenowenisme Nov 16, 2025 •

edited

Loading