diff --git a/lib/python3.12/site-packages/ray/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..388c1abd1a3d3bc13f6ae375f3505796ddb7867a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/_version.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/_version.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..814faed91431870099fcfc9623e28f030995b39f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/_version.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/actor.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/actor.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e5dde12c47cef2322a6a0ede5f78049f5ade455f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/actor.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/client_builder.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/client_builder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7d98a3a4f3449da8e371d7a1b31d38ba48b4ab56
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/client_builder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/cluster_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/cluster_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c17832854258fc19b583c742d13eb11b2f92de1b
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/cluster_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/cross_language.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/cross_language.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d5f26d208cd326aa4d8ad4b86233c0d5fbf4bd62
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/cross_language.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/exceptions.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..73fe04a52195f3592d266149cd2373830e51fd84
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/exceptions.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/job_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/job_config.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..309f4bc686ef27c4e25b65697e4eaae3ef3fc884
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/job_config.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/remote_function.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/remote_function.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..946cf3f17f01fa7badd007aa67ef548df5264002
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/remote_function.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/runtime_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/runtime_context.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..07e21ca4af7077af4d1613cad8806df23d124826
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/runtime_context.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/setup-dev.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/setup-dev.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e3413f679b01175b52107f5061026a4a0cf78c2a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/setup-dev.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/ray/__pycache__/types.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..29448ee4298a14e7d7c596a19a70b7360e441082
Binary files /dev/null and b/lib/python3.12/site-packages/ray/__pycache__/types.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..736fb3275da31bf93cf4041cd8659d4d1a892da7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/aggregate.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/aggregate.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..03527b8bcbeb748e4923803cfcfb43022e65bd80
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/aggregate.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/block.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/block.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..460007aeb41e887e3ce44df277630bcfd29a671f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/block.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/collate_fn.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/collate_fn.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..aebafc88261c70161dea55fefb2ba944de0e6d68
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/collate_fn.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/context.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/context.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1cd11729ad6cc7e5c3e9e26000a48e92b79d7019
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/context.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/datatype.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/datatype.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d3849be34ef12c90dff3c39a7a3d64faa94ed676
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/datatype.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/exceptions.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..126179ca223fe6c4d3347e7cffc8304cf5d8e0f2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/exceptions.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/expressions.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/expressions.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3a65772039c1421e0a51d784166df0d2eee1d641
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/expressions.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/grouped_data.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/grouped_data.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f7469cb2eb160c3753f74a2815c0df37899c4714
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/grouped_data.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/iterator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/iterator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..576282b6754416b0e2ee4e0e9982d290de7a5de3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/iterator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/llm.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/llm.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2de242e232675448a72cf630cd622057f19a1f0a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/llm.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/preprocessor.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/preprocessor.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e2e41a9378177f48e611b7b815c4bed12e9179ef
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/preprocessor.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/random_access_dataset.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/random_access_dataset.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5b4050c0c533216202c7b60ddb2b351591564949
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/random_access_dataset.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/__pycache__/stats.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/__pycache__/stats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..afe720319bb46af73b0da2d0a6aea1118c36bf4c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/__pycache__/stats.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..94233d7e2491d60dca4fd88f742fd3bbeba60dd7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/aggregate.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/aggregate.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..09e80e6dd4486a11bc2b52cb86ba99646253f855
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/aggregate.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/arrow_block.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/arrow_block.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..526741a321910f7db0e8740d273f41ee00b86ff0
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/arrow_block.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/batcher.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/batcher.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..019fbcf91953b83f9d4c3c036a47db130b01ad11
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/batcher.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/block_builder.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/block_builder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..38b34c831487c33de7a62fd9596c87a15cc42469
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/block_builder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/block_list.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/block_list.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d03d4ea84a207df88a1274c8bea7e69406bd513e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/block_list.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/collections.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/collections.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9fa7bfbab49fc9f4547937159fd64d648c6dcc50
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/collections.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/compute.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/compute.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..16501aca0fc4c83d6c49b0342bfa423c8eae935e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/compute.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/delegating_block_builder.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/delegating_block_builder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7ff68b8ba503e2000216a1c500434f22da2291c9
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/delegating_block_builder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/equalize.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/equalize.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ba5f9ef6e7d690b37d3341a3c69101c5b84003fd
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/equalize.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/logging.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/logging.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..34540e0470291ac1764b3ce03e05049b86682667
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/logging.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/memory_tracing.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/memory_tracing.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1ccbf7e3d75d7bb117bb1720b4ee3272de7ed1b8
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/memory_tracing.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/metadata_exporter.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/metadata_exporter.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f9fa66ed53769cfb69534783a433953b2186e217
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/metadata_exporter.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/numpy_support.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/numpy_support.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..91ce1d0153329137b56b0fe49eef20eb001d3e57
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/numpy_support.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/operator_event_exporter.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/operator_event_exporter.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a98eee531adab09f8a04b4b63d86f25615869e3f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/operator_event_exporter.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/output_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/output_buffer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..517a6db572f879bf6e3102ff2bbbd94cdae0ae13
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/output_buffer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/pandas_block.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/pandas_block.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d3d732ff34cdd22d3853182bc63969d1bc72f0f5
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/pandas_block.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/plan.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/plan.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..68b4d975e512b4f2d558a4489fdcc83603db9b91
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/plan.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/progress_bar.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/progress_bar.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b7864de36e99ef49fd8e053138dac05b38d02875
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/progress_bar.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/remote_fn.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/remote_fn.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1e8b238e053f2512569fdc5c98d01e403046fad3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/remote_fn.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/row.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/row.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1e88ecc28a77eb8cfa284ace4536f2f9ca2d5117
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/row.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/savemode.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/savemode.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..862a9a98901df3993c41f62afc48cbc4ff108ef2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/savemode.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/size_estimator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/size_estimator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..05374bb0dbe5c422a5da9e0c69b02724bb7dcdd1
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/size_estimator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/split.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/split.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ccdc8b2a38896cf20548a9a6a4e35d13f17b1a09
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/split.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/stats.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/stats.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ce17c9a4a7079dae130be4da55a9646ce4c619ee
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/stats.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/streaming_repartition.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/streaming_repartition.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4ce766c976dfb87bb1bccbde48c01f2babb2c03c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/streaming_repartition.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/table_block.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/table_block.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0110b500ef2b4a24bd15c395a83cebbedf3ec791
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/table_block.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/torch_iterable_dataset.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/torch_iterable_dataset.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2cb6ba27efa87a45c1689ac070e97ddf846d5c88
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/torch_iterable_dataset.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/util.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0e6ee58067f241d6d4d9a822ea6e155f3795d754
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/__pycache__/util.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d29cbc9e78c5ecd284f4167a8acb16a4d3bc1c2
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__init__.py
@@ -0,0 +1,30 @@
+from typing import TYPE_CHECKING
+
+from .autoscaling_actor_pool import ActorPoolScalingRequest, AutoscalingActorPool
+from .base_actor_autoscaler import ActorAutoscaler
+from .default_actor_autoscaler import DefaultActorAutoscaler
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+    from ray.data.context import AutoscalingConfig
+
+
+def create_actor_autoscaler(
+    topology: "Topology",
+    resource_manager: "ResourceManager",
+    config: "AutoscalingConfig",
+) -> ActorAutoscaler:
+    return DefaultActorAutoscaler(
+        topology,
+        resource_manager,
+        config=config,
+    )
+
+
+__all__ = [
+    "ActorAutoscaler",
+    "ActorPoolScalingRequest",
+    "AutoscalingActorPool",
+    "create_actor_autoscaler",
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6f0943329f57ac2a626737f509b1aa8026ec9119
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/autoscaling_actor_pool.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/autoscaling_actor_pool.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f179b545aa099b1af12f531a0f507cd5452ff69b
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/autoscaling_actor_pool.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/base_actor_autoscaler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/base_actor_autoscaler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c36162add98a2e4e5c1688bce5bf24383434a4c9
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/base_actor_autoscaler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/default_actor_autoscaler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/default_actor_autoscaler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3930f17883b60daa784fa475dc4c9dfcf06ecdfe
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/__pycache__/default_actor_autoscaler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/autoscaling_actor_pool.py b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/autoscaling_actor_pool.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6145a4c175f6612ff133c22559955951ea7a064
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/autoscaling_actor_pool.py
@@ -0,0 +1,114 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional
+
+from ray.data._internal.execution.interfaces.execution_options import ExecutionResources
+from ray.util.annotations import DeveloperAPI
+
+
+@dataclass(frozen=True)
+class ActorPoolScalingRequest:
+
+    delta: int
+    force: bool = field(default=False)
+    reason: Optional[str] = field(default=None)
+
+    @classmethod
+    def no_op(cls, *, reason: Optional[str] = None) -> "ActorPoolScalingRequest":
+        return ActorPoolScalingRequest(delta=0, reason=reason)
+
+    @classmethod
+    def upscale(cls, *, delta: int, reason: Optional[str] = None):
+        assert delta > 0
+        return ActorPoolScalingRequest(delta=delta, reason=reason)
+
+    @classmethod
+    def downscale(
+        cls, *, delta: int, force: bool = False, reason: Optional[str] = None
+    ):
+        assert delta < 0, "For scale down delta is expected to be negative!"
+        return ActorPoolScalingRequest(delta=delta, force=force, reason=reason)
+
+
+@DeveloperAPI
+class AutoscalingActorPool(ABC):
+    """Abstract interface of an autoscaling actor pool.
+
+    A `PhysicalOperator` can manage one or more `AutoscalingActorPool`s.
+    `Autoscaler` is responsible for deciding autoscaling of these actor
+    pools.
+    """
+
+    @abstractmethod
+    def min_size(self) -> int:
+        """Min size of the actor pool."""
+        ...
+
+    @abstractmethod
+    def max_size(self) -> int:
+        """Max size of the actor pool."""
+        ...
+
+    @abstractmethod
+    def current_size(self) -> int:
+        """Current size of the actor pool."""
+        ...
+
+    @abstractmethod
+    def num_running_actors(self) -> int:
+        """Number of running actors."""
+        ...
+
+    @abstractmethod
+    def num_active_actors(self) -> int:
+        """Number of actors with at least one active task."""
+        ...
+
+    @abstractmethod
+    def num_pending_actors(self) -> int:
+        """Number of actors pending creation."""
+        ...
+
+    @abstractmethod
+    def max_tasks_in_flight_per_actor(self) -> int:
+        """Max number of in-flight tasks per actor."""
+        ...
+
+    @abstractmethod
+    def max_actor_concurrency(self) -> int:
+        """Returns max number of tasks single actor could run concurrently."""
+        ...
+
+    @abstractmethod
+    def num_tasks_in_flight(self) -> int:
+        """Number of current in-flight tasks (ie total nubmer of tasks that have been
+        submitted to the actor pool)."""
+        ...
+
+    def num_free_task_slots(self) -> int:
+        """Number of free slots to run tasks.
+
+        This doesn't include task slots for pending actors.
+        """
+        return (
+            self.max_tasks_in_flight_per_actor() * self.num_running_actors()
+            - self.num_tasks_in_flight()
+        )
+
+    @abstractmethod
+    def scale(self, req: ActorPoolScalingRequest):
+        """Applies autoscaling action"""
+        ...
+
+    @abstractmethod
+    def per_actor_resource_usage(self) -> ExecutionResources:
+        """Per actor resource usage."""
+        ...
+
+    @abstractmethod
+    def get_pool_util(self) -> float:
+        """Calculate the utilization of the given actor pool."""
+        ...
+
+    def max_concurrent_tasks(self) -> int:
+        return self.max_actor_concurrency() * self.num_running_actors()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/base_actor_autoscaler.py b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/base_actor_autoscaler.py
new file mode 100644
index 0000000000000000000000000000000000000000..aebdb89bb4314c06858631f5319aa0fba3d99391
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/base_actor_autoscaler.py
@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+@DeveloperAPI
+class ActorAutoscaler(ABC):
+    """Abstract interface for Ray Data actor autoscaler."""
+
+    def __init__(
+        self,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        self._topology = topology
+        self._resource_manager = resource_manager
+
+    @abstractmethod
+    def try_trigger_scaling(self):
+        """Try trigger autoscaling.
+
+        This method will be called each time when StreamingExecutor makes
+        a scheduling decision. A subclass should override this method to
+        handle the autoscaling of `AutoscalingActorPool`s.
+        """
+        ...
diff --git a/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/default_actor_autoscaler.py b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/default_actor_autoscaler.py
new file mode 100644
index 0000000000000000000000000000000000000000..e84c063994c3e416ce2e0e6a475b9032bf9521d7
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/actor_autoscaler/default_actor_autoscaler.py
@@ -0,0 +1,194 @@
+import logging
+import math
+from typing import TYPE_CHECKING, Optional
+
+from .autoscaling_actor_pool import ActorPoolScalingRequest, AutoscalingActorPool
+from .base_actor_autoscaler import ActorAutoscaler
+from ray.data._internal.execution.interfaces.execution_options import ExecutionResources
+from ray.data.context import WARN_PREFIX, AutoscalingConfig
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import PhysicalOperator
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import OpState, Topology
+
+logger = logging.getLogger(__name__)
+
+
+class DefaultActorAutoscaler(ActorAutoscaler):
+    def __init__(
+        self,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+        *,
+        config: AutoscalingConfig,
+    ):
+        super().__init__(topology, resource_manager)
+
+        self._actor_pool_scaling_up_threshold = (
+            config.actor_pool_util_upscaling_threshold
+        )
+        self._actor_pool_scaling_down_threshold = (
+            config.actor_pool_util_downscaling_threshold
+        )
+
+        self._validate_autoscaling_config()
+
+    def try_trigger_scaling(self):
+        for op, state in self._topology.items():
+            actor_pools = op.get_autoscaling_actor_pools()
+            for actor_pool in actor_pools:
+                # Trigger auto-scaling
+                actor_pool.scale(
+                    self._derive_target_scaling_config(actor_pool, op, state)
+                )
+
+    def _derive_target_scaling_config(
+        self,
+        actor_pool: "AutoscalingActorPool",
+        op: "PhysicalOperator",
+        op_state: "OpState",
+    ) -> ActorPoolScalingRequest:
+        # If all inputs have been consumed, short-circuit
+        if op.completed() or (
+            op._inputs_complete and op_state.total_enqueued_input_blocks() == 0
+        ):
+            return ActorPoolScalingRequest.downscale(
+                delta=-1, force=True, reason="consumed all inputs"
+            )
+
+        if actor_pool.current_size() < actor_pool.min_size():
+            # Scale up, if the actor pool is below min size.
+            return ActorPoolScalingRequest.upscale(
+                delta=actor_pool.min_size() - actor_pool.current_size(),
+                reason="pool below min size",
+            )
+        elif actor_pool.current_size() > actor_pool.max_size():
+            # Do not scale up, if the actor pool is already at max size.
+            return ActorPoolScalingRequest.downscale(
+                # NOTE: For scale down delta has to be negative
+                delta=-(actor_pool.current_size() - actor_pool.max_size()),
+                reason="pool exceeding max size",
+            )
+
+        # Determine whether to scale up based on the actor pool utilization.
+        util = actor_pool.get_pool_util()
+        if util >= self._actor_pool_scaling_up_threshold:
+            # Do not scale up if either
+            #   - Previous scale up has not finished yet
+            #   - Actor Pool is at max size already
+            #   - Op is throttled (ie exceeding allocated resource quota)
+            #   - Actor Pool has sufficient amount of slots available to handle
+            #   pending tasks
+            if actor_pool.num_pending_actors() > 0:
+                return ActorPoolScalingRequest.no_op(reason="pending actors")
+            elif actor_pool.current_size() >= actor_pool.max_size():
+                return ActorPoolScalingRequest.no_op(reason="reached max size")
+            if not op_state._scheduling_status.under_resource_limits:
+                return ActorPoolScalingRequest.no_op(
+                    reason="operator exceeding resource quota"
+                )
+            budget = self._resource_manager.get_budget(op)
+            if _get_max_scale_up(actor_pool, budget) == 0:
+                return ActorPoolScalingRequest.no_op(reason="exceeded resource limits")
+
+            return ActorPoolScalingRequest.upscale(
+                delta=1,
+                reason=(
+                    f"utilization of {util} >= "
+                    f"{self._actor_pool_scaling_up_threshold}"
+                ),
+            )
+        elif util <= self._actor_pool_scaling_down_threshold:
+            if actor_pool.current_size() <= actor_pool.min_size():
+                return ActorPoolScalingRequest.no_op(reason="reached min size")
+
+            return ActorPoolScalingRequest.downscale(
+                delta=-1,
+                reason=(
+                    f"utilization of {util} <= "
+                    f"{self._actor_pool_scaling_down_threshold}"
+                ),
+            )
+        else:
+            return ActorPoolScalingRequest.no_op(
+                reason=(
+                    f"utilization of {util} w/in limits "
+                    f"[{self._actor_pool_scaling_down_threshold}, "
+                    f"{self._actor_pool_scaling_up_threshold}]"
+                )
+            )
+
+    def _validate_autoscaling_config(self):
+        for op, state in self._topology.items():
+            for actor_pool in op.get_autoscaling_actor_pools():
+                self._validate_actor_pool_autoscaling_config(actor_pool, op)
+
+    def _validate_actor_pool_autoscaling_config(
+        self,
+        actor_pool: "AutoscalingActorPool",
+        op: "PhysicalOperator",
+    ) -> None:
+        """Validate autoscaling configuration.
+
+        Args:
+            actor_pool: Actor pool to validate configuration thereof.
+            op: ``PhysicalOperator`` using target actor pool.
+        """
+        max_tasks_in_flight_per_actor = actor_pool.max_tasks_in_flight_per_actor()
+        max_concurrency = actor_pool.max_actor_concurrency()
+
+        if (
+            max_tasks_in_flight_per_actor / max_concurrency
+            < self._actor_pool_scaling_up_threshold
+        ):
+            logger.warning(
+                f"{WARN_PREFIX} Actor Pool configuration of the {op} will not allow it to scale up: "
+                f"configured utilization threshold ({self._actor_pool_scaling_up_threshold * 100}%) "
+                f"couldn't be reached with configured max_concurrency={max_concurrency} "
+                f"and max_tasks_in_flight_per_actor={max_tasks_in_flight_per_actor} "
+                f"(max utilization will be max_tasks_in_flight_per_actor / max_concurrency = {(max_tasks_in_flight_per_actor / max_concurrency) * 100:g}%)"
+            )
+
+
+def _get_max_scale_up(
+    actor_pool: AutoscalingActorPool,
+    budget: Optional[ExecutionResources],
+) -> Optional[int]:
+    """Get the maximum number of actors that can be scaled up.
+
+    Args:
+        actor_pool: The actor pool to scale up.
+        budget: The budget to scale up.
+
+    Returns:
+        The maximum number of actors that can be scaled up, or `None` if you can
+        scale up infinitely.
+    """
+    if budget is None:
+        return None
+
+    assert budget.cpu >= 0 and budget.gpu >= 0
+
+    num_cpus_per_actor = actor_pool.per_actor_resource_usage().cpu
+    num_gpus_per_actor = actor_pool.per_actor_resource_usage().gpu
+    assert num_cpus_per_actor >= 0 and num_gpus_per_actor >= 0
+
+    max_cpu_scale_up: float = float("inf")
+    if num_cpus_per_actor > 0 and not math.isinf(budget.cpu):
+        max_cpu_scale_up = budget.cpu // num_cpus_per_actor
+
+    max_gpu_scale_up: float = float("inf")
+    if num_gpus_per_actor > 0 and not math.isinf(budget.gpu):
+        max_gpu_scale_up = budget.gpu // num_gpus_per_actor
+
+    max_scale_up = min(max_cpu_scale_up, max_gpu_scale_up)
+    if math.isinf(max_scale_up):
+        return None
+    else:
+        assert not math.isnan(max_scale_up), (
+            budget,
+            num_cpus_per_actor,
+            num_gpus_per_actor,
+        )
+        return int(max_scale_up)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..29c77f6219805d9ca6f73284a3e2a0f9dfeed269
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__init__.py
@@ -0,0 +1,3 @@
+from ray.data._internal.block_batching.block_batching import batch_blocks
+
+__all__ = ["batch_blocks"]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d7d4b9a4a5c7af4185060b3512c1982732f814cf
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/block_batching.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/block_batching.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9417effc2a24ff8e4e7815cc2d50c2e8e7d9ed64
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/block_batching.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/interfaces.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/interfaces.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2015ef70e6c747ded926c906dc992c78b4039a62
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/interfaces.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/iter_batches.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/iter_batches.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..eb0cf3f405a036c1e85cd8007317837513d6a9b4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/iter_batches.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/util.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0dd32dca8da3e1666f501bbdb9b69cd569191544
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/block_batching/__pycache__/util.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/block_batching.py b/lib/python3.12/site-packages/ray/data/_internal/block_batching/block_batching.py
new file mode 100644
index 0000000000000000000000000000000000000000..b90dd7809949703827a827491545e9425aa8630a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/block_batching/block_batching.py
@@ -0,0 +1,60 @@
+from contextlib import nullcontext
+from typing import Callable, Iterator, Optional, TypeVar
+
+from ray.data._internal.block_batching.util import (
+    blocks_to_batches,
+    collate,
+    extract_data_from_batch,
+    format_batches,
+)
+from ray.data._internal.stats import DatasetStats
+from ray.data.block import Block, DataBatch
+
+T = TypeVar("T")
+
+
+def batch_blocks(
+    blocks: Iterator[Block],
+    *,
+    stats: Optional[DatasetStats] = None,
+    batch_size: Optional[int] = None,
+    batch_format: str = "default",
+    drop_last: bool = False,
+    collate_fn: Optional[Callable[[DataBatch], DataBatch]] = None,
+    shuffle_buffer_min_size: Optional[int] = None,
+    shuffle_seed: Optional[int] = None,
+    ensure_copy: bool = False,
+) -> Iterator[DataBatch]:
+    """Create formatted batches of data from 1 or more blocks.
+
+    This function takes in an iterator of already fetched blocks. Consequently, this
+    function doesn't support block prefetching.
+    """
+
+    def _iterator_fn(base_iterator: Iterator[Block]) -> Iterator[DataBatch]:
+        batch_iter = format_batches(
+            blocks_to_batches(
+                block_iter=base_iterator,
+                stats=stats,
+                batch_size=batch_size,
+                drop_last=drop_last,
+                shuffle_buffer_min_size=shuffle_buffer_min_size,
+                shuffle_seed=shuffle_seed,
+                ensure_copy=ensure_copy,
+            ),
+            batch_format=batch_format,
+            stats=stats,
+        )
+
+        if collate_fn is not None:
+            batch_iter = collate(batch_iter, collate_fn=collate_fn, stats=stats)
+
+        batch_iter = extract_data_from_batch(batch_iter)
+        yield from batch_iter
+
+    batch_iter = _iterator_fn(blocks)
+
+    for formatted_batch in batch_iter:
+        user_timer = stats.iter_user_s.timer() if stats else nullcontext()
+        with user_timer:
+            yield formatted_batch
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/interfaces.py b/lib/python3.12/site-packages/ray/data/_internal/block_batching/interfaces.py
new file mode 100644
index 0000000000000000000000000000000000000000..4f0bed6b3dd42b6a1722006c1fc2681c1d534d8c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/block_batching/interfaces.py
@@ -0,0 +1,55 @@
+import abc
+from dataclasses import dataclass
+from typing import Any, List
+
+from ray.data.block import Block, DataBatch
+from ray.types import ObjectRef
+
+
+@dataclass
+class BatchMetadata:
+    """Metadata associated with a batch.
+
+    Attributes:
+        batch_idx: The global index of this batch so that downstream operations can
+            maintain ordering.
+    """
+
+    batch_idx: int
+
+
+@dataclass
+class Batch:
+    """A batch of data.
+
+    Attributes:
+        metadata: Metadata associated with this batch.
+        data: The batch of data.
+    """
+
+    metadata: BatchMetadata
+    data: DataBatch
+
+
+class CollatedBatch(Batch):
+    """A batch of collated data.
+
+    Attributes:
+        data: The batch of data which is the output of a user provided collate_fn
+            Therefore, the type of this data can be Any.
+    """
+
+    data: Any
+
+
+class BlockPrefetcher(metaclass=abc.ABCMeta):
+    """Interface for prefetching blocks."""
+
+    @abc.abstractmethod
+    def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
+        """Prefetch the provided blocks to this node."""
+        pass
+
+    def stop(self):
+        """Stop prefetching and release resources."""
+        pass
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/iter_batches.py b/lib/python3.12/site-packages/ray/data/_internal/block_batching/iter_batches.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0f19f96f983ed981fa34ce3ae7eaf8693ac957a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/block_batching/iter_batches.py
@@ -0,0 +1,432 @@
+import collections
+import time
+from contextlib import contextmanager, nullcontext
+from typing import Any, Callable, Dict, Iterator, Optional
+
+import ray
+from ray.data._internal.block_batching.interfaces import Batch, BlockPrefetcher
+from ray.data._internal.block_batching.util import (
+    ActorBlockPrefetcher,
+    WaitBlockPrefetcher,
+    blocks_to_batches,
+    collate,
+    finalize_batches,
+    format_batches,
+    resolve_block_refs,
+)
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.memory_tracing import trace_deallocation
+from ray.data._internal.stats import DatasetStats, _StatsManager
+from ray.data._internal.util import make_async_gen
+from ray.data.block import Block, DataBatch
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+
+class BatchIterator:
+    """Defines an iterator pipeline to convert a stream of block object references
+    into a stream of formatted batches ready to be consumed by the user.
+
+    This takes a block iterator and creates batch_size batches, slicing,
+    unioning, shuffling, prefetching, and formatting blocks as needed.
+
+    This involves both pipeline parallelism (e.g. prefetching)
+    and data parallelism (e.g. threadpool operations):
+
+    If prefetch_batches=2, these are all the batches in flight:
+
+    [User thread] trains on Batch 0
+    - [Fetch thread] Batch 1 finalization + move to output queue
+            - [Worker thread 1] Batch 2 formatting + collating
+            - [Worker thread 2] Batch 3 formatting + collating
+            - [Raylet] Batches 4 + 5 fetched to local object store memory
+
+    At any point in time there are prefetch_batches+1 batches in local heap memory.
+    And the next set of prefetch_batches in local object store memory.
+
+    The actual steps are as follows:
+
+    In a single async thread, do the following:
+        1. Trigger Ray local prefetching of `prefetch_batches` worth of block object
+            references.
+        2. Resolve (i.e. call `ray.get()`) on the block references.
+        3. Perform the necessary batch slicing to construct full batches, possibly
+            shuffling if necessary.
+        4. Then, in a threadpool consisting of `prefetch_batches` threads:
+            a. Format the batches to the provided batch format.
+            b. Apply the collate function.
+        5. Finalize each of the collated batches
+        6. Fetch outputs from the threadpool, maintaining order of the batches.
+
+    Args:
+        ref_bundles: An iterator over RefBundles.
+        stats: DatasetStats object to record timing and other statistics.
+        dataset_tag: The tag of the dataset to record timing and other statistics.
+        clear_block_after_read: Whether to clear the block from object store
+            manually (i.e. without waiting for Python's automatic GC) after it
+            is read. Doing so will reclaim memory faster and hence reduce the
+            memory footprint. However, the caller has to ensure the safety, i.e.
+            the block will never be accessed again.
+        batch_size: Record batch size, or None to let the system pick.
+        batch_format: The format in which to return each batch.
+            Specify "default" to use the current block format (promoting
+            Arrow to pandas automatically), "pandas" to
+            select ``pandas.DataFrame`` or "pyarrow" to select
+            ``pyarrow.Table``, or None to use entire blocks
+            as batches. Default is "default".
+        drop_last: Whether to drop the last batch if it's incomplete.
+        collate_fn: A function to apply to each data batch before returning it.
+        finalize_fn: A function to apply to each data batch after it has been collated.
+            This function is not run in a threadpool so it can be used for
+            memory-intensive operations such as GPU preloading.
+        shuffle_buffer_min_size: If non-None, the data will be randomly shuffled using a
+            local in-memory shuffle buffer, and this value will serve as the minimum
+            number of rows that must be in the local in-memory shuffle buffer in order
+            to yield a batch.
+        shuffle_seed: The seed to use for the local random shuffle.
+        ensure_copy: Whether batches are always copied from the underlying base
+            blocks (not zero-copy views).
+        prefetch_batches: The number of batches to fetch ahead of the current batch to
+            process. If set to greater than 0, a separate thread will be used to fetch
+            the specified amount of formatted batches from blocks. This improves
+            performance for non-CPU bound UDFs, allowing batch fetching compute and
+            formatting to be overlapped with the UDF. Defaults to 1.
+    """
+
+    UPDATE_METRICS_INTERVAL_S: float = 5.0
+
+    def __init__(
+        self,
+        ref_bundles: Iterator[RefBundle],
+        *,
+        stats: Optional[DatasetStats] = None,
+        dataset_tag: Optional[str] = None,
+        clear_block_after_read: bool = False,
+        batch_size: Optional[int] = None,
+        batch_format: Optional[str] = "default",
+        drop_last: bool = False,
+        collate_fn: Optional[Callable[[DataBatch], Any]] = None,
+        finalize_fn: Optional[Callable[[Any], Any]] = None,
+        shuffle_buffer_min_size: Optional[int] = None,
+        shuffle_seed: Optional[int] = None,
+        ensure_copy: bool = False,
+        prefetch_batches: int = 1,
+    ):
+        self._ref_bundles = ref_bundles
+        self._stats = stats
+        self._dataset_tag = dataset_tag
+        self._batch_size = batch_size
+        self._batch_format = batch_format
+        self._drop_last = drop_last
+        self._collate_fn = collate_fn
+        self._finalize_fn = finalize_fn
+        self._shuffle_buffer_min_size = shuffle_buffer_min_size
+        self._shuffle_seed = shuffle_seed
+        self._ensure_copy = ensure_copy
+        self._prefetch_batches = prefetch_batches
+        self._eager_free = (
+            clear_block_after_read and DataContext.get_current().eager_free
+        )
+
+        actor_prefetcher_enabled = (
+            prefetch_batches > 0
+            and DataContext.get_current().actor_prefetcher_enabled
+            and not ray.util.client.ray.is_connected()
+        )
+        self._prefetcher = (
+            ActorBlockPrefetcher()
+            if actor_prefetcher_enabled
+            else WaitBlockPrefetcher()
+        )
+        self._yielded_first_batch = False
+
+        # This stores the last time we updated the metrics.
+        # This allows us to update metrics on some interval,
+        # by comparing it with the current timestamp.
+        self._metrics_last_updated: float = 0.0
+
+    def _prefetch_blocks(
+        self, ref_bundles: Iterator[RefBundle]
+    ) -> Iterator[ObjectRef[Block]]:
+        return prefetch_batches_locally(
+            ref_bundles=ref_bundles,
+            prefetcher=self._prefetcher,
+            num_batches_to_prefetch=self._prefetch_batches,
+            batch_size=self._batch_size,
+            eager_free=self._eager_free,
+            stats=self._stats,
+        )
+
+    def _resolve_block_refs(
+        self, block_refs: Iterator[ObjectRef[Block]]
+    ) -> Iterator[Block]:
+        return resolve_block_refs(block_ref_iter=block_refs, stats=self._stats)
+
+    def _blocks_to_batches(self, blocks: Iterator[Block]) -> Iterator[Batch]:
+        return blocks_to_batches(
+            block_iter=blocks,
+            stats=self._stats,
+            batch_size=self._batch_size,
+            drop_last=self._drop_last,
+            shuffle_buffer_min_size=self._shuffle_buffer_min_size,
+            shuffle_seed=self._shuffle_seed,
+            ensure_copy=self._ensure_copy,
+        )
+
+    def _format_batches(self, batches: Iterator[Batch]) -> Iterator[Batch]:
+        return _format_in_threadpool(
+            batch_iter=batches,
+            stats=self._stats,
+            batch_format=self._batch_format,
+            collate_fn=self._collate_fn,
+            num_threadpool_workers=self._prefetch_batches,
+        )
+
+    def _finalize_batches(
+        self,
+        batch_iter: Iterator[Batch],
+    ) -> Iterator[Batch]:
+        if self._finalize_fn is None:
+            return batch_iter
+
+        return finalize_batches(
+            batch_iter, finalize_fn=self._finalize_fn, stats=self._stats
+        )
+
+    def _restore_original_batch_order(
+        self, batches: Iterator[Batch]
+    ) -> Iterator[Batch]:
+        return restore_original_order(batches)
+
+    def _pipeline(self, ref_bundles: Iterator[RefBundle]) -> Iterator[Batch]:
+        # Step 1: Prefetch logical batches locally.
+        block_iter = self._prefetch_blocks(ref_bundles)
+
+        # Step 2: Resolve the blocks.
+        block_iter = self._resolve_block_refs(block_iter)
+
+        # Step 3: Batch and shuffle the resolved blocks.
+        batch_iter = self._blocks_to_batches(block_iter)
+
+        # Step 4: Format and collate the batches in a threadpool.
+        batch_iter = self._format_batches(batch_iter)
+
+        # Step 5: Finalize the batches (e.g., move to GPU).
+        batch_iter = self._finalize_batches(batch_iter)
+
+        # Step 6: Restore the original order of the batches, as the prior
+        # threadpool operations may have reordered the batches non-deterministically.
+        batch_iter = self._restore_original_batch_order(batch_iter)
+
+        yield from batch_iter
+
+    def _iter_batches(self) -> Iterator[DataBatch]:
+        async_batch_iter = make_async_gen(
+            self._ref_bundles,
+            fn=self._pipeline,
+            num_workers=1,
+            preserve_ordering=False,
+        )
+
+        self.before_epoch_start()
+
+        while True:
+            with self.get_next_batch_context():
+                try:
+                    batch = next(async_batch_iter)
+                except StopIteration:
+                    break
+            with self.yield_batch_context(batch):
+                yield batch.data
+
+        self.after_epoch_end()
+
+    def __iter__(self) -> Iterator[DataBatch]:
+        return self._iter_batches()
+
+    def before_epoch_start(self):
+        self._yielded_first_batch = False
+
+    def after_epoch_end(self):
+        if self._stats is None:
+            return
+
+        _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag)
+
+    @contextmanager
+    def get_next_batch_context(self):
+        try:
+            if self._stats:
+                # Always track total blocked time
+                total_timer = self._stats.iter_total_blocked_s.timer()
+                # Also track the time until the first batch is ready
+                first_batch_ready_timer = (
+                    self._stats.iter_time_to_first_batch_s.timer()
+                    if not self._yielded_first_batch
+                    else nullcontext()
+                )
+                with total_timer, first_batch_ready_timer:
+                    yield
+            else:
+                yield
+        finally:
+            self._yielded_first_batch = True
+
+    @contextmanager
+    def yield_batch_context(self, batch: Batch):
+        with self._stats.iter_user_s.timer() if self._stats else nullcontext():
+            yield
+
+        if self._stats is None:
+            return
+        now = time.time()
+        if (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S:
+            _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag)
+            self._metrics_last_updated = now
+
+
+def _format_in_threadpool(
+    batch_iter: Iterator[Batch],
+    stats: DatasetStats,
+    batch_format: Optional[str],
+    collate_fn: Optional[Callable[[DataBatch], Any]],
+    num_threadpool_workers: int,
+) -> Iterator[Batch]:
+    """Executes the batching, formatting, and collation logic in a threadpool.
+
+    Args:
+        logical_batch_iterator: An iterator over logical batches.
+        stats: DatasetStats object to record timing and other statistics.
+        batch_format: The format in which to return each batch.
+            Specify "default" to use the current block format (promoting
+            Arrow to pandas automatically), "pandas" to
+            select ``pandas.DataFrame`` or "pyarrow" to select
+            ``pyarrow.Table``, or None to use entire blocks
+            as batches.
+        collate_fn: A function to apply to each data batch before returning it.
+        num_threadpool_workers: The number of threads to use in the threadpool.
+    """
+
+    def threadpool_computations_format_collate(
+        batch_iter: Iterator[Batch],
+    ) -> Iterator[Batch]:
+        # Step 4a: Format the batches.
+        formatted_batch_iter = format_batches(
+            batch_iter, batch_format=batch_format, stats=stats
+        )
+
+        # Step 4b: Apply the collate function if applicable.
+        if collate_fn is not None:
+            formatted_batch_iter = collate(
+                formatted_batch_iter, collate_fn=collate_fn, stats=stats
+            )
+        yield from formatted_batch_iter
+
+    if num_threadpool_workers > 0:
+        collated_iter = make_async_gen(
+            base_iterator=batch_iter,
+            fn=threadpool_computations_format_collate,
+            preserve_ordering=False,
+            num_workers=num_threadpool_workers,
+        )
+    else:
+        collated_iter = threadpool_computations_format_collate(batch_iter)
+    return collated_iter
+
+
+def prefetch_batches_locally(
+    ref_bundles: Iterator[RefBundle],
+    prefetcher: BlockPrefetcher,
+    num_batches_to_prefetch: int,
+    batch_size: Optional[int],
+    eager_free: bool = False,
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[ObjectRef[Block]]:
+    """Given an iterator of batched RefBundles, returns an iterator over the
+    corresponding block references while prefetching `num_batches_to_prefetch`
+    batches in advance.
+
+    Args:
+        ref_bundles: An iterator over batched RefBundles.
+        prefetcher: The prefetcher to use.
+        num_batches_to_prefetch: The number of batches to prefetch ahead of the
+            current batch during the scan.
+        batch_size: User specified batch size, or None to let the system pick.
+        eager_free: Whether to eagerly free the object reference from the object store.
+        stats: Dataset stats object used to store ref bundle retrieval time.
+    """
+
+    def get_next_ref_bundle() -> RefBundle:
+        with stats.iter_get_ref_bundles_s.timer() if stats else nullcontext():
+            return next(ref_bundles)
+
+    sliding_window = collections.deque()
+    current_window_size = 0
+
+    if num_batches_to_prefetch <= 0:
+        for ref_bundle in ref_bundles:
+            for block_ref in ref_bundle.block_refs:
+                yield block_ref
+        return
+
+    if batch_size is not None:
+        num_rows_to_prefetch = num_batches_to_prefetch * batch_size
+    else:
+        num_rows_to_prefetch = None
+
+    # Create and fetch the initial window.
+    # Stop adding if the number of rows in this window is greater than requested
+    # batch size, or if the batch size is None and the number of blocks in this window
+    # is greater than requested batches to prefetch.
+    while (batch_size is not None and current_window_size < num_rows_to_prefetch) or (
+        batch_size is None and len(sliding_window) < num_batches_to_prefetch
+    ):
+        try:
+            next_ref_bundle = get_next_ref_bundle()
+            sliding_window.extend(next_ref_bundle.blocks)
+            current_window_size += next_ref_bundle.num_rows()
+        except StopIteration:
+            break
+
+    prefetcher.prefetch_blocks([block_ref for block_ref, _ in list(sliding_window)])
+
+    while sliding_window:
+        block_ref, metadata = sliding_window.popleft()
+        current_window_size -= metadata.num_rows
+        if batch_size is None or current_window_size < num_rows_to_prefetch:
+            try:
+                next_ref_bundle = get_next_ref_bundle()
+                for block_ref_and_md in next_ref_bundle.blocks:
+                    sliding_window.append(block_ref_and_md)
+                    current_window_size += block_ref_and_md[1].num_rows
+                prefetcher.prefetch_blocks(
+                    [block_ref for block_ref, _ in list(sliding_window)]
+                )
+            except StopIteration:
+                pass
+        yield block_ref
+        trace_deallocation(block_ref, loc="iter_batches", free=eager_free)
+    prefetcher.stop()
+
+
+def restore_original_order(batch_iter: Iterator[Batch]) -> Iterator[Batch]:
+    """Restores the original order of the provided `batch_iter`
+
+    This function will yield items from `base_iterator` in the correct order based on
+    each batch's batch_idx. All indexes are expected to be unique.
+
+    `batch_iter` is expected to not have any missing indexes. All indexes from 0 to len
+    (base_iterator) must be present.
+    """
+    next_index_required = 0
+    buffer: Dict[int, Batch] = {}
+    for batch in batch_iter:
+        assert batch.metadata.batch_idx not in buffer
+        buffer[batch.metadata.batch_idx] = batch
+        while next_index_required in buffer:
+            yield buffer.pop(next_index_required)
+            next_index_required += 1
+
+    while next_index_required in buffer:
+        yield buffer.pop(next_index_required)
+        next_index_required += 1
diff --git a/lib/python3.12/site-packages/ray/data/_internal/block_batching/util.py b/lib/python3.12/site-packages/ray/data/_internal/block_batching/util.py
new file mode 100644
index 0000000000000000000000000000000000000000..8896b1199ec9c784add6351e10c37d924a8cb68c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/block_batching/util.py
@@ -0,0 +1,303 @@
+import dataclasses
+import logging
+import threading
+from contextlib import nullcontext
+from typing import Any, Callable, Iterator, List, Optional, Tuple
+
+import ray
+from ray.actor import ActorHandle
+from ray.data._internal.batcher import Batcher, ShufflingBatcher
+from ray.data._internal.block_batching.interfaces import (
+    Batch,
+    BatchMetadata,
+    BlockPrefetcher,
+    CollatedBatch,
+)
+from ray.data._internal.stats import DatasetStats
+from ray.data.block import Block, BlockAccessor, DataBatch
+from ray.types import ObjectRef
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+logger = logging.getLogger(__name__)
+
+
+def _calculate_ref_hits(refs: List[ObjectRef[Any]]) -> Tuple[int, int, int]:
+    """Given a list of object references, returns how many are already on the local
+    node, how many require fetching from another node, and how many have unknown
+    locations. If `DataContext.get_current().enable_get_object_locations_for_metrics` is
+    False, this will return `(-1, -1, -1)` as getting object locations is disabled."""
+    current_node_id = ray.get_runtime_context().get_node_id()
+
+    ctx = ray.data.context.DataContext.get_current()
+    if ctx.enable_get_object_locations_for_metrics:
+        locs = ray.experimental.get_object_locations(refs)
+        nodes: List[List[str]] = [loc["node_ids"] for loc in locs.values()]
+        hits = sum(current_node_id in node_ids for node_ids in nodes)
+        unknowns = sum(1 for node_ids in nodes if not node_ids)
+        misses = len(nodes) - hits - unknowns
+        return hits, misses, unknowns
+
+    return -1, -1, -1
+
+
+def resolve_block_refs(
+    block_ref_iter: Iterator[ObjectRef[Block]],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[Block]:
+    """Resolves the block references for each logical batch.
+
+    Args:
+        block_ref_iter: An iterator over block object references.
+        stats: An optional stats object to recording block hits and misses.
+    """
+    hits = 0
+    misses = 0
+    unknowns = 0
+
+    for block_ref in block_ref_iter:
+        current_hit, current_miss, current_unknown = _calculate_ref_hits([block_ref])
+        hits += current_hit
+        misses += current_miss
+        unknowns += current_unknown
+
+        # TODO(amogkam): Optimized further by batching multiple references in a single
+        # `ray.get()` call.
+        with stats.iter_get_s.timer() if stats else nullcontext():
+            block = ray.get(block_ref)
+        yield block
+
+    if stats:
+        stats.iter_blocks_local = hits
+        stats.iter_blocks_remote = misses
+        stats.iter_unknown_location = unknowns
+
+
+def blocks_to_batches(
+    block_iter: Iterator[Block],
+    stats: Optional[DatasetStats] = None,
+    batch_size: Optional[int] = None,
+    drop_last: bool = False,
+    shuffle_buffer_min_size: Optional[int] = None,
+    shuffle_seed: Optional[int] = None,
+    ensure_copy: bool = False,
+) -> Iterator[Batch]:
+    """Given an iterator over blocks, returns an iterator over blocks
+    of the appropriate bacth size.
+
+    If the shuffling configurations are specified, then the
+    output blocks contain shuffled data.
+
+    Args:
+        block_iter: An iterator over blocks.
+        stats: Dataset stats object used to store block batching time.
+        batch_size: Record batch size, or None to let the system pick.
+        drop_last: Whether to drop the last batch if it's incomplete.
+        shuffle_buffer_min_size: If non-None, the data will be randomly shuffled
+            using a local in-memory shuffle buffer, and this value will serve as the
+            minimum number of rows that must be in the local in-memory shuffle buffer in
+            order to yield a batch.
+        shuffle_seed: The seed to use for the local random shuffle.
+        ensure_copy: Whether batches are always copied from the underlying base
+            blocks (not zero-copy views).
+
+    Returns:
+        An iterator over blocks of the given size that are potentially shuffled.
+    """
+    if shuffle_buffer_min_size is not None:
+        batcher = ShufflingBatcher(
+            batch_size=batch_size,
+            shuffle_buffer_min_size=shuffle_buffer_min_size,
+            shuffle_seed=shuffle_seed,
+        )
+    else:
+        batcher = Batcher(batch_size=batch_size, ensure_copy=ensure_copy)
+
+    def get_iter_next_batch_s_timer():
+        return stats.iter_next_batch_s.timer() if stats else nullcontext()
+
+    global_counter = 0
+
+    for block in block_iter:
+        batcher.add(block)
+        while batcher.has_batch():
+            with get_iter_next_batch_s_timer():
+                batch = batcher.next_batch()
+            yield Batch(metadata=BatchMetadata(batch_idx=global_counter), data=batch)
+            global_counter += 1
+
+    # Signal to the batcher that there are no more blocks to add.
+    batcher.done_adding()
+
+    # Get any leftover batches in ShufflingBatcher.
+    while batcher.has_batch():
+        with get_iter_next_batch_s_timer():
+            batch = batcher.next_batch()
+        yield Batch(metadata=BatchMetadata(batch_idx=global_counter), data=batch)
+        global_counter += 1
+
+    # Get any remaining data.
+    if not drop_last and batcher.has_any():
+        with get_iter_next_batch_s_timer():
+            batch = batcher.next_batch()
+        yield Batch(metadata=BatchMetadata(batch_idx=global_counter), data=batch)
+        global_counter += 1
+
+
+def format_batches(
+    batch_iter: Iterator[Batch],
+    batch_format: Optional[str],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[Batch]:
+    """Given an iterator of blocks, returns an iterator of formatted batches.
+
+    Args:
+        batch_iter: An iterator over batches.
+        batch_format: The batch format to use.
+        stats: An optional stats object to record formatting times.
+
+    Returns:
+        An iterator over batch index and the formatted batch.
+    """
+    for batch in batch_iter:
+        with stats.iter_format_batch_s.timer() if stats else nullcontext():
+            formatted_batch = BlockAccessor.for_block(batch.data).to_batch_format(
+                batch_format
+            )
+        yield dataclasses.replace(batch, data=formatted_batch)
+
+
+def collate(
+    batch_iter: Iterator[Batch],
+    collate_fn: Optional[Callable[[DataBatch], Any]],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[CollatedBatch]:
+    """Returns an iterator with the provided collate_fn applied to items of the batch
+    iterator.
+
+    Args:
+        batch_iter: An iterator over formatted batches.
+        collate_fn: A function to apply to each batch.
+        stats: An optional stats object to record formatting times.
+    """
+    for batch in batch_iter:
+        with stats.iter_collate_batch_s.timer() if stats else nullcontext():
+            collated_batch = collate_fn(batch.data)
+        yield CollatedBatch(metadata=batch.metadata, data=collated_batch)
+
+
+def finalize_batches(
+    batch_iter: Iterator[CollatedBatch],
+    finalize_fn: Callable[[Any], Any],
+    stats: Optional[DatasetStats] = None,
+) -> Iterator[CollatedBatch]:
+    """Returns an iterator with the provided finalize_fn applied to items of the batch
+    iterator.
+
+    This is the same as `collate` except the input batches can be of type Any.
+
+    Args:
+        batch_iter: An iterator over processed batches.
+        finalize_fn: A function to apply to each batch.
+        stats: An optional stats object to record formatting times.
+
+    Returns:
+        An iterator over batch index and the finalized batch.
+    """
+    for batch in batch_iter:
+        with stats.iter_finalize_batch_s.timer() if stats else nullcontext():
+            finalized_batch = finalize_fn(batch.data)
+        yield dataclasses.replace(batch, data=finalized_batch)
+
+
+def extract_data_from_batch(batch_iter: Iterator[Batch]) -> Iterator[Any]:
+    for batch in batch_iter:
+        yield batch.data
+
+
+PREFETCHER_ACTOR_NAMESPACE = "ray.dataset"
+
+
+class WaitBlockPrefetcher(BlockPrefetcher):
+    """Block prefetcher using ray.wait."""
+
+    def __init__(self):
+        self._blocks = []
+        self._stopped = False
+        self._condition = threading.Condition()
+        self._thread = threading.Thread(
+            target=self._run,
+            name="Prefetcher",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def _run(self):
+        while not self._stopped:
+            try:
+                with self._condition:
+                    if len(self._blocks) == 0:
+                        # Park, waiting for notification that prefetching
+                        # should resume
+                        self._condition.wait()
+
+                    blocks_to_fetch, self._blocks = self._blocks[:], []
+
+                if len(blocks_to_fetch) > 0:
+                    ray.wait(
+                        blocks_to_fetch,
+                        num_returns=1,
+                        # NOTE: We deliberately setting timeout to 0 to avoid
+                        #       blocking the fetching thread unnecessarily
+                        timeout=0,
+                        fetch_local=True,
+                    )
+            except Exception:
+                logger.exception("Error in prefetcher thread.")
+
+        logger.info("Exiting prefetcher's background thread")
+
+    def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
+        with self._condition:
+            if self._stopped:
+                raise RuntimeError("Prefetcher is stopped.")
+            self._blocks = blocks
+            self._condition.notify()
+
+    def stop(self):
+        with self._condition:
+            if self._stopped:
+                return
+            self._stopped = True
+            self._condition.notify()
+
+    def __del__(self):
+        self.stop()
+
+
+class ActorBlockPrefetcher(BlockPrefetcher):
+    """Block prefetcher using a local actor."""
+
+    def __init__(self):
+        self.prefetch_actor = self._get_or_create_actor_prefetcher()
+
+    @staticmethod
+    def _get_or_create_actor_prefetcher() -> "ActorHandle":
+        node_id = ray.get_runtime_context().get_node_id()
+        actor_name = f"dataset-block-prefetcher-{node_id}"
+        return _BlockPretcher.options(
+            scheduling_strategy=NodeAffinitySchedulingStrategy(node_id, soft=False),
+            name=actor_name,
+            namespace=PREFETCHER_ACTOR_NAMESPACE,
+            get_if_exists=True,
+        ).remote()
+
+    def prefetch_blocks(self, blocks: List[ObjectRef[Block]]):
+        self.prefetch_actor.prefetch.remote(*blocks)
+
+
+@ray.remote(num_cpus=0)
+class _BlockPretcher:
+    """Helper actor that prefetches blocks asynchronously."""
+
+    def prefetch(self, *blocks) -> None:
+        pass
diff --git a/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..01e15270a0c47a85bba34d164658fecf9ba9fc5e
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__init__.py
@@ -0,0 +1,19 @@
+from typing import TYPE_CHECKING
+
+from .base_cluster_autoscaler import ClusterAutoscaler
+from .default_cluster_autoscaler import DefaultClusterAutoscaler
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+def create_cluster_autoscaler(
+    topology: "Topology", resource_manager: "ResourceManager", *, execution_id: str
+) -> ClusterAutoscaler:
+    return DefaultClusterAutoscaler(
+        topology, resource_manager, execution_id=execution_id
+    )
+
+
+__all__ = ["ClusterAutoscaler"]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d3d6db0e4e81cdab68f5306189618c4467fa72dd
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/base_cluster_autoscaler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/base_cluster_autoscaler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bfa023b76a873e09a2cd0a145949adaac4b457a3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/base_cluster_autoscaler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/default_cluster_autoscaler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/default_cluster_autoscaler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f2592a1bc36a84f768adb735a533215714df887d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/__pycache__/default_cluster_autoscaler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/base_cluster_autoscaler.py b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/base_cluster_autoscaler.py
new file mode 100644
index 0000000000000000000000000000000000000000..cdca9187a70f052deac2a8b48c6e6430b8434481
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/base_cluster_autoscaler.py
@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.execution_options import (
+        ExecutionResources,
+    )
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+@DeveloperAPI
+class ClusterAutoscaler(ABC):
+    """Abstract interface for Ray Data cluster autoscaler."""
+
+    def __init__(
+        self,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+        execution_id: str,
+    ):
+        self._topology = topology
+        self._resource_manager = resource_manager
+        self._execution_id = execution_id
+
+    @abstractmethod
+    def try_trigger_scaling(self):
+        """Try trigger autoscaling.
+
+        This method will be called each time when StreamingExecutor makes
+        a scheduling decision. A subclass should override this method to
+        handle the autoscaling of the cluster.
+        """
+        ...
+
+    @abstractmethod
+    def on_executor_shutdown(self):
+        """Callback when the StreamingExecutor is shutting down."""
+        ...
+
+    @abstractmethod
+    def get_total_resources(self) -> "ExecutionResources":
+        """Get the total resources that are available to this data execution."""
+        ...
diff --git a/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/default_cluster_autoscaler.py b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/default_cluster_autoscaler.py
new file mode 100644
index 0000000000000000000000000000000000000000..b2a0e0d94c548e91662ab4135789c75fb1bb8989
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/cluster_autoscaler/default_cluster_autoscaler.py
@@ -0,0 +1,106 @@
+import logging
+import math
+import time
+from typing import TYPE_CHECKING, Dict
+
+import ray
+from .base_cluster_autoscaler import ClusterAutoscaler
+from ray.data._internal.execution.autoscaling_requester import (
+    get_or_create_autoscaling_requester_actor,
+)
+from ray.data._internal.execution.interfaces import ExecutionResources
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+logger = logging.getLogger(__name__)
+
+
+class DefaultClusterAutoscaler(ClusterAutoscaler):
+
+    # Min number of seconds between two autoscaling requests.
+    MIN_GAP_BETWEEN_AUTOSCALING_REQUESTS = 20
+
+    def __init__(
+        self,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+        *,
+        execution_id: str,
+    ):
+        super().__init__(topology, resource_manager, execution_id)
+
+        # Last time when a request was sent to Ray's autoscaler.
+        self._last_request_time = 0
+
+    def try_trigger_scaling(self):
+        """Try to scale up the cluster to accommodate the provided in-progress workload.
+
+        This makes a resource request to Ray's autoscaler consisting of the current,
+        aggregate usage of all operators in the DAG + the incremental usage of all
+        operators that are ready for dispatch (i.e. that have inputs queued). If the
+        autoscaler were to grant this resource request, it would allow us to dispatch
+        one task for every ready operator.
+
+        Note that this resource request does not take the global resource limits or the
+        liveness policy into account; it only tries to make the existing resource usage
+        + one more task per ready operator feasible in the cluster.
+        """
+        # Limit the frequency of autoscaling requests.
+        now = time.time()
+        if now - self._last_request_time < self.MIN_GAP_BETWEEN_AUTOSCALING_REQUESTS:
+            return
+
+        # Scale up the cluster, if no ops are allowed to run, but there are still data
+        # in the input queues.
+        no_runnable_op = all(
+            not op_state._scheduling_status.runnable
+            for _, op_state in self._topology.items()
+        )
+        any_has_input = any(
+            op_state.has_pending_bundles() for _, op_state in self._topology.items()
+        )
+        if not (no_runnable_op and any_has_input):
+            return
+
+        self._last_request_time = now
+
+        # Get resource usage for all ops + additional resources needed to launch one
+        # more task for each ready op.
+        resource_request = []
+
+        def to_bundle(resource: ExecutionResources) -> Dict:
+            req = {}
+            if resource.cpu:
+                req["CPU"] = math.ceil(resource.cpu)
+            if resource.gpu:
+                req["GPU"] = math.ceil(resource.gpu)
+            return req
+
+        for op, state in self._topology.items():
+            per_task_resource = op.incremental_resource_usage()
+            task_bundle = to_bundle(per_task_resource)
+            resource_request.extend([task_bundle] * op.num_active_tasks())
+            # Only include incremental resource usage for ops that are ready for
+            # dispatch.
+            if state.has_pending_bundles():
+                # TODO(Clark): Scale up more aggressively by adding incremental resource
+                # usage for more than one bundle in the queue for this op?
+                resource_request.append(task_bundle)
+
+        self._send_resource_request(resource_request)
+
+    def _send_resource_request(self, resource_request):
+        # Make autoscaler resource request.
+        actor = get_or_create_autoscaling_requester_actor()
+        actor.request_resources.remote(resource_request, self._execution_id)
+
+    def on_executor_shutdown(self):
+        # Make request for zero resources to autoscaler for this execution.
+        actor = get_or_create_autoscaling_requester_actor()
+        actor.request_resources.remote({}, self._execution_id)
+
+    def get_total_resources(self) -> ExecutionResources:
+        return ExecutionResources.from_resource_dict(ray.cluster_resources())
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3e7e2e2d27b07707b0340e71d15995dd91082906
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/audio_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/audio_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7d85254af0361d2a20d46a4572825eb314bea124
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/audio_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/avro_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/avro_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..10be5dca7aec710e3311efbac844c540e65ff619
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/avro_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/bigquery_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/bigquery_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a91c39a104a575b43e6d6d3f2c4bd38034d3f196
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/bigquery_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/bigquery_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/bigquery_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3d1c0a97a18b83c87a5864b6e30a450204b77633
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/bigquery_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/binary_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/binary_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..14066e933c1be6d15d7df84dcab0e0ba3c90d42d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/binary_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/clickhouse_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/clickhouse_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..02fd6f5746e0436b56e915760ae0438369150a9c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/clickhouse_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/clickhouse_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/clickhouse_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..99cbf04241589e5ab33c45cb3c46a01795b86c07
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/clickhouse_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/csv_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/csv_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..187c3546220b974523fb72c77233f6223ebafae7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/csv_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/csv_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/csv_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cebb9ae80291239e0c9953cf7cc989dad1ce15c2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/csv_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/databricks_uc_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/databricks_uc_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..49a1c64c8a656d36254becc44ca1fb80ed340fc7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/databricks_uc_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/delta_sharing_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/delta_sharing_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..01fae9cf6b7dd1cebd77fc844332911ab3c68143
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/delta_sharing_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/hudi_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/hudi_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c3ca77fa67a1c51261c394af6e67694cc59a5549
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/hudi_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/huggingface_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/huggingface_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9289fa8f23cb46d07ffcd44ad8214e06b4af78a0
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/huggingface_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/iceberg_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/iceberg_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..03383e4415a4896feeebdc199572bba81a4ace05
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/iceberg_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/iceberg_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/iceberg_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c06992dace33050af204c7ce457a9a229b0d762a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/iceberg_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/image_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/image_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9cb07bcf8a507660d3cec309be25886cf6f28fd9
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/image_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/image_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/image_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..46c4087b759c0d194316a37a899d0561eea4996e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/image_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/json_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/json_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b6268914a8f978d6ebc538bc0ae9327f625f77d0
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/json_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/json_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/json_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..edad1b249f51f979bdf194797eac4a2121ea3e38
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/json_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/lance_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/lance_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d1d425f12ad3210f8f760c3e20e81ffa498e4d66
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/lance_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/lance_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/lance_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5d14e577baec0c2f5ca0382a0d9cc4b845818983
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/lance_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mcap_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mcap_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fd9b78ff02e5e35c4308d289b2a075f6ac633eff
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mcap_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mongo_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mongo_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8a03cef557db741426cb90ba2849308fa25017a3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mongo_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mongo_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mongo_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0631e583778a89e6366ac369111bce6d8c27b48d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/mongo_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/numpy_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/numpy_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9d13962016128b6ca22598cdd1058161d754dae9
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/numpy_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/numpy_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/numpy_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..875b40a8638763731798cc950d78a1536541a06f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/numpy_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_bulk_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_bulk_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2f3806adfa1d4e368ce5921e0d6a971281bc7466
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_bulk_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..153d8127eb68aec58fbb21b9cf718ef45605f55d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2bd2b2a8fef525eb00b83419a21d8a9320f10a2e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/parquet_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/range_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/range_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4b3c2e3b928d07af8912e04ea764e8d6dbc2b37c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/range_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/sql_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/sql_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2dff44b7dfa430e412d6d9468de0d6711b88b2e2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/sql_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/sql_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/sql_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4735e7add45a9e44629794278dedff50d3008b93
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/sql_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/text_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/text_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..281fef9e43a0e1ea54c86f6307f5cc4e152af820
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/text_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/tfrecords_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/tfrecords_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8e6881d4114c2c3f1936a98891f907bb8b179450
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/tfrecords_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/tfrecords_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/tfrecords_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..304692cd4610f24bb62e0722c682f98c7272d934
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/tfrecords_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/torch_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/torch_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..71fc0956911da6d710c42793c8c8eaef93d684da
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/torch_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/uc_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/uc_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ae862f184b85ed06eba0feba88c659911d0f3983
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/uc_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/video_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/video_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e7ce9e72f87ffebe4abeae949a089bc657fecf53
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/video_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/webdataset_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/webdataset_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..65a0c6dd52043a84ae3382caa2b813954baeb182
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/webdataset_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/webdataset_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/webdataset_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e6c19232e69adc3dd20ca027bd8a819dde456851
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/datasource/__pycache__/webdataset_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/audio_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/audio_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..3a3a28d26f45fdee6cd63ccce32cd4b8f9bb6e81
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/audio_datasource.py
@@ -0,0 +1,57 @@
+import io
+from typing import TYPE_CHECKING, Iterator, List, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class AudioDatasource(FileBasedDatasource):
+    _FILE_EXTENSIONS = [
+        "mp3",
+        "wav",
+        "aac",
+        "flac",
+        "ogg",
+        "m4a",
+        "wma",
+        "alac",
+        "aiff",
+        "pcm",
+        "amr",
+        "opus",
+        "ra",
+        "rm",
+        "au",
+        "mid",
+        "midi",
+        "caf",
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="soundfile", package="soundfile")
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import soundfile
+
+        # `soundfile` doesn't support reading from a `pyarrow.NativeFile` directly, so
+        # we need to read the file into memory first.
+        stream = io.BytesIO(f.read())
+        amplitude, sample_rate = soundfile.read(stream, always_2d=True, dtype="float32")
+
+        # (amplitude, channels) -> (channels, amplitude)
+        amplitude = amplitude.transpose((1, 0))
+
+        builder = DelegatingBlockBuilder()
+        builder.add({"amplitude": amplitude, "sample_rate": sample_rate})
+        yield builder.build()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/avro_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/avro_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..21f956698b051ff558ccc81ef54aea9e493abd93
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/avro_datasource.py
@@ -0,0 +1,45 @@
+from typing import TYPE_CHECKING, Iterator, List, Union
+
+from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.context import DataContext
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class AvroDatasource(FileBasedDatasource):
+    """A datasource that reads Avro files."""
+
+    _FILE_EXTENSIONS = ["avro"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="fastavro", package="fastavro")
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import fastavro
+
+        # Read the Avro file. This assumes the Avro file includes its schema.
+        reader = fastavro.reader(f)
+
+        ctx = DataContext.get_current()
+        output_block_size_option = OutputBlockSizeOption.of(
+            target_max_block_size=ctx.target_max_block_size
+        )
+        output_buffer = BlockOutputBuffer(output_block_size_option)
+        for record in reader:
+            output_buffer.add(record)
+            while output_buffer.has_next():
+                yield output_buffer.next()
+
+        output_buffer.finalize()
+        while output_buffer.has_next():
+            yield output_buffer.next()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/bigquery_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/bigquery_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..96be61825b44705ac483406b8d1eeda58be83092
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/bigquery_datasink.py
@@ -0,0 +1,129 @@
+import logging
+import os
+import tempfile
+import time
+import uuid
+from typing import Iterable, Optional
+
+import pyarrow.parquet as pq
+
+import ray
+from ray.data._internal.datasource import bigquery_datasource
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_RETRY_CNT = 10
+RATE_LIMIT_EXCEEDED_SLEEP_TIME = 11
+
+
+class BigQueryDatasink(Datasink[None]):
+    def __init__(
+        self,
+        project_id: str,
+        dataset: str,
+        max_retry_cnt: int = DEFAULT_MAX_RETRY_CNT,
+        overwrite_table: Optional[bool] = True,
+    ) -> None:
+        _check_import(self, module="google.cloud", package="bigquery")
+        _check_import(self, module="google.cloud", package="bigquery_storage")
+        _check_import(self, module="google.api_core", package="exceptions")
+
+        self.project_id = project_id
+        self.dataset = dataset
+        self.max_retry_cnt = max_retry_cnt
+        self.overwrite_table = overwrite_table
+
+    def on_write_start(self) -> None:
+        from google.api_core import exceptions
+
+        if self.project_id is None or self.dataset is None:
+            raise ValueError("project_id and dataset are required args")
+
+        # Set up datasets to write
+        client = bigquery_datasource._create_client(project_id=self.project_id)
+        dataset_id = self.dataset.split(".", 1)[0]
+        try:
+            client.get_dataset(dataset_id)
+        except exceptions.NotFound:
+            client.create_dataset(f"{self.project_id}.{dataset_id}", timeout=30)
+            logger.info("Created dataset " + dataset_id)
+
+        # Delete table if overwrite_table is True
+        if self.overwrite_table:
+            logger.info(
+                f"Attempting to delete table {self.dataset}"
+                + " if it already exists since kwarg overwrite_table = True."
+            )
+            client.delete_table(f"{self.project_id}.{self.dataset}", not_found_ok=True)
+        else:
+            logger.info(
+                f"The write will append to table {self.dataset}"
+                + " if it already exists since kwarg overwrite_table = False."
+            )
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        def _write_single_block(block: Block, project_id: str, dataset: str) -> None:
+            from google.api_core import exceptions
+            from google.cloud import bigquery
+
+            block = BlockAccessor.for_block(block).to_arrow()
+
+            client = bigquery_datasource._create_client(project_id=project_id)
+            job_config = bigquery.LoadJobConfig(autodetect=True)
+            job_config.source_format = bigquery.SourceFormat.PARQUET
+            job_config.write_disposition = bigquery.WriteDisposition.WRITE_APPEND
+
+            with tempfile.TemporaryDirectory() as temp_dir:
+                fp = os.path.join(temp_dir, f"block_{uuid.uuid4()}.parquet")
+                pq.write_table(block, fp, compression="SNAPPY")
+
+                retry_cnt = 0
+                while retry_cnt <= self.max_retry_cnt:
+                    with open(fp, "rb") as source_file:
+                        job = client.load_table_from_file(
+                            source_file, dataset, job_config=job_config
+                        )
+                    try:
+                        logger.info(job.result())
+                        break
+                    except (exceptions.Forbidden, exceptions.TooManyRequests) as e:
+                        retry_cnt += 1
+                        if retry_cnt > self.max_retry_cnt:
+                            break
+                        logger.info(
+                            "A block write encountered a rate limit exceeded error"
+                            + f" {retry_cnt} time(s). Sleeping to try again."
+                        )
+                        logging.debug(e)
+                        time.sleep(RATE_LIMIT_EXCEEDED_SLEEP_TIME)
+
+                # Raise exception if retry_cnt exceeds max_retry_cnt
+                if retry_cnt > self.max_retry_cnt:
+                    logger.info(
+                        f"Maximum ({self.max_retry_cnt}) retry count exceeded. Ray"
+                        + " will attempt to retry the block write via fault tolerance."
+                    )
+                    raise RuntimeError(
+                        f"Write failed due to {retry_cnt}"
+                        + " repeated API rate limit exceeded responses. Consider"
+                        + " specifiying the max_retry_cnt kwarg with a higher value."
+                    )
+
+        _write_single_block = cached_remote_fn(_write_single_block)
+
+        # Launch a remote task for each block within this write task
+        ray.get(
+            [
+                _write_single_block.remote(block, self.project_id, self.dataset)
+                for block in blocks
+            ]
+        )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/bigquery_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/bigquery_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2e01ec70a03b86d8201e8ba8efb87a58168eec5
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/bigquery_datasource.py
@@ -0,0 +1,158 @@
+import logging
+from typing import List, Optional
+
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+logger = logging.getLogger(__name__)
+
+
+def _create_user_agent() -> str:
+    import ray
+
+    return f"ray/{ray.__version__}"
+
+
+def _create_client_info():
+    from google.api_core.client_info import ClientInfo
+
+    return ClientInfo(
+        user_agent=_create_user_agent(),
+    )
+
+
+def _create_client_info_gapic():
+    from google.api_core.gapic_v1.client_info import ClientInfo
+
+    return ClientInfo(
+        user_agent=_create_user_agent(),
+    )
+
+
+def _create_client(project_id: str):
+    from google.cloud import bigquery
+
+    return bigquery.Client(
+        project=project_id,
+        client_info=_create_client_info(),
+    )
+
+
+def _create_read_client():
+    from google.cloud import bigquery_storage
+
+    return bigquery_storage.BigQueryReadClient(
+        client_info=_create_client_info_gapic(),
+    )
+
+
+class BigQueryDatasource(Datasource):
+    def __init__(
+        self,
+        project_id: str,
+        dataset: Optional[str] = None,
+        query: Optional[str] = None,
+    ):
+        _check_import(self, module="google.cloud", package="bigquery")
+        _check_import(self, module="google.cloud", package="bigquery_storage")
+        _check_import(self, module="google.api_core", package="exceptions")
+
+        self._project_id = project_id
+        self._dataset = dataset
+        self._query = query
+
+        if query is not None and dataset is not None:
+            raise ValueError(
+                "Query and dataset kwargs cannot both be provided "
+                + "(must be mutually exclusive)."
+            )
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        from google.cloud import bigquery_storage
+
+        def _read_single_partition(stream) -> Block:
+            client = _create_read_client()
+            reader = client.read_rows(stream.name)
+            return reader.to_arrow()
+
+        if self._query:
+            query_client = _create_client(project_id=self._project_id)
+            query_job = query_client.query(self._query)
+            query_job.result()
+            destination = str(query_job.destination)
+            dataset_id = destination.split(".")[-2]
+            table_id = destination.split(".")[-1]
+        else:
+            self._validate_dataset_table_exist(self._project_id, self._dataset)
+            dataset_id = self._dataset.split(".")[0]
+            table_id = self._dataset.split(".")[1]
+
+        bqs_client = _create_read_client()
+        table = f"projects/{self._project_id}/datasets/{dataset_id}/tables/{table_id}"
+
+        if parallelism == -1:
+            parallelism = None
+        requested_session = bigquery_storage.types.ReadSession(
+            table=table,
+            data_format=bigquery_storage.types.DataFormat.ARROW,
+        )
+        read_session = bqs_client.create_read_session(
+            parent=f"projects/{self._project_id}",
+            read_session=requested_session,
+            max_stream_count=parallelism,
+        )
+
+        read_tasks = []
+        logger.info("Created streams: " + str(len(read_session.streams)))
+        if len(read_session.streams) < parallelism:
+            logger.info(
+                "The number of streams created by the "
+                + "BigQuery Storage Read API is less than the requested "
+                + "parallelism due to the size of the dataset."
+            )
+
+        for stream in read_session.streams:
+            # Create a metadata block object to store schema, etc.
+            metadata = BlockMetadata(
+                num_rows=None,
+                size_bytes=None,
+                input_files=None,
+                exec_stats=None,
+            )
+
+            # Create the read task and pass the no-arg wrapper and metadata in
+            read_task = ReadTask(
+                lambda stream=stream: [_read_single_partition(stream)],
+                metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def _validate_dataset_table_exist(self, project_id: str, dataset: str) -> None:
+        from google.api_core import exceptions
+
+        client = _create_client(project_id=project_id)
+        dataset_id = dataset.split(".")[0]
+        try:
+            client.get_dataset(dataset_id)
+        except exceptions.NotFound:
+            raise ValueError(
+                "Dataset {} is not found. Please ensure that it exists.".format(
+                    dataset_id
+                )
+            )
+
+        try:
+            client.get_table(dataset)
+        except exceptions.NotFound:
+            raise ValueError(
+                "Table {} is not found. Please ensure that it exists.".format(dataset)
+            )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/binary_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/binary_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..31e6d89969c61cc9d12b6dc289bee35bd7cc9f55
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/binary_datasource.py
@@ -0,0 +1,24 @@
+from typing import TYPE_CHECKING
+
+from ray.data._internal.arrow_block import ArrowBlockBuilder
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class BinaryDatasource(FileBasedDatasource):
+    """Binary datasource, for reading and writing binary files."""
+
+    _COLUMN_NAME = "bytes"
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        data = f.readall()
+
+        builder = ArrowBlockBuilder()
+        item = {self._COLUMN_NAME: data}
+        builder.add(item)
+        yield builder.build()
+
+    def _rows_per_file(self):
+        return 1
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/clickhouse_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/clickhouse_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c80081dcd8a3d561ea8443028e3904c69edb06d
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/clickhouse_datasink.py
@@ -0,0 +1,430 @@
+import logging
+import re
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Optional,
+)
+
+import pyarrow
+import pyarrow.types as pat
+
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink, WriteReturnType
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DECIMAL_PRECISION = 38
+DEFAULT_DECIMAL_SCALE = 10
+
+
+def _pick_best_arrow_field_for_order_by(schema: pyarrow.Schema) -> str:
+    if len(schema) == 0:
+        return "tuple()"
+    # Prefer a timestamp column if available
+    for f in schema:
+        if pat.is_timestamp(f.type):
+            return f.name
+    # Next prefer a non-string column
+    for f in schema:
+        if not (pat.is_string(f.type) or pat.is_large_string(f.type)):
+            return f.name
+    # Otherwise pick the first column
+    return schema[0].name
+
+
+def _arrow_to_clickhouse_type(field: pyarrow.Field) -> str:
+    """Convert a PyArrow field to an appropriate ClickHouse column type."""
+    t = field.type
+    if pat.is_decimal(t):
+        precision = t.precision or DEFAULT_DECIMAL_PRECISION
+        scale = t.scale or DEFAULT_DECIMAL_SCALE
+        return f"Decimal({precision}, {scale})"
+    if pat.is_boolean(t):
+        return "UInt8"
+    if pat.is_int8(t):
+        return "Int8"
+    if pat.is_int16(t):
+        return "Int16"
+    if pat.is_int32(t):
+        return "Int32"
+    if pat.is_int64(t):
+        return "Int64"
+    if pat.is_uint8(t):
+        return "UInt8"
+    if pat.is_uint16(t):
+        return "UInt16"
+    if pat.is_uint32(t):
+        return "UInt32"
+    if pat.is_uint64(t):
+        return "UInt64"
+    if pat.is_float16(t):
+        return "Float32"
+    if pat.is_float32(t):
+        return "Float32"
+    if pat.is_float64(t):
+        return "Float64"
+    if pat.is_timestamp(t):
+        return "DateTime64(3)"
+    return "String"
+
+
+@dataclass
+class ClickHouseTableSettings:
+    """
+    Additional table creation instructions for ClickHouse.
+
+    Attributes:
+        engine: The engine definition for the created table. Defaults
+            to "MergeTree()".
+        order_by: The ORDER BY clause for the table.
+        partition_by: The PARTITION BY clause for the table.
+        primary_key: The PRIMARY KEY clause for the table.
+        settings: Additional SETTINGS clause for the table
+            (comma-separated or any valid string).
+    """
+
+    engine: str = "MergeTree()"
+    order_by: Optional[str] = None
+    partition_by: Optional[str] = None
+    primary_key: Optional[str] = None
+    settings: Optional[str] = None
+
+
+@PublicAPI(stability="alpha")
+class SinkMode(IntEnum):
+    """
+    Enum of possible modes for sinking data
+
+    Attributes:
+        CREATE: Create a new table; fail if that table already exists.
+        APPEND: Use an existing table if present, otherwise create one; then append data.
+        OVERWRITE: Drop the table if it already exists, then re-create it and write.
+    """
+
+    # Create a new table and fail if that table already exists.
+    CREATE = 1
+
+    # Append data to an existing table, or create one if it does not exist.
+    APPEND = 2
+
+    # Drop the table if it already exists, then re-create it and write.
+    OVERWRITE = 3
+
+
+@DeveloperAPI
+class ClickHouseDatasink(Datasink):
+    """ClickHouse Ray Datasink.
+
+    A Ray Datasink for writing data into ClickHouse, with support for distributed
+    writes and mode-based table management (create, append, or overwrite).
+
+    Args:
+        table: Fully qualified table identifier (e.g., "default.my_table").
+        dsn: A string in DSN (Data Source Name) HTTP format
+            (e.g., "clickhouse+http://username:password@host:8123/default").
+            For more information, see `ClickHouse Connection String doc
+            <https://clickhouse.com/docs/en/integrations/sql-clients/cli#connection_string>`_.
+        mode: One of SinkMode.CREATE, SinkMode.APPEND,
+            or SinkMode.OVERWRITE.
+            - **CREATE**: Create a new table; fail if that table already exists.
+              Requires a user-supplied schema if the table doesn’t already exist.
+            - **APPEND**: Use an existing table if present, otherwise create one.
+              If the table does not exist, the user must supply a schema. Data
+              is then appended to the table.
+            - **OVERWRITE**: Drop the table if it exists, then re-create it.
+              **Always requires** a user-supplied schema to define the new table.
+        schema: An optional PyArrow schema object that, if provided, will
+            override any inferred schema for table creation.
+            - If you are creating a new table (CREATE or APPEND when the table
+              doesn’t exist) or overwriting an existing table, you **must**
+              provide a schema.
+            - If you’re appending to an already-existing table, the schema is
+              not strictly required unless you want to cast data or enforce
+              column types. If omitted, the existing table definition is used.
+        client_settings: Optional ClickHouse server settings to be used with the
+            session/every request.
+        client_kwargs: Additional keyword arguments to pass to the
+            ClickHouse client.
+        table_settings: An optional dataclass with additional table creation
+            instructions (e.g., engine, order_by, partition_by, primary_key, settings).
+        max_insert_block_rows: If you have extremely large blocks, specifying
+            a limit here will chunk the insert into multiple smaller insert calls.
+            Defaults to None (no chunking).
+    """
+
+    NAME = "ClickHouse"
+
+    _CREATE_TABLE_TEMPLATE = """
+        CREATE TABLE IF NOT EXISTS {table_name} (
+            {columns}
+        )
+        ENGINE = {engine}
+        ORDER BY {order_by}
+        {additional_props}
+    """
+    _DROP_TABLE_TEMPLATE = """DROP TABLE IF EXISTS {table_name}"""
+    _CHECK_TABLE_EXISTS_TEMPLATE = """EXISTS {table_name}"""
+    _SHOW_CREATE_TABLE_TEMPLATE = """SHOW CREATE TABLE {table_name}"""
+
+    def __init__(
+        self,
+        table: str,
+        dsn: str,
+        mode: SinkMode = SinkMode.CREATE,
+        schema: Optional[pyarrow.Schema] = None,
+        client_settings: Optional[Dict[str, Any]] = None,
+        client_kwargs: Optional[Dict[str, Any]] = None,
+        table_settings: Optional[ClickHouseTableSettings] = None,
+        max_insert_block_rows: Optional[int] = None,
+    ) -> None:
+        self._table = table
+        self._dsn = dsn
+        self._mode = mode
+        self._schema = schema
+        self._client_settings = client_settings or {}
+        self._client_kwargs = client_kwargs or {}
+        self._table_settings = table_settings or ClickHouseTableSettings()
+        self._max_insert_block_rows = max_insert_block_rows
+        self._table_dropped = False
+
+    def _init_client(self):
+        _check_import(self, module="clickhouse_connect", package="clickhouse-connect")
+        import clickhouse_connect
+
+        return clickhouse_connect.get_client(
+            dsn=self._dsn,
+            settings=self._client_settings,
+            **self._client_kwargs,
+        )
+
+    def _generate_create_table_sql(
+        self,
+        schema: pyarrow.Schema,
+    ) -> str:
+        engine = self._table_settings.engine
+        if self._table_settings.order_by is not None:
+            order_by = self._table_settings.order_by
+        else:
+            order_by = _pick_best_arrow_field_for_order_by(schema)
+        additional_clauses = []
+        if self._table_settings.partition_by is not None:
+            additional_clauses.append(
+                f"PARTITION BY {self._table_settings.partition_by}"
+            )
+        if self._table_settings.primary_key is not None:
+            additional_clauses.append(
+                f"PRIMARY KEY ({self._table_settings.primary_key})"
+            )
+        if self._table_settings.settings is not None:
+            additional_clauses.append(f"SETTINGS {self._table_settings.settings}")
+        additional_props = ""
+        if additional_clauses:
+            additional_props = "\n" + "\n".join(additional_clauses)
+        columns_def = []
+        for field in schema:
+            ch_type = _arrow_to_clickhouse_type(field)
+            columns_def.append(f"`{field.name}` {ch_type}")
+        columns_str = ",\n    ".join(columns_def)
+        return self._CREATE_TABLE_TEMPLATE.format(
+            table_name=self._table,
+            columns=columns_str,
+            engine=engine,
+            order_by=order_by,
+            additional_props=additional_props,
+        )
+
+    def _table_exists(self, client) -> bool:
+        try:
+            result = client.query(
+                self._CHECK_TABLE_EXISTS_TEMPLATE.format(table_name=self._table)
+            )
+            if result is None:
+                return False
+            for helper in ("scalar", "first_item", "first_value"):
+                _check_import(
+                    self, module="clickhouse_connect", package="clickhouse-connect"
+                )
+                from clickhouse_connect.driver.exceptions import Error as CHError
+
+                if hasattr(result, helper):
+                    try:
+                        return bool(getattr(result, helper)())
+                    except (TypeError, ValueError, CHError) as exc:
+                        # Helper exists but failed – continue probing.
+                        logger.debug(
+                            "Helper %s failed: %s; will try fallbacks", helper, exc
+                        )
+            # Fallback to inspecting common container attributes.
+            for attr in ("result_rows", "rows", "data"):
+                rows = getattr(result, attr, None)
+                if rows:
+                    first = rows[0]
+                    # Unwrap an extra layer if present (i.e. [[1]] or [(1,)])
+                    if isinstance(first, (list, tuple)):
+                        first = first[0] if first else 0
+                    return bool(first)
+            return False
+        except Exception as e:
+            logger.warning(f"Could not verify if table {self._table} exists: {e}")
+            return False
+
+    def _get_existing_order_by(self, client) -> Optional[str]:
+        logger.debug(
+            f"Retrieving ORDER BY clause from SHOW CREATE TABLE for {self._table}"
+        )
+        try:
+            show_create_sql = self._SHOW_CREATE_TABLE_TEMPLATE.format(
+                table_name=self._table
+            )
+            result = client.command(show_create_sql)
+            ddl_str = str(result)
+            pattern = r"(?is)\border\s+by\s+(.*?)(?=\bengine\b|$)"
+            match = re.search(pattern, ddl_str)
+            if match:
+                return match.group(1).strip()
+            return None
+        except Exception as e:
+            logger.warning(
+                f"Could not retrieve SHOW CREATE TABLE for {self._table}: {e}"
+            )
+            return None
+
+    @property
+    def supports_distributed_writes(self) -> bool:
+        return True
+
+    def on_write_start(self) -> None:
+        client = None
+        try:
+            client = self._init_client()
+            table_exists = self._table_exists(client)
+            schema_required = (
+                # Overwrite always needs a schema because it recreates the table
+                self._mode == SinkMode.OVERWRITE
+                # For CREATE or APPEND we need a schema only when the table is
+                # absent and will therefore be created in this call.
+                or (
+                    self._mode in (SinkMode.CREATE, SinkMode.APPEND)
+                    and not table_exists
+                )
+            )
+            if schema_required and self._schema is None:
+                if self._mode == SinkMode.OVERWRITE:
+                    raise ValueError(
+                        f"Overwriting table {self._table} requires a user‑provided schema."
+                    )
+                else:
+                    raise ValueError(
+                        f"Table {self._table} does not exist in mode='{self._mode.name}' and "
+                        "no schema was provided. Cannot create the table without a schema."
+                    )
+            # OVERWRITE MODE
+            if self._mode == SinkMode.OVERWRITE:
+                # If we plan to overwrite, drop the table if it exists,
+                # then re-create it using the user-provided schema.
+                if table_exists and self._table_settings.order_by is None:
+                    # Collect existing ORDER BY. This lets us preserve it if the user
+                    # hasn't specified one explicitly.
+                    existing_order_by = self._get_existing_order_by(client)
+                    if existing_order_by is not None:
+                        self._table_settings.order_by = existing_order_by
+                        logger.info(
+                            f"Reusing old ORDER BY from overwritten table:"
+                            f" {existing_order_by}"
+                        )
+                # DROP and CREATE the table.
+                drop_sql = self._DROP_TABLE_TEMPLATE.format(table_name=self._table)
+                logger.info(f"Mode=OVERWRITE => {drop_sql}")
+                client.command(drop_sql)
+                self._table_dropped = True
+                create_sql = self._generate_create_table_sql(self._schema)
+                client.command(create_sql)
+            # CREATE MODE
+            elif self._mode == SinkMode.CREATE:
+                # If table already exists in CREATE mode, fail immediately.
+                if table_exists:
+                    msg = (
+                        f"Table {self._table} already exists in mode='CREATE'. "
+                        "Use mode='APPEND' or 'OVERWRITE' instead."
+                    )
+                    logger.error(msg)
+                    raise ValueError(msg)
+                # Otherwise, create it (requires user-provided schema).
+                create_sql = self._generate_create_table_sql(self._schema)
+                client.command(create_sql)
+            # APPEND MODE
+            elif self._mode == SinkMode.APPEND:
+                if table_exists:
+                    # Table exists – validate or adopt ORDER BY.
+                    existing_order_by = self._get_existing_order_by(client)
+                    user_order_by = self._table_settings.order_by
+                    if user_order_by is not None:
+                        # The user explicitly set an order_by. Check if it conflicts:
+                        if existing_order_by and existing_order_by != user_order_by:
+                            raise ValueError(
+                                f"Conflict with order_by. The existing table {self._table} "
+                                f"has ORDER BY {existing_order_by}, but the user specified "
+                                f"ORDER BY {user_order_by}. Please drop or overwrite the table, "
+                                f"or use the same ordering."
+                            )
+                    elif existing_order_by:
+                        self._table_settings.order_by = existing_order_by
+                        logger.info(
+                            f"Reusing existing ORDER BY for table {self._table}: {existing_order_by}"
+                        )
+                else:
+                    # Table doesn't exist, so create it with user schema
+                    create_sql = self._generate_create_table_sql(self._schema)
+                    client.command(create_sql)
+
+        except Exception as e:
+            logger.error(
+                f"Could not complete on_write_start for table {self._table}: {e}"
+            )
+            raise e
+        finally:
+            if client:
+                client.close()
+
+    def get_name(self) -> str:
+        return self.NAME
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> WriteReturnType:
+        client = self._init_client()
+        total_inserted = 0
+        try:
+            for block in blocks:
+                arrow_table = BlockAccessor.for_block(block).to_arrow()
+                row_count = arrow_table.num_rows
+                if self._schema is not None:
+                    arrow_table = arrow_table.cast(self._schema, safe=True)
+                if self._max_insert_block_rows is not None:
+                    max_chunk_size = self._max_insert_block_rows
+                else:
+                    # If max_insert_block_rows is not set, insert all rows in one go
+                    max_chunk_size = row_count if row_count > 0 else 1
+                offsets = list(range(0, row_count, max_chunk_size))
+                offsets.append(row_count)
+                for i in range(len(offsets) - 1):
+                    start = offsets[i]
+                    end = offsets[i + 1]
+                    chunk = arrow_table.slice(start, end - start)
+                    client.insert_arrow(self._table, chunk)
+                    total_inserted += chunk.num_rows
+        except Exception as e:
+            logger.error(f"Failed to write block(s) to table {self._table}: {e}")
+            raise e
+        finally:
+            client.close()
+        return [total_inserted]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/clickhouse_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/clickhouse_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..b82ec4bf7d0c9bb04971d88d9917bf2def30946a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/clickhouse_datasource.py
@@ -0,0 +1,355 @@
+import logging
+import math
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
+
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+def _is_filter_string_safe(filter_str: str) -> bool:
+    in_string = False
+    escape_next = False
+    for c in filter_str:
+        if in_string:
+            # If we're inside a string, check if we're closing it.
+            if c == "'" and not escape_next:
+                in_string = False
+            escape_next = (c == "\\") and not escape_next
+        else:
+            # If we're not in a string, entering one if we see a single quote
+            if c == "'":
+                in_string = True
+                escape_next = False
+            # Disallow semicolon if we're not in a string
+            elif c == ";":
+                return False
+            else:
+                escape_next = False
+    # If we end inside a string, it's suspicious, but let's allow
+    # it to be further validated by the DB. Just return True here.
+    return True
+
+
+@DeveloperAPI
+class ClickHouseDatasource(Datasource):
+    """
+    A Ray datasource for reading from ClickHouse.
+
+    Args:
+        table: Fully qualified table or view identifier (e.g.,
+            "default.table_name").
+        dsn: A string in DSN (Data Source Name) HTTP format (e.g.,
+            "clickhouse+http://username:password@host:8124/default").
+            For more information, see `ClickHouse Connection String doc
+            <https://clickhouse.com/docs/en/integrations/sql-clients/cli#connection_string>`_.
+        columns: Optional List of columns to select from the data source.
+            If no columns are specified, all columns will be selected by default.
+        filter: Optional SQL filter string that will be used in the
+            WHERE statement (e.g., "label = 2 AND text IS NOT NULL").
+            The filter must be valid for use in a ClickHouse SQL WHERE clause.
+            Note: Parallel reads are not currently supported when a filter is set.
+            Specifying a filter forces the parallelism to 1 to ensure deterministic
+            and consistent results. For more information, see
+            `ClickHouse SQL WHERE Clause doc
+            <https://clickhouse.com/docs/en/sql-reference/statements/select/where>`_.
+        order_by: Optional Tuple containing a list of columns to order by
+            and a boolean indicating the order. Note: order_by is required to
+            support parallelism.
+        client_settings: Optional ClickHouse server settings to be used with the
+            session/every request. For more information, see
+            `ClickHouse Client Settings doc
+            <https://clickhouse.com/docs/en/integrations/python#settings-argument>`_.
+        client_kwargs: Optional Additional keyword arguments to pass to the
+            ClickHouse client. For more information,
+            see `ClickHouse Core Settings doc
+            <https://clickhouse.com/docs/en/integrations/python#additional-options>`_.
+    """
+
+    NUM_SAMPLE_ROWS = 100
+    MIN_ROWS_PER_READ_TASK = 50
+    _BASE_QUERY = "SELECT {select_clause} FROM {table}"
+    _EXPLAIN_FILTERS_QUERY = "EXPLAIN SELECT 1 FROM {table} WHERE {filter_clause}"
+    _SIZE_ESTIMATE_QUERY = "SELECT SUM(byteSize(*)) AS estimate FROM ({query})"
+    _COUNT_ESTIMATE_QUERY = "SELECT COUNT(*) AS estimate FROM ({query})"
+    _SAMPLE_BLOCK_QUERY = "{query} LIMIT {limit_row_count}"
+    _FIRST_BLOCK_QUERY = """
+        {query}
+        FETCH FIRST {fetch_row_count} {fetch_row_or_rows} ONLY
+    """
+    _NEXT_BLOCK_QUERY = """
+        {query}
+        OFFSET {offset_row_count} {offset_row_or_rows}
+        FETCH NEXT {fetch_row_count} {fetch_row_or_rows} ONLY
+    """
+
+    def __init__(
+        self,
+        table: str,
+        dsn: str,
+        columns: Optional[List[str]] = None,
+        filter: Optional[str] = None,
+        order_by: Optional[Tuple[List[str], bool]] = None,
+        client_settings: Optional[Dict[str, Any]] = None,
+        client_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        self._table = table
+        self._dsn = dsn
+        self._columns = columns
+        self._filter = filter
+        self._order_by = order_by
+        self._client_settings = client_settings or {}
+        self._client_kwargs = client_kwargs or {}
+        self._query = self._generate_query()
+
+    def _init_client(self):
+        _check_import(self, module="clickhouse_connect", package="clickhouse-connect")
+        import clickhouse_connect
+
+        return clickhouse_connect.get_client(
+            dsn=self._dsn,
+            settings=self._client_settings or {},
+            **self._client_kwargs or {},
+        )
+
+    def _validate_filter(self):
+        if not self._filter:
+            return
+        # Minimal lexical check (regex or manual approach for semicolons, etc.).
+        if not _is_filter_string_safe(self._filter):
+            raise ValueError(
+                f"Invalid characters outside of "
+                f"string literals in filter: {self._filter}"
+            )
+        # Test "EXPLAIN" query to confirm parse-ability and catch expression errors.
+        client = self._init_client()
+        try:
+            test_query = self._EXPLAIN_FILTERS_QUERY.format(
+                table=self._table,
+                filter_clause=self._filter,
+            )
+            client.query(test_query)
+        except Exception as e:
+            raise ValueError(
+                f"Invalid filter expression: {self._filter}. Error: {e}",
+            )
+        finally:
+            client.close()
+
+    def _generate_query(self) -> str:
+        query = self._BASE_QUERY.format(
+            select_clause=", ".join(self._columns) if self._columns else "*",
+            table=self._table,
+        )
+        if self._filter:
+            self._validate_filter()
+            query += f" WHERE {self._filter}"
+        if self._order_by:
+            columns, desc = self._order_by
+            direction = " DESC" if desc else ""
+            if len(columns) == 1:
+                query += f" ORDER BY {columns[0]}{direction}"
+            elif len(columns) > 1:
+                columns_clause = ", ".join(columns)
+                query += f" ORDER BY ({columns_clause}){direction}"
+        return query
+
+    def _build_block_query(self, limit_row_count: int, offset_row_count: int) -> str:
+        if offset_row_count == 0:
+            # The first block query is optimized to use FETCH FIRST clause
+            # with an OFFSET specified.
+            return self._FIRST_BLOCK_QUERY.format(
+                query=self._query,
+                fetch_row_count=limit_row_count,
+                fetch_row_or_rows="ROWS" if limit_row_count > 1 else "ROW",
+            )
+        # Subsequent block queries use OFFSET and FETCH NEXT clauses to read the
+        # next block of data.
+        return self._NEXT_BLOCK_QUERY.format(
+            query=self._query,
+            offset_row_count=offset_row_count,
+            offset_row_or_rows="ROWS" if offset_row_count > 1 else "ROW",
+            fetch_row_count=limit_row_count,
+            fetch_row_or_rows="ROWS" if limit_row_count > 1 else "ROW",
+        )
+
+    def _create_read_fn(
+        self,
+        query: str,
+    ) -> Callable[[], Iterable[Block]]:
+        def read_fn() -> Iterable[Block]:
+            return [self._execute_block_query(query)]
+
+        return read_fn
+
+    def _get_sampled_estimates(self):
+        if self._order_by is not None:
+            # If the query is ordered, we can use a FETCH clause to get a sample.
+            # This reduces the CPU overhead on ClickHouse and speeds up the
+            # estimation query.
+            query = self._FIRST_BLOCK_QUERY.format(
+                query=self._query,
+                fetch_row_count=self.NUM_SAMPLE_ROWS,
+                fetch_row_or_rows="ROWS" if self.NUM_SAMPLE_ROWS > 1 else "ROW",
+            )
+        else:
+            # If the query is not ordered, we need to use a LIMIT clause to
+            # get a sample.
+            query = self._SAMPLE_BLOCK_QUERY.format(
+                query=self._query,
+                limit_row_count=self.NUM_SAMPLE_ROWS,
+            )
+        sample_block_accessor = BlockAccessor.for_block(
+            self._execute_block_query(query)
+        )
+        estimated_size_bytes_per_row = math.ceil(
+            sample_block_accessor.size_bytes() / sample_block_accessor.num_rows()
+        )
+        sample_block_schema = sample_block_accessor.schema()
+        return estimated_size_bytes_per_row, sample_block_schema
+
+    def _get_estimate_count(self) -> Optional[int]:
+        return self._execute_estimate_query(self._COUNT_ESTIMATE_QUERY)
+
+    def _get_estimate_size(self) -> Optional[int]:
+        return self._execute_estimate_query(self._SIZE_ESTIMATE_QUERY)
+
+    def _execute_estimate_query(self, estimate_query: str) -> Optional[int]:
+        client = self._init_client()
+        try:
+            # Estimate queries wrap around the primary query, self._query.
+            # This allows us to use self._query as a sub-query to efficiently
+            # and accurately estimate the size or count of the result set.
+            query = estimate_query.format(query=self._query)
+            result = client.query(query)
+            if result and len(result.result_rows) > 0:
+                estimate = result.result_rows[0][0]
+                return int(estimate) if estimate is not None else None
+        except Exception as e:
+            logger.warning(f"Failed to execute estimate query: {e}")
+        finally:
+            client.close()
+        return None
+
+    def _execute_block_query(self, query: str) -> Block:
+        import pyarrow as pa
+
+        client = self._init_client()
+        try:
+            with client.query_arrow_stream(query) as stream:
+                record_batches = list(stream)  # Collect all record batches
+            return pa.Table.from_batches(record_batches)
+        except Exception as e:
+            raise RuntimeError(f"Failed to execute block query: {e}")
+        finally:
+            client.close()
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        """
+        Estimate the in-memory data size for the query.
+
+        Returns:
+            Estimated in-memory data size in bytes, or
+            None if the estimation cannot be performed.
+        """
+        return self._get_estimate_size()
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        """
+        Create read tasks for the ClickHouse query.
+
+        Args:
+            parallelism: The desired number of partitions to read the data into.
+                - If ``order_by`` is not set, parallelism will be forced to 1.
+                - If ``filter`` is set, parallelism will also be forced to 1
+                    to ensure deterministic results.
+            per_task_row_limit: Maximum number of rows allowed in each emitted
+                task.  Blocks larger than this limit will be sliced before
+                being yielded downstream.
+
+        Returns:
+            A list of read tasks to be executed.
+        """
+        num_rows_total = self._get_estimate_count()
+        if num_rows_total == 0 or num_rows_total is None:
+            return []
+        parallelism = min(
+            parallelism, math.ceil(num_rows_total / self.MIN_ROWS_PER_READ_TASK)
+        )
+        # To ensure consistent order of query results, self._order_by
+        # must be specified and self.filter must not be specified
+        # in order to support parallelism.
+        if self._filter is not None and parallelism > 1:
+            logger.warning(
+                "ClickHouse datasource does not currently support parallel reads "
+                "when a filter is set; falling back to parallelism of 1."
+            )
+            # When filter is specified and parallelism is greater than 1,
+            # we need to reduce parallelism to 1 to ensure consistent results.
+            parallelism = 1
+        # To ensure consistent order of query results, self._order_by
+        # must be specified in order to support parallelism.
+        if self._order_by is None and parallelism > 1:
+            logger.warning(
+                "ClickHouse datasource requires dataset to be explicitly ordered "
+                "to support parallelism; falling back to parallelism of 1."
+            )
+            # When order_by is not specified and parallelism is greater than 1,
+            # we need to reduce parallelism to 1 to ensure consistent results.
+            parallelism = 1
+        # By reducing parallelism to 1 when either of the conditions above are met,
+        # we ensure the downstream process is treated exactly as a non-parallelized
+        # (single block) process would be, thus ensuring output consistency.
+        num_rows_per_block = num_rows_total // parallelism
+        num_blocks_with_extra_row = num_rows_total % parallelism
+        (
+            estimated_size_bytes_per_row,
+            sample_block_schema,
+        ) = self._get_sampled_estimates()
+
+        def _get_read_task(
+            block_rows: int, offset_rows: int, parallelized: bool
+        ) -> ReadTask:
+            if parallelized:
+                # When parallelized, we need to build a block query with OFFSET
+                # and FETCH clauses.
+                query = self._build_block_query(block_rows, offset_rows)
+            else:
+                # When not parallelized, we can use the original query without
+                # OFFSET and FETCH clauses.
+                query = self._query
+            return ReadTask(
+                self._create_read_fn(query),
+                BlockMetadata(
+                    num_rows=block_rows,
+                    size_bytes=estimated_size_bytes_per_row * block_rows,
+                    input_files=None,
+                    exec_stats=None,
+                ),
+                schema=sample_block_schema,
+                per_task_row_limit=per_task_row_limit,
+            )
+
+        if parallelism == 1:
+            # When parallelism is 1, we can read the entire dataset in a single task.
+            # We then optimize this scenario by using self._query directly without
+            # unnecessary OFFSET and FETCH clauses.
+            return [_get_read_task(num_rows_total, 0, False)]
+
+        # Otherwise we need to split the dataset into multiple tasks.
+        # Each task will include OFFSET and FETCH clauses to efficiently
+        # read a subset of the dataset.
+        read_tasks = []
+        offset = 0
+        for i in range(parallelism):
+            this_block_size = num_rows_per_block
+            if i < num_blocks_with_extra_row:
+                this_block_size += 1
+            read_tasks.append(_get_read_task(this_block_size, offset, True))
+            offset += this_block_size
+        return read_tasks
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/csv_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/csv_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..565abf85fe2a4d15ad16e530dba3b8eec9c643d0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/csv_datasink.py
@@ -0,0 +1,36 @@
+from typing import Any, Callable, Dict, Optional
+
+import pyarrow
+
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_based_datasource import _resolve_kwargs
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class CSVDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        arrow_csv_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        arrow_csv_args: Optional[Dict[str, Any]] = None,
+        file_format="csv",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        if arrow_csv_args_fn is None:
+            arrow_csv_args_fn = lambda: {}  # noqa: E731
+
+        if arrow_csv_args is None:
+            arrow_csv_args = {}
+
+        self.arrow_csv_args_fn = arrow_csv_args_fn
+        self.arrow_csv_args = arrow_csv_args
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        from pyarrow import csv
+
+        writer_args = _resolve_kwargs(self.arrow_csv_args_fn, **self.arrow_csv_args)
+        write_options = writer_args.pop("write_options", None)
+        csv.write_csv(block.to_arrow(), file, write_options, **writer_args)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/csv_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/csv_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b280279177da177776ddee28cf0e96af84010d9
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/csv_datasource.py
@@ -0,0 +1,82 @@
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class CSVDatasource(FileBasedDatasource):
+    """CSV datasource, for reading and writing CSV files."""
+
+    _FILE_EXTENSIONS = [
+        "csv",
+        "csv.gz",  # gzip-compressed files
+        "csv.br",  # Brotli-compressed files
+        "csv.zst",  # Zstandard-compressed files
+        "csv.lz4",  # lz4-compressed files
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        arrow_csv_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        from pyarrow import csv
+
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if arrow_csv_args is None:
+            arrow_csv_args = {}
+
+        self.read_options = arrow_csv_args.pop(
+            "read_options", csv.ReadOptions(use_threads=False)
+        )
+        self.parse_options = arrow_csv_args.pop("parse_options", csv.ParseOptions())
+        self.arrow_csv_args = arrow_csv_args
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import pyarrow as pa
+        from pyarrow import csv
+
+        # Re-init invalid row handler: https://issues.apache.org/jira/browse/ARROW-17641
+        if hasattr(self.parse_options, "invalid_row_handler"):
+            self.parse_options.invalid_row_handler = (
+                self.parse_options.invalid_row_handler
+            )
+
+        filter_expr = (
+            self._predicate_expr.to_pyarrow()
+            if self._predicate_expr is not None
+            else None
+        )
+
+        try:
+            reader = csv.open_csv(
+                f,
+                read_options=self.read_options,
+                parse_options=self.parse_options,
+                **self.arrow_csv_args,
+            )
+            schema = None
+            while True:
+                try:
+                    batch = reader.read_next_batch()
+                    table = pa.Table.from_batches([batch], schema=schema)
+                    if schema is None:
+                        schema = table.schema
+                    if filter_expr is not None:
+                        table = table.filter(filter_expr)
+
+                    yield table
+                except StopIteration:
+                    return
+        except pa.lib.ArrowInvalid as e:
+            raise ValueError(
+                f"Failed to read CSV file: {path}. "
+                "Please check the CSV file has correct format, or filter out non-CSV "
+                "file with 'partition_filter' field. See read_csv() documentation for "
+                "more details."
+            ) from e
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/databricks_uc_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/databricks_uc_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a5232bd21a37838fc11be9a2fb781bce4b286f3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/databricks_uc_datasource.py
@@ -0,0 +1,222 @@
+import json
+import logging
+import os
+import time
+from typing import List, Optional
+from urllib.parse import urljoin
+
+import numpy as np
+import pyarrow
+import requests
+
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+_STATEMENT_EXEC_POLL_TIME_S = 1
+
+
+@PublicAPI(stability="alpha")
+class DatabricksUCDatasource(Datasource):
+    def __init__(
+        self,
+        host: str,
+        token: str,
+        warehouse_id: str,
+        catalog: str,
+        schema: str,
+        query: str,
+    ):
+        self.host = host
+        self.token = token
+        self.warehouse_id = warehouse_id
+        self.catalog = catalog
+        self.schema = schema
+        self.query = query
+
+        if not host.startswith(("http://", "https://")):
+            self.host = f"https://{host}"
+
+        url_base = f"{self.host}/api/2.0/sql/statements/"
+
+        payload = json.dumps(
+            {
+                "statement": self.query,
+                "warehouse_id": self.warehouse_id,
+                "wait_timeout": "0s",
+                "disposition": "EXTERNAL_LINKS",
+                "format": "ARROW_STREAM",
+                "catalog": self.catalog,
+                "schema": self.schema,
+            }
+        )
+
+        req_headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer " + self.token,
+        }
+
+        response = requests.post(
+            url_base,
+            headers=req_headers,
+            data=payload,
+        )
+        response.raise_for_status()
+        statement_id = response.json()["statement_id"]
+
+        state = response.json()["status"]["state"]
+
+        logger.info(f"Waiting for query {query!r} execution result.")
+        try:
+            while state in ["PENDING", "RUNNING"]:
+                time.sleep(_STATEMENT_EXEC_POLL_TIME_S)
+                response = requests.get(
+                    urljoin(url_base, statement_id) + "/",
+                    headers=req_headers,
+                )
+                response.raise_for_status()
+                state = response.json()["status"]["state"]
+        except KeyboardInterrupt:
+            # User cancel the command, so we cancel query execution.
+            requests.post(
+                urljoin(url_base, f"{statement_id}/cancel"),
+                headers=req_headers,
+            )
+            try:
+                response.raise_for_status()
+            except Exception as e:
+                logger.warning(
+                    f"Canceling query {query!r} execution failed, reason: {repr(e)}."
+                )
+            raise
+
+        if state != "SUCCEEDED":
+            raise RuntimeError(
+                f"Query {self.query!r} execution failed.\n{response.json()}"
+            )
+
+        manifest = response.json()["manifest"]
+        self.is_truncated = manifest.get("truncated", False)
+
+        if self.is_truncated:
+            logger.warning(
+                f"The resulting size of the dataset of '{query!r}' exceeds "
+                "100GiB and it is truncated."
+            )
+
+        chunks = manifest.get("chunks", [])
+
+        # Make chunks metadata are ordered by index.
+        chunks = sorted(chunks, key=lambda x: x["chunk_index"])
+        num_chunks = len(chunks)
+        self.num_chunks = num_chunks
+        self._estimate_inmemory_data_size = sum(chunk["byte_count"] for chunk in chunks)
+
+        def get_read_task(
+            task_index: int, parallelism: int, per_task_row_limit: Optional[int] = None
+        ):
+            # Handle empty chunk list by yielding an empty PyArrow table
+            if num_chunks == 0:
+                import pyarrow as pa
+
+                metadata = BlockMetadata(
+                    num_rows=0,
+                    size_bytes=0,
+                    input_files=None,
+                    exec_stats=None,
+                )
+
+                def empty_read_fn():
+                    yield pa.Table.from_pydict({})
+
+                return ReadTask(read_fn=empty_read_fn, metadata=metadata)
+
+            # get chunk list to be read in this task and preserve original chunk order
+            chunk_index_list = list(
+                np.array_split(range(num_chunks), parallelism)[task_index]
+            )
+
+            num_rows = sum(
+                chunks[chunk_index]["row_count"] for chunk_index in chunk_index_list
+            )
+            size_bytes = sum(
+                chunks[chunk_index]["byte_count"] for chunk_index in chunk_index_list
+            )
+
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=size_bytes,
+                input_files=None,
+                exec_stats=None,
+            )
+
+            def _read_fn():
+                for chunk_index in chunk_index_list:
+                    resolve_external_link_url = urljoin(
+                        url_base, f"{statement_id}/result/chunks/{chunk_index}"
+                    )
+
+                    resolve_response = requests.get(
+                        resolve_external_link_url, headers=req_headers
+                    )
+                    resolve_response.raise_for_status()
+                    external_url = resolve_response.json()["external_links"][0][
+                        "external_link"
+                    ]
+                    # NOTE: do _NOT_ send the authorization header to external urls
+                    raw_response = requests.get(external_url, auth=None, headers=None)
+                    raw_response.raise_for_status()
+
+                    with pyarrow.ipc.open_stream(raw_response.content) as reader:
+                        arrow_table = reader.read_all()
+
+                    yield arrow_table
+
+            def read_fn():
+                if mock_setup_fn_path := os.environ.get(
+                    "RAY_DATABRICKS_UC_DATASOURCE_READ_FN_MOCK_TEST_SETUP_FN_PATH"
+                ):
+                    import ray.cloudpickle as pickle
+
+                    # This is for testing.
+                    with open(mock_setup_fn_path, "rb") as f:
+                        mock_setup = pickle.load(f)
+                    with mock_setup():
+                        yield from _read_fn()
+                else:
+                    yield from _read_fn()
+
+            return ReadTask(
+                read_fn=read_fn,
+                metadata=metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+
+        self._get_read_task = get_read_task
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._estimate_inmemory_data_size
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        # Handle empty dataset case
+        if self.num_chunks == 0:
+            return [self._get_read_task(0, 1, per_task_row_limit)]
+
+        assert parallelism > 0, f"Invalid parallelism {parallelism}"
+
+        if parallelism > self.num_chunks:
+            parallelism = self.num_chunks
+            logger.info(
+                "The parallelism is reduced to chunk number due to "
+                "insufficient chunk parallelism."
+            )
+
+        return [
+            self._get_read_task(index, parallelism, per_task_row_limit)
+            for index in range(parallelism)
+        ]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/delta_sharing_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/delta_sharing_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a6914e5606dcff1afdc8d74754e0141a617a587
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/delta_sharing_datasource.py
@@ -0,0 +1,132 @@
+import logging
+from json import loads
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+logger = logging.getLogger(__name__)
+
+
+class DeltaSharingDatasource(Datasource):
+    def __init__(
+        self,
+        url: str,
+        json_predicate_hints: Optional[str] = None,
+        limit: Optional[int] = None,
+        version: Optional[int] = None,
+        timestamp: Optional[str] = None,
+    ):
+        _check_import(self, module="delta_sharing", package="delta-sharing")
+
+        if limit is not None:
+            assert (
+                isinstance(limit, int) and limit >= 0
+            ), "'limit' must be a non-negative int"
+
+        self._url = url
+        self._json_predicate_hints = json_predicate_hints
+        self._limit = limit
+        self._version = version
+        self._timestamp = timestamp
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def _read_files(self, files, converters):
+        """Read files with Delta Sharing."""
+        from delta_sharing.reader import DeltaSharingReader
+
+        for file in files:
+            yield DeltaSharingReader._to_pandas(
+                action=file, converters=converters, for_cdf=False, limit=None
+            )
+
+    def setup_delta_sharing_connections(self, url: str):
+        """
+        Set up delta sharing connections based on the url.
+
+        Args:
+            url: A URL under the format "<profile>#<share>.<schema>.<table>"
+
+        Returns:
+            A tuple of (table, rest_client) where table is a delta_sharing Table
+            object and rest_client is a DataSharingRestClient instance.
+        """
+        from delta_sharing.protocol import DeltaSharingProfile, Table
+        from delta_sharing.rest_client import DataSharingRestClient
+
+        profile_str, share, schema, table_str = _parse_delta_sharing_url(url)
+        table = Table(name=table_str, share=share, schema=schema)
+
+        profile = DeltaSharingProfile.read_from_file(profile_str)
+        rest_client = DataSharingRestClient(profile)
+        return table, rest_client
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        assert parallelism > 0, f"Invalid parallelism {parallelism}"
+        from delta_sharing.converter import to_converters
+
+        self._table, self._rest_client = self.setup_delta_sharing_connections(self._url)
+        self._response = self._rest_client.list_files_in_table(
+            self._table,
+            jsonPredicateHints=self._json_predicate_hints,
+            limitHint=self._limit,
+            version=self._version,
+            timestamp=self._timestamp,
+        )
+
+        if len(self._response.add_files) == 0 or self._limit == 0:
+            logger.warning("No files found from the delta sharing table or limit is 0")
+
+        schema_json = loads(self._response.metadata.schema_string)
+        self._converters = to_converters(schema_json)
+
+        read_tasks = []
+        # get file list to be read in this task and preserve original chunk order
+        for files in np.array_split(self._response.add_files, parallelism):
+            files = files.tolist()
+            metadata = BlockMetadata(
+                num_rows=None,
+                input_files=files,
+                size_bytes=None,
+                exec_stats=None,
+            )
+            converters = self._converters
+            read_task = ReadTask(
+                lambda f=files: self._read_files(f, converters),
+                metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+
+def _parse_delta_sharing_url(url: str) -> Tuple[str, str, str, str]:
+    """
+    Developed from delta_sharing's _parse_url function.
+    https://github.com/delta-io/delta-sharing/blob/main/python/delta_sharing/delta_sharing.py#L36
+
+    Args:
+        url: a url under the format "<profile>#<share>.<schema>.<table>"
+
+    Returns:
+        a tuple with parsed (profile, share, schema, table)
+    """
+    shape_index = url.rfind("#")
+    if shape_index < 0:
+        raise ValueError(f"Invalid 'url': {url}")
+    profile = url[0:shape_index]
+    fragments = url[shape_index + 1 :].split(".")
+    if len(fragments) != 3:
+        raise ValueError(f"Invalid 'url': {url}")
+    share, schema, table = fragments
+    if len(profile) == 0 or len(share) == 0 or len(schema) == 0 or len(table) == 0:
+        raise ValueError(f"Invalid 'url': {url}")
+    return (profile, share, schema, table)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/hudi_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/hudi_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..800c747798e495796b5f809683d299c8fe7ca178
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/hudi_datasource.py
@@ -0,0 +1,148 @@
+import logging
+import os
+from enum import Enum
+from typing import Dict, Iterator, List, Optional, Tuple
+
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+logger = logging.getLogger(__name__)
+
+
+class HudiQueryType(Enum):
+    SNAPSHOT = "snapshot"
+    INCREMENTAL = "incremental"
+
+    @classmethod
+    def supported_types(cls) -> List[str]:
+        return [e.value for e in cls]
+
+
+class HudiDatasource(Datasource):
+    """Hudi datasource, for reading Apache Hudi table."""
+
+    def __init__(
+        self,
+        table_uri: str,
+        query_type: str,
+        filters: Optional[List[Tuple[str, str, str]]] = None,
+        hudi_options: Optional[Dict[str, str]] = None,
+        storage_options: Optional[Dict[str, str]] = None,
+    ):
+        _check_import(self, module="hudi", package="hudi-python")
+
+        self._table_uri = table_uri
+        self._query_type = HudiQueryType(query_type.lower())
+        self._filters = filters or []
+        self._hudi_options = hudi_options or {}
+        self._storage_options = storage_options or {}
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List["ReadTask"]:
+        import numpy as np
+        import pyarrow
+        from hudi import HudiTableBuilder
+
+        def _perform_read(
+            table_uri: str,
+            base_file_paths: List[str],
+            options: Dict[str, str],
+        ) -> Iterator["pyarrow.Table"]:
+            from hudi import HudiFileGroupReader
+
+            for p in base_file_paths:
+                file_group_reader = HudiFileGroupReader(table_uri, options)
+                batch = file_group_reader.read_file_slice_by_base_file_path(p)
+                yield pyarrow.Table.from_batches([batch])
+
+        hudi_table = (
+            HudiTableBuilder.from_base_uri(self._table_uri)
+            .with_hudi_options(self._hudi_options)
+            .with_storage_options(self._storage_options)
+            # Although hudi-rs supports MOR snapshot, we need to add an API in
+            # the next release to allow file group reader to take in a list of
+            # files. Hence, setting this config for now to restrain reading
+            # only on parquet files (read optimized mode).
+            # This won't affect reading COW.
+            .with_hudi_option("hoodie.read.use.read_optimized.mode", "true")
+            .build()
+        )
+
+        logger.info("Collecting file slices for Hudi table at: %s", self._table_uri)
+
+        if self._query_type == HudiQueryType.SNAPSHOT:
+            file_slices_splits = hudi_table.get_file_slices_splits(
+                parallelism, self._filters
+            )
+        elif self._query_type == HudiQueryType.INCREMENTAL:
+            start_ts = self._hudi_options.get("hoodie.read.file_group.start_timestamp")
+            end_ts = self._hudi_options.get("hoodie.read.file_group.end_timestamp")
+            # TODO(xushiyan): add table API to return splits of file slices
+            file_slices = hudi_table.get_file_slices_between(start_ts, end_ts)
+            file_slices_splits = np.array_split(file_slices, parallelism)
+        else:
+            raise ValueError(
+                f"Unsupported query type: {self._query_type}. Supported types are: {HudiQueryType.supported_types()}."
+            )
+
+        logger.info("Creating read tasks for Hudi table at: %s", self._table_uri)
+
+        reader_options = {
+            **hudi_table.storage_options(),
+            **hudi_table.hudi_options(),
+        }
+
+        schema = hudi_table.get_schema()
+        read_tasks = []
+
+        for file_slices_split in file_slices_splits:
+            num_rows = 0
+            relative_paths = []
+            input_files = []
+            size_bytes = 0
+            for file_slice in file_slices_split:
+                # A file slice in a Hudi table is a logical group of data files
+                # within a physical partition. Records stored in a file slice
+                # are associated with a commit on the Hudi table's timeline.
+                # For more info, see https://hudi.apache.org/docs/file_layouts
+                num_rows += file_slice.num_records
+                relative_path = file_slice.base_file_relative_path()
+                relative_paths.append(relative_path)
+                full_path = os.path.join(self._table_uri, relative_path)
+                input_files.append(full_path)
+                size_bytes += file_slice.base_file_size
+
+            if self._query_type == HudiQueryType.SNAPSHOT:
+                metadata = BlockMetadata(
+                    num_rows=num_rows,
+                    input_files=input_files,
+                    size_bytes=size_bytes,
+                    exec_stats=None,
+                )
+            elif self._query_type == HudiQueryType.INCREMENTAL:
+                # need the check due to
+                # https://github.com/apache/hudi-rs/issues/401
+                metadata = BlockMetadata(
+                    num_rows=None,
+                    input_files=input_files,
+                    size_bytes=None,
+                    exec_stats=None,
+                )
+
+            read_task = ReadTask(
+                read_fn=lambda paths=relative_paths: _perform_read(
+                    self._table_uri, paths, reader_options
+                ),
+                metadata=metadata,
+                schema=schema,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO(xushiyan) add APIs to provide estimated in-memory size
+        return None
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/huggingface_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/huggingface_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..54ebaea3f9637b563b7740aec70512a6017ea920
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/huggingface_datasource.py
@@ -0,0 +1,191 @@
+import sys
+from typing import TYPE_CHECKING, Iterable, List, Optional, Union
+
+from ray.air.util.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.dataset import Dataset
+from ray.data.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import datasets
+
+
+TRANSFORMERS_IMPORT_ERROR: Optional[ImportError] = None
+
+try:
+    # Due to HF Dataset's dynamic module system, we need to dynamically import the
+    # datasets_modules module on every actor when training.
+    # We accomplish this by simply running the following bit of code directly
+    # in the module you are currently viewing. This ensures that when we
+    # unpickle the Dataset, it runs before pickle tries to
+    # import datasets_modules and prevents an exception from being thrown.
+    # Same logic is present inside HF Transformers Ray
+    # integration: https://github.com/huggingface/transformers/blob/\
+    # 7d5fde991d598370d961be8cb7add6541e2b59ce/src/transformers/integrations.py#L271
+    # Also see https://github.com/ray-project/ray/issues/28084
+    from transformers.utils import is_datasets_available
+
+    if "datasets_modules" not in sys.modules and is_datasets_available():
+        import importlib
+        import importlib.metadata
+        import os
+
+        import datasets.load
+        from packaging.version import parse
+
+        # Datasets >= 4.0 removed dataset scripts support and the dynamic-modules cache.
+        # Only initialize dynamic modules on <= 3.x where the initializer `init_dynamic_modules` exists.
+        DATASETS_VERSION = parse(importlib.metadata.version("datasets"))
+        DATASETS_VERSION_WITHOUT_SCRIPT_SUPPORT = parse("4.0.0")
+
+        if DATASETS_VERSION < DATASETS_VERSION_WITHOUT_SCRIPT_SUPPORT:
+            dynamic_modules_path = os.path.join(
+                datasets.load.init_dynamic_modules(), "__init__.py"
+            )
+            # load dynamic_modules from path
+            spec = importlib.util.spec_from_file_location(
+                "datasets_modules", dynamic_modules_path
+            )
+            datasets_modules = importlib.util.module_from_spec(spec)
+            sys.modules[spec.name] = datasets_modules
+            spec.loader.exec_module(datasets_modules)
+except ImportError as e:
+    TRANSFORMERS_IMPORT_ERROR = e
+
+
+class HuggingFaceDatasource(Datasource):
+    """Hugging Face Dataset datasource, for reading from a
+    `Hugging Face Datasets Dataset <https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.Dataset/>`_.
+    This Datasource implements a streamed read using a
+    single read task, most beneficial for a
+    `Hugging Face Datasets IterableDataset <https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.IterableDataset/>`_
+    or datasets which are too large to fit in-memory.
+    For an in-memory Hugging Face Dataset (`datasets.Dataset`), use :meth:`~ray.data.from_huggingface`
+    directly for faster performance.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        dataset: Union["datasets.Dataset", "datasets.IterableDataset"],
+        batch_size: int = 4096,
+    ):
+        if TRANSFORMERS_IMPORT_ERROR is not None:
+            raise TRANSFORMERS_IMPORT_ERROR
+
+        self._dataset = dataset
+        self._batch_size = batch_size
+
+    @classmethod
+    def list_parquet_urls_from_dataset(
+        cls, dataset: Union["datasets.Dataset", "datasets.IterableDataset"]
+    ) -> Dataset:
+        """Return list of Hugging Face hosted parquet file URLs if they
+        exist for the data (i.e. if the dataset is a public dataset that
+        has not been transformed) else return an empty list."""
+        import datasets
+
+        # We can use the dataset name, config name, and split name to load
+        # public hugging face datasets from the Hugging Face Hub. More info
+        # here: https://huggingface.co/docs/datasets-server/parquet
+        dataset_name = dataset.info.dataset_name
+        config_name = dataset.info.config_name
+        split_name = str(dataset.split)
+
+        # If a dataset is not an iterable dataset, we will check if the
+        # dataset with the matching dataset name, config name, and split name
+        # on the Hugging Face Hub has the same fingerprint as the
+        # dataset passed into this function. If it is not matching, transforms
+        # or other operations have been performed so we cannot use the parquet
+        # files on the Hugging Face Hub, so we return an empty list.
+        if not isinstance(dataset, datasets.IterableDataset):
+            from datasets import load_dataset
+
+            try:
+                ds = load_dataset(dataset_name, config_name, split=split_name)
+                if ds._fingerprint != dataset._fingerprint:
+                    return []
+            except Exception:
+                # If an exception is thrown when trying to reload the dataset
+                # we should exit gracefully by returning an empty list.
+                return []
+
+        import requests
+
+        public_url = (
+            f"https://huggingface.co/api/datasets/{dataset_name}"
+            f"/parquet/{config_name}/{split_name}"
+        )
+        resp = requests.get(public_url)
+        if resp.status_code == requests.codes["ok"]:
+            # dataset corresponds to a public dataset, return list of parquet_files
+            return resp.json()
+        else:
+            return []
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._dataset.dataset_size
+
+    def _read_dataset(self) -> Iterable[Block]:
+        # Note: This is a method instead of a higher level function because
+        # we need to capture `self`. This will trigger the try-import logic at
+        # the top of file to avoid import error of dataset_modules.
+        import numpy as np
+        import pandas as pd
+        import pyarrow
+
+        for batch in self._dataset.with_format("arrow").iter(
+            batch_size=self._batch_size
+        ):
+            # HuggingFace IterableDatasets do not fully support methods like
+            # `set_format`, `with_format`, and `formatted_as`, so the dataset
+            # can return whatever is the default configured batch type, even if
+            # the format is manually overriden before iterating above.
+            # Therefore, we limit support to batch formats which have native
+            # block types in Ray Data (pyarrow.Table, pd.DataFrame),
+            # or can easily be converted to such (dict, np.array).
+            # See: https://github.com/huggingface/datasets/issues/3444
+            if not isinstance(batch, (pyarrow.Table, pd.DataFrame, dict, np.array)):
+                raise ValueError(
+                    f"Batch format {type(batch)} isn't supported. Only the "
+                    f"following batch formats are supported: "
+                    f"dict (corresponds to `None` in `dataset.with_format()`), "
+                    f"pyarrow.Table, np.array, pd.DataFrame."
+                )
+            # Ensure np.arrays are wrapped in a dict
+            # (subsequently converted to a pyarrow.Table).
+            if isinstance(batch, np.ndarray):
+                batch = {"item": batch}
+            if isinstance(batch, dict):
+                batch = pyarrow_table_from_pydict(batch)
+            # Ensure that we return the default block type.
+            block = BlockAccessor.for_block(batch).to_default()
+            yield block
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+    ) -> List[ReadTask]:
+        # Note: `parallelism` arg is currently not used by HuggingFaceDatasource.
+        # We always generate a single ReadTask to perform the read.
+        _check_pyarrow_version()
+
+        # TODO(scottjlee): IterableDataset doesn't provide APIs
+        # for getting number of rows, byte size, etc., so the
+        # BlockMetadata is currently empty. Properly retrieve
+        # or calculate these so that progress bars have meaning.
+        meta = BlockMetadata(
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        read_tasks: List[ReadTask] = [
+            ReadTask(
+                self._read_dataset,
+                meta,
+                per_task_row_limit=per_task_row_limit,
+            )
+        ]
+        return read_tasks
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/iceberg_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/iceberg_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b1b46fb48cb651d411aeba345a0a57eabc2d40a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/iceberg_datasink.py
@@ -0,0 +1,470 @@
+"""
+Module to write a Ray Dataset into an iceberg table, by using the Ray Datasink API.
+"""
+import logging
+import uuid
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
+
+import ray
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.savemode import SaveMode
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink, WriteResult
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.manifest import DataFile
+    from pyiceberg.table import Table, Transaction
+
+    from ray.data.expressions import Expr
+
+
+logger = logging.getLogger(__name__)
+
+
+def _append_and_commit(
+    txn: "Transaction",
+    data_files: List["DataFile"],
+    snapshot_properties: Dict[str, str],
+) -> None:
+    """Helper to append data files to a transaction and commit.
+
+    Args:
+        txn: PyIceberg transaction object
+        data_files: List of DataFile objects to append
+        snapshot_properties: Custom properties to write to snapshot summary
+    """
+    with txn._append_snapshot_producer(snapshot_properties) as append_files:
+        for data_file in data_files:
+            append_files.append_data_file(data_file)
+    txn.commit_transaction()
+
+
+@ray.remote(num_cpus=1)
+def _commit_upsert_task(
+    table_identifier: str,
+    catalog_name: str,
+    catalog_kwargs: Dict[str, Any],
+    data_files: List["DataFile"],
+    join_keys_dicts: List[Dict[str, List[Any]]],
+    join_cols: List[str],
+    snapshot_properties: Dict[str, str],
+) -> None:
+    """
+    Ray task to commit upsert transaction with copy-on-write strategy.
+
+    This runs as a Ray task to avoid OOM on the driver when:
+    - There are many join keys to merge
+    - The transaction commit is memory or time intensive
+
+    Args:
+        table_identifier: Iceberg table identifier
+        catalog_name: Catalog name
+        catalog_kwargs: Catalog initialization kwargs
+        data_files: List of DataFile objects to commit
+        join_keys_dicts: List of dictionaries mapping column names to lists of key values
+        join_cols: List of join column names
+        snapshot_properties: Custom properties to write to snapshot summary
+    """
+    import functools
+    from collections import defaultdict
+
+    import pyarrow as pa
+    from pyiceberg import catalog
+    from pyiceberg.table.upsert_util import create_match_filter
+
+    # Load the table
+    cat = catalog.load_catalog(catalog_name, **catalog_kwargs)
+    table = cat.load_table(table_identifier)
+
+    # Merge all join keys dictionaries
+    merged_keys_dict = defaultdict(list)
+    for join_keys_dict in join_keys_dicts:
+        for col, values in join_keys_dict.items():
+            merged_keys_dict[col].extend(values)
+
+    # Start transaction
+    txn = table.transaction()
+
+    # Create delete filter if we have join keys
+    if merged_keys_dict:
+        # Create PyArrow table from join keys
+        keys_table = pa.table(dict(merged_keys_dict))
+
+        # Filter out rows with any NULL values in join columns
+        # (NULL != NULL in SQL semantics)
+        masks = (pa.compute.is_valid(keys_table[col]) for col in join_cols)
+        mask = functools.reduce(pa.compute.and_, masks)
+        keys_table = keys_table.filter(mask)
+
+        # Only delete if we have non-NULL keys
+        if len(keys_table) > 0:
+            # Use PyIceberg's helper to build delete filter
+            delete_filter = create_match_filter(keys_table, join_cols)
+            txn.delete(
+                delete_filter=delete_filter,
+                snapshot_properties=snapshot_properties,
+            )
+
+    # Append new data files (includes updates and inserts) and commit
+    _append_and_commit(txn, data_files, snapshot_properties)
+
+
+@DeveloperAPI
+class IcebergDatasink(
+    Datasink[Union[List["DataFile"], tuple[List["DataFile"], Dict[str, List[Any]]]]]
+):
+    """
+    Iceberg datasink to write a Ray Dataset into an existing Iceberg table.
+    This datasink handles concurrent writes by:
+    - Each worker writes Parquet files to storage and returns DataFile metadata
+    - The driver collects all DataFile objects and performs a single commit
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+        snapshot_properties: Optional[Dict[str, str]] = None,
+        mode: SaveMode = SaveMode.APPEND,
+        overwrite_filter: Optional["Expr"] = None,
+        upsert_kwargs: Optional[Dict[str, Any]] = None,
+        overwrite_kwargs: Optional[Dict[str, Any]] = None,
+        upsert_commit_memory: Optional[int] = None,
+    ):
+        """
+        Initialize the IcebergDatasink
+
+        Args:
+            table_identifier: The identifier of the table such as `default.taxi_dataset`
+            catalog_kwargs: Optional arguments to use when setting up the Iceberg catalog
+            snapshot_properties: Custom properties to write to snapshot summary
+            mode: Write mode - APPEND, UPSERT, or OVERWRITE. Defaults to APPEND.
+                - APPEND: Add new data without checking for duplicates
+                - UPSERT: Update existing rows or insert new ones based on a join condition
+                - OVERWRITE: Replace table data (all data or filtered subset)
+            overwrite_filter: Optional filter for OVERWRITE mode to perform partial overwrites.
+                Must be a Ray Data expression from `ray.data.expressions`. Only rows matching
+                this filter are replaced. If None with OVERWRITE mode, replaces all table data.
+            upsert_kwargs: Optional arguments for upsert operations.
+                Supported parameters: join_cols (List[str]), case_sensitive (bool),
+                branch (str). Note: This implementation uses a copy-on-write strategy
+                that always updates all columns for matched keys and inserts all new keys.
+            overwrite_kwargs: Optional arguments to pass through to PyIceberg's table.overwrite()
+                method. Supported parameters include case_sensitive (bool) and branch (str).
+                See PyIceberg documentation for details.
+            upsert_commit_memory: [For UPSERT mode only] The heap memory in bytes
+                to reserve for the upsert commit operation. If None,
+                uses Ray's default memory allocation.
+
+        Note:
+            Schema evolution is automatically enabled. New columns in the incoming data
+            are automatically added to the table schema.
+        """
+        self.table_identifier = table_identifier
+        self._catalog_kwargs = (catalog_kwargs or {}).copy()
+        self._snapshot_properties = snapshot_properties or {}
+        self._mode = mode
+        self._overwrite_filter = overwrite_filter
+        self._upsert_kwargs = (upsert_kwargs or {}).copy()
+        self._overwrite_kwargs = (overwrite_kwargs or {}).copy()
+        self._upsert_commit_memory = upsert_commit_memory
+
+        # Validate kwargs are only set for relevant modes
+        if self._upsert_kwargs and self._mode != SaveMode.UPSERT:
+            raise ValueError(
+                f"upsert_kwargs can only be specified when mode is SaveMode.UPSERT, but mode is {self._mode}"
+            )
+        if self._overwrite_kwargs and self._mode != SaveMode.OVERWRITE:
+            raise ValueError(
+                f"overwrite_kwargs can only be specified when mode is SaveMode.OVERWRITE, but mode is {self._mode}"
+            )
+
+        # Remove invalid parameters from overwrite_kwargs if present
+        for invalid_param, reason in [
+            (
+                "overwrite_filter",
+                "should be passed as a separate parameter to write_iceberg()",
+            ),
+            (
+                "delete_filter",
+                "is an internal PyIceberg parameter; use 'overwrite_filter' instead",
+            ),
+        ]:
+            if self._overwrite_kwargs.pop(invalid_param, None) is not None:
+                logger.warning(
+                    f"Removed '{invalid_param}' from overwrite_kwargs: {reason}"
+                )
+
+        if "name" in self._catalog_kwargs:
+            self._catalog_name = self._catalog_kwargs.pop("name")
+        else:
+            self._catalog_name = "default"
+
+        self._table: "Table" = None
+        self._write_uuid: uuid.UUID = None
+
+    def __getstate__(self) -> dict:
+        """Exclude `_table` during pickling."""
+        state = self.__dict__.copy()
+        state.pop("_table", None)
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self._table = None
+
+    def _get_catalog(self) -> "Catalog":
+        from pyiceberg import catalog
+
+        return catalog.load_catalog(self._catalog_name, **self._catalog_kwargs)
+
+    def _reload_table(self) -> None:
+        """Reload the Iceberg table from the catalog."""
+        catalog = self._get_catalog()
+        self._table = catalog.load_table(self.table_identifier)
+
+    def _get_join_cols(self) -> List[str]:
+        """Get join columns for upsert, using table identifier fields as fallback."""
+        join_cols = self._upsert_kwargs.get("join_cols", [])
+        if not join_cols:
+            # Use table's identifier fields as fallback
+            for field_id in self._table.metadata.schema().identifier_field_ids:
+                col_name = self._table.metadata.schema().find_column_name(field_id)
+                if col_name:
+                    join_cols.append(col_name)
+        return join_cols
+
+    def _update_schema(self, incoming_schema: "pa.Schema") -> None:
+        """
+        Update the table schema to accommodate incoming data using union-by-name semantics.
+
+        This automatically handles:
+        - Adding new columns from the incoming schema
+        - Type promotion (e.g., int32 -> int64) where compatible
+        - Preserving existing columns not in the incoming schema
+        - Concurrent schema updates from multiple workers (with retry logic)
+
+        Note:
+            Each worker calls this once with the first block's schema. All blocks
+            within a worker are validated to have the same schema before calling this.
+
+        Args:
+            incoming_schema: The PyArrow schema from the incoming data
+        """
+        from pyiceberg.exceptions import CommitFailedException
+
+        max_retries = 3
+
+        for attempt in range(max_retries):
+            try:
+                with self._table.update_schema() as update:
+                    update.union_by_name(incoming_schema)
+                # Succeeded, reload to get latest table version and exit.
+                self._reload_table()
+                return
+            except CommitFailedException:
+                if attempt < max_retries - 1:
+                    logger.debug(
+                        f"Schema update conflict - another worker modified schema, "
+                        f"reloading and retrying (attempt {attempt + 1}/{max_retries})"
+                    )
+                    self._reload_table()
+                else:
+                    logger.error(
+                        "Failed to update schema after %d retries due to conflicts.",
+                        max_retries,
+                    )
+                    raise
+
+    def on_write_start(self) -> None:
+        """Initialize table for writing and create a shared write UUID."""
+        self._table = self._get_catalog().load_table(self.table_identifier)
+        self._write_uuid = uuid.uuid4()
+
+        # Validate join_cols for UPSERT mode before writing any files
+        if self._mode == SaveMode.UPSERT:
+            join_cols = self._upsert_kwargs.get("join_cols", [])
+            if not join_cols:
+                # Check if table has identifier fields as fallback
+                identifier_field_ids = (
+                    self._table.metadata.schema().identifier_field_ids
+                )
+                if not identifier_field_ids:
+                    raise ValueError(
+                        "join_cols must be specified in upsert_kwargs for UPSERT mode "
+                        "when table has no identifier fields"
+                    )
+
+    def write(
+        self, blocks: Iterable[Block], ctx: TaskContext
+    ) -> Union[List["DataFile"], tuple[List["DataFile"], Dict[str, List[Any]]]]:
+        """
+        Write blocks to Parquet files in storage and return DataFile metadata.
+
+        This runs on each worker in parallel. Files are written directly to storage
+        (S3, HDFS, etc.) and only metadata is returned to the driver.
+
+        Args:
+            blocks: Iterable of Ray Data blocks to write
+            ctx: TaskContext object containing task-specific information
+
+        Returns:
+            For APPEND/OVERWRITE: List of DataFile objects.
+            For UPSERT: Tuple of (List of DataFile objects, Dict mapping col names to key values).
+        """
+        from pyiceberg.io.pyarrow import _dataframe_to_data_files
+
+        if self._table is None:
+            self._reload_table()
+
+        all_data_files = []
+        join_keys_dict = defaultdict(list)
+        first_schema = None
+
+        # Extract join keys for copy-on-write upsert
+        extract_join_keys = self._mode == SaveMode.UPSERT
+
+        for block in blocks:
+            pa_table = BlockAccessor.for_block(block).to_arrow()
+            if pa_table.num_rows > 0:
+                # Schema validation and update strategy:
+                # 1. Each worker validates all its blocks have the same schema
+                # 2. Schema is updated once per worker (not per block) to minimize conflicts
+                # 3. Concurrent updates from different workers are handled by retry logic
+                if first_schema is None:
+                    first_schema = pa_table.schema
+                    # Update schema once per worker (handles concurrent updates from other workers)
+                    self._update_schema(first_schema)
+                elif not pa_table.schema.equals(first_schema):
+                    raise ValueError(
+                        f"Schema mismatch within worker: expected {first_schema}, "
+                        f"got {pa_table.schema}. All blocks must have the same schema."
+                    )
+
+                # Extract join key values for copy-on-write upsert
+                # These will be used to build a delete filter for matching rows
+                # Example: Upserting [{col_a: 2}, {col_a: 3}] extracts {col_a: [2, 3]}
+                #          which becomes filter: col_a IN (2, 3)
+                if extract_join_keys:
+                    join_cols = self._get_join_cols()
+                    for col in join_cols:
+                        join_keys_dict[col].extend(pa_table[col].to_pylist())
+
+                # Write data files to storage (distributed!)
+                # _dataframe_to_data_files writes Parquet files and returns DataFile metadata
+                data_files = list(
+                    _dataframe_to_data_files(
+                        table_metadata=self._table.metadata,
+                        write_uuid=self._write_uuid,
+                        df=pa_table,
+                        io=self._table.io,
+                    )
+                )
+                all_data_files.extend(data_files)
+
+        # Return appropriate type based on whether we extracted join keys
+        if extract_join_keys:
+            return (all_data_files, join_keys_dict)
+        return all_data_files
+
+    def on_write_complete(self, write_result: WriteResult) -> None:
+        """
+        Complete the write by committing all data files in a single transaction.
+
+        This runs on the driver after all workers finish writing files.
+        Collects all DataFile objects from all workers and performs a single
+        atomic commit based on the configured mode.
+        """
+        # Reload table to get latest metadata
+        self._reload_table()
+
+        # Collect all data files and join keys (if applicable) from all workers
+        all_data_files = []
+        join_keys_dicts = []
+
+        # Check if we're in upsert mode (returns tuple with join keys)
+        use_copy_on_write_upsert = self._mode == SaveMode.UPSERT
+
+        for write_return in write_result.write_returns:
+            if not write_return:
+                continue
+
+            if use_copy_on_write_upsert:
+                # For copy-on-write upsert, write() returns (data_files, join_keys_dict)
+                data_files, join_keys_dict = write_return
+                all_data_files.extend(data_files)
+                # Collect join keys dicts (will be merged in Ray task)
+                join_keys_dicts.append(join_keys_dict)
+            else:
+                # For other modes, write() returns just data_files
+                all_data_files.extend(write_return)
+
+        if not all_data_files:
+            return
+
+        # Create transaction and commit based on mode
+        if self._mode == SaveMode.APPEND:
+            self._commit_append(all_data_files)
+        elif self._mode == SaveMode.OVERWRITE:
+            self._commit_overwrite(all_data_files)
+        elif self._mode == SaveMode.UPSERT:
+            # Execute entire commit in Ray task to avoid OOM on driver
+            join_cols = self._get_join_cols()
+
+            # Configure task options with memory if specified
+            task_options = {}
+            if self._upsert_commit_memory is not None:
+                task_options["memory"] = self._upsert_commit_memory
+
+            ray.get(
+                _commit_upsert_task.options(**task_options).remote(
+                    table_identifier=self.table_identifier,
+                    catalog_name=self._catalog_name,
+                    catalog_kwargs=self._catalog_kwargs,
+                    data_files=all_data_files,
+                    join_keys_dicts=join_keys_dicts,
+                    join_cols=join_cols,
+                    snapshot_properties=self._snapshot_properties,
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported mode: {self._mode}")
+
+    def _commit_append(self, data_files: List["DataFile"]) -> None:
+        """Commit data files using APPEND mode."""
+        txn = self._table.transaction()
+        _append_and_commit(txn, data_files, self._snapshot_properties)
+
+    def _commit_overwrite(self, data_files: List["DataFile"]) -> None:
+        """Commit data files using OVERWRITE mode."""
+        txn = self._table.transaction()
+        # Delete matching data if filter provided
+        if self._overwrite_filter is not None:
+            from ray.data._internal.datasource.iceberg_datasource import (
+                _IcebergExpressionVisitor,
+            )
+
+            visitor = _IcebergExpressionVisitor()
+            pyi_filter = visitor.visit(self._overwrite_filter)
+
+            txn.delete(
+                delete_filter=pyi_filter,
+                snapshot_properties=self._snapshot_properties,
+                **self._overwrite_kwargs,
+            )
+        else:
+            # Full overwrite - delete all
+            from pyiceberg.expressions import AlwaysTrue
+
+            txn.delete(
+                delete_filter=AlwaysTrue(),
+                snapshot_properties=self._snapshot_properties,
+                **self._overwrite_kwargs,
+            )
+
+        # Append new data files and commit
+        _append_and_commit(txn, data_files, self._snapshot_properties)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/iceberg_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/iceberg_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..106cd5c99b17206efc6648be848e027b734cbd0c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/iceberg_datasource.py
@@ -0,0 +1,493 @@
+"""
+Module to read an iceberg table into a Ray Dataset, by using the Ray Datasource API.
+"""
+
+import heapq
+import itertools
+import logging
+from functools import partial
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import pyarrow as pa
+from packaging import version
+
+from ray.data._internal.planner.plan_expression.expression_visitors import _ExprVisitor
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    LiteralExpr,
+    Operation,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+)
+from ray.util import log_once
+from ray.util.annotations import DeveloperAPI
+
+try:
+    from pyiceberg.expressions import (
+        And,
+        EqualTo,
+        GreaterThan,
+        GreaterThanOrEqual,
+        In,
+        IsNull,
+        LessThan,
+        LessThanOrEqual,
+        Literal,
+        Not,
+        NotEqualTo,
+        NotIn,
+        NotNull,
+        Or,
+        Reference,
+        UnboundTerm,
+        literal,
+    )
+
+    RAY_DATA_OPERATION_TO_ICEBERG = {
+        Operation.EQ: EqualTo,
+        Operation.NE: NotEqualTo,
+        Operation.GT: GreaterThan,
+        Operation.GE: GreaterThanOrEqual,
+        Operation.LT: LessThan,
+        Operation.LE: LessThanOrEqual,
+        Operation.AND: And,
+        Operation.OR: Or,
+        Operation.IN: In,
+        Operation.NOT_IN: NotIn,
+        Operation.IS_NULL: IsNull,
+        Operation.IS_NOT_NULL: NotNull,
+        Operation.NOT: Not,
+    }
+except ImportError:
+    log_once("pyiceberg.expressions not found. Please install pyiceberg >= 0.9.0")
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.expressions import BooleanExpression
+    from pyiceberg.io import FileIO
+    from pyiceberg.manifest import DataFile
+    from pyiceberg.schema import Schema
+    from pyiceberg.table import DataScan, FileScanTask, Table
+    from pyiceberg.table.metadata import TableMetadata
+
+logger = logging.getLogger(__name__)
+
+
+class _IcebergExpressionVisitor(
+    _ExprVisitor["BooleanExpression | UnboundTerm[Any] | Literal[Any]"]
+):
+    """
+    Visitor that converts Ray Data expressions to PyIceberg expressions.
+
+    This enables Ray Data users to write filters using the familiar col() syntax
+    while leveraging Iceberg's native filtering capabilities.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> ray_expr = (col("date") >= "2024-01-01") & (col("status") == "active")
+        >>> iceberg_expr = _IcebergExpressionVisitor().visit(ray_expr)
+        >>> # iceberg_expr can now be used with PyIceberg's filter APIs
+    """
+
+    def visit_column(self, expr: "ColumnExpr") -> "UnboundTerm[Any]":
+        """Convert a column reference to an Iceberg reference."""
+        return Reference(expr.name)
+
+    def visit_literal(self, expr: "LiteralExpr") -> "Literal[Any]":
+        """Convert a literal value to an Iceberg literal."""
+        return literal(expr.value)
+
+    def visit_binary(self, expr: "BinaryExpr") -> "BooleanExpression":
+        """Convert a binary operation to an Iceberg expression."""
+        # Handle IN/NOT_IN specially since they don't visit the right operand
+        # (the right operand is a list literal that can't be converted)
+        if expr.op in (Operation.IN, Operation.NOT_IN):
+            left = self.visit(expr.left)
+            if not isinstance(expr.right, LiteralExpr):
+                raise ValueError(
+                    f"{expr.op.name} operation requires right operand to be a literal list, "
+                    f"got {type(expr.right).__name__}"
+                )
+            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](left, expr.right.value)
+
+        # For all other operations, visit both operands
+        left = self.visit(expr.left)
+        right = self.visit(expr.right)
+
+        if expr.op in RAY_DATA_OPERATION_TO_ICEBERG:
+            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](left, right)
+        else:
+            # Arithmetic operations are not supported in filter expressions
+            raise ValueError(
+                f"Unsupported binary operation for Iceberg filters: {expr.op}. "
+                f"Iceberg filters support: {RAY_DATA_OPERATION_TO_ICEBERG.keys()}. "
+                f"Arithmetic operations (ADD, SUB, MUL, DIV) cannot be used in filters."
+            )
+
+    def visit_unary(self, expr: "UnaryExpr") -> "BooleanExpression":
+        """Convert a unary operation to an Iceberg expression."""
+        operand = self.visit(expr.operand)
+
+        if expr.op in RAY_DATA_OPERATION_TO_ICEBERG:
+            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](operand)
+        else:
+            raise ValueError(
+                f"Unsupported unary operation for Iceberg: {expr.op}. "
+                f"Supported operations: {RAY_DATA_OPERATION_TO_ICEBERG.keys()}"
+            )
+
+    def visit_alias(
+        self, expr: "AliasExpr"
+    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
+        """Convert an aliased expression (just unwrap the alias)."""
+        return self.visit(expr.expr)
+
+    def visit_udf(
+        self, expr: "UDFExpr"
+    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
+        """UDF expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "UDF expressions cannot be converted to Iceberg expressions. "
+            "Iceberg filters must use simple column comparisons and boolean operations."
+        )
+
+    def visit_download(
+        self, expr: "DownloadExpr"
+    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
+        """Download expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "Download expressions cannot be converted to Iceberg expressions."
+        )
+
+    def visit_star(
+        self, expr: "StarExpr"
+    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
+        """Star expressions cannot be converted to Iceberg expressions."""
+        raise TypeError(
+            "Star expressions cannot be converted to Iceberg filter expressions."
+        )
+
+
+def _get_read_task(
+    tasks: Iterable["FileScanTask"],
+    table_io: "FileIO",
+    table_metadata: "TableMetadata",
+    row_filter: "BooleanExpression",
+    case_sensitive: bool,
+    limit: Optional[int],
+    schema: "Schema",
+    column_rename_map: Optional[Dict[str, str]],
+) -> Iterable[Block]:
+    # Determine the PyIceberg version to handle backward compatibility
+    import pyiceberg
+
+    from ray.data.datasource.datasource import _DatasourceProjectionPushdownMixin
+
+    def _generate_tables() -> Iterable[pa.Table]:
+        """Inner generator that yields tables without renaming."""
+        if version.parse(pyiceberg.__version__) >= version.parse("0.9.0"):
+            # Modern implementation using ArrowScan (PyIceberg 0.9.0+)
+            from pyiceberg.io.pyarrow import ArrowScan
+
+            # Initialize scanner with Iceberg metadata and query parameters
+            scanner = ArrowScan(
+                table_metadata=table_metadata,
+                io=table_io,
+                row_filter=row_filter,
+                projected_schema=schema,
+                case_sensitive=case_sensitive,
+                limit=limit,
+            )
+
+            # Convert scanned data to Arrow Table format
+            result_table = scanner.to_table(tasks=tasks)
+
+            # Stream results as RecordBatches for memory efficiency
+            for batch in result_table.to_batches():
+                yield pa.Table.from_batches([batch])
+
+        else:
+            # Legacy implementation using project_table (PyIceberg <0.9.0)
+            from pyiceberg.io import pyarrow as pyi_pa_io
+
+            # Use the PyIceberg API to read only a single task (specifically, a
+            # FileScanTask) - note that this is not as simple as reading a single
+            # parquet file, as there might be delete files, etc. associated, so we
+            # must use the PyIceberg API for the projection.
+            table = pyi_pa_io.project_table(
+                tasks=tasks,
+                table_metadata=table_metadata,
+                io=table_io,
+                row_filter=row_filter,
+                projected_schema=schema,
+                case_sensitive=case_sensitive,
+                limit=limit,
+            )
+            yield table
+
+    # Apply renames to all tables from the generator
+    yield from _DatasourceProjectionPushdownMixin._apply_rename_to_tables(
+        _generate_tables(), column_rename_map
+    )
+
+
+@DeveloperAPI
+class IcebergDatasource(Datasource):
+    """
+    Iceberg datasource to read Iceberg tables into a Ray Dataset. This module heavily
+    uses PyIceberg to read iceberg tables. All the routines in this class override
+    `ray.data.Datasource`.
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        row_filter: Union[str, "BooleanExpression"] = None,
+        selected_fields: Tuple[str, ...] = ("*",),
+        snapshot_id: Optional[int] = None,
+        scan_kwargs: Optional[Dict[str, Any]] = None,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize an IcebergDatasource.
+
+        Args:
+            table_identifier: Fully qualified table identifier (i.e.,
+                "db_name.table_name")
+            row_filter: A PyIceberg BooleanExpression to use to filter the data *prior*
+                 to reading
+            selected_fields: Which columns from the data to read, passed directly to
+                PyIceberg's load functions
+            snapshot_id: Optional snapshot ID for the Iceberg table
+            scan_kwargs: Optional arguments to pass to PyIceberg's Table.scan()
+                function
+            catalog_kwargs: Optional arguments to use when setting up the Iceberg
+                catalog
+        """
+        # Initialize parent class to set up predicate pushdown mixin
+        super().__init__()
+
+        _check_import(self, module="pyiceberg", package="pyiceberg")
+        from pyiceberg.expressions import AlwaysTrue
+
+        self._scan_kwargs = scan_kwargs if scan_kwargs is not None else {}
+        self._catalog_kwargs = catalog_kwargs if catalog_kwargs is not None else {}
+
+        if "name" in self._catalog_kwargs:
+            self._catalog_name = self._catalog_kwargs.pop("name")
+        else:
+            self._catalog_name = "default"
+
+        self.table_identifier = table_identifier
+
+        self._row_filter = row_filter if row_filter is not None else AlwaysTrue()
+        # Convert selected_fields to projection_map (identity mapping if specified)
+        # Note: Empty tuple () means no columns, None/"*" means all columns
+        if selected_fields is None or selected_fields == ("*",):
+            self._projection_map = None
+        else:
+            self._projection_map = {col: col for col in selected_fields}
+
+        if snapshot_id:
+            self._scan_kwargs["snapshot_id"] = snapshot_id
+
+        self._plan_files = None
+        self._table = None
+
+    def _get_catalog(self) -> "Catalog":
+        from pyiceberg import catalog
+
+        return catalog.load_catalog(self._catalog_name, **self._catalog_kwargs)
+
+    @property
+    def table(self) -> "Table":
+        """
+        Return the table reference from the catalog
+        """
+        if self._table is None:
+            catalog = self._get_catalog()
+            self._table = catalog.load_table(self.table_identifier)
+        return self._table
+
+    @property
+    def plan_files(self) -> List["FileScanTask"]:
+        """
+        Return the plan files specified by this query
+        """
+        # Calculate and cache the plan_files if they don't already exist
+        if self._plan_files is None:
+            data_scan = self._get_data_scan()
+            self._plan_files = data_scan.plan_files()
+
+        return self._plan_files
+
+    def _get_combined_filter(self) -> "BooleanExpression":
+        """Get the combined filter including both row_filter and pushed-down predicates."""
+        combined_filter = self._row_filter
+
+        if self._predicate_expr is not None:
+            # Convert Ray Data expression to PyIceberg expression using internal visitor
+            visitor = _IcebergExpressionVisitor()
+            iceberg_filter = visitor.visit(self._predicate_expr)
+
+            # Combine with existing row_filter using AND
+            from pyiceberg.expressions import AlwaysTrue, And
+
+            if not isinstance(combined_filter, AlwaysTrue):
+                combined_filter = And(combined_filter, iceberg_filter)
+            else:
+                combined_filter = iceberg_filter
+
+        return combined_filter
+
+    def _get_data_scan(self) -> "DataScan":
+        # Get the combined filter
+        combined_filter = self._get_combined_filter()
+
+        # Convert back to tuple for PyIceberg API (None -> ("*",))
+        data_columns = self._get_data_columns()
+        selected_fields = ("*",) if data_columns is None else tuple(data_columns)
+
+        data_scan = self.table.scan(
+            row_filter=combined_filter,
+            selected_fields=selected_fields,
+            **self._scan_kwargs,
+        )
+
+        return data_scan
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # Approximate the size by using the plan files - this will not
+        # incorporate the deletes, but that's a reasonable approximation
+        # task
+        return sum(task.file.file_size_in_bytes for task in self.plan_files)
+
+    def supports_predicate_pushdown(self) -> bool:
+        """Returns True to indicate this datasource supports predicate pushdown."""
+        return True
+
+    def supports_projection_pushdown(self) -> bool:
+        """Returns True to indicate this datasource supports projection pushdown."""
+        return True
+
+    @staticmethod
+    def _distribute_tasks_into_equal_chunks(
+        plan_files: Iterable["FileScanTask"], n_chunks: int
+    ) -> List[List["FileScanTask"]]:
+        """
+        Implement a greedy knapsack algorithm to distribute the files in the scan
+        across tasks, based on their file size, as evenly as possible
+        """
+        chunks = [list() for _ in range(n_chunks)]
+
+        chunk_sizes = [(0, chunk_id) for chunk_id in range(n_chunks)]
+        heapq.heapify(chunk_sizes)
+
+        # From largest to smallest, add the plan files to the smallest chunk one at a
+        # time
+        for plan_file in sorted(
+            plan_files, key=lambda f: f.file.file_size_in_bytes, reverse=True
+        ):
+            smallest_chunk = heapq.heappop(chunk_sizes)
+            chunks[smallest_chunk[1]].append(plan_file)
+            heapq.heappush(
+                chunk_sizes,
+                (
+                    smallest_chunk[0] + plan_file.file.file_size_in_bytes,
+                    smallest_chunk[1],
+                ),
+            )
+
+        return chunks
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        from pyiceberg.io import pyarrow as pyi_pa_io
+        from pyiceberg.manifest import DataFileContent
+
+        # Get the PyIceberg scan
+        data_scan = self._get_data_scan()
+        # Get the plan files in this query
+        plan_files = self.plan_files
+
+        # Get the projected schema for this scan, given all the row filters,
+        # snapshot ID, etc.
+        projected_schema = data_scan.projection()
+        # Get the arrow schema, to set in the metadata
+        pya_schema = pyi_pa_io.schema_to_pyarrow(projected_schema)
+
+        # Set the n_chunks to the min of the number of plan files and the actual
+        # requested n_chunks, so that there are no empty tasks
+        if parallelism > len(list(plan_files)):
+            parallelism = len(list(plan_files))
+            logger.warning(
+                f"Reducing the parallelism to {parallelism}, as that is the"
+                "number of files"
+            )
+
+        # Get required properties for reading tasks - table IO, table metadata,
+        # row filter, case sensitivity,limit and projected schema to pass
+        # them directly to `_get_read_task` to avoid capture of `self` reference
+        # within the closure carrying substantial overhead invoking these tasks
+        #
+        # See https://github.com/ray-project/ray/issues/49107 for more context
+        table_io = self.table.io
+        table_metadata = self.table.metadata
+        row_filter = self._get_combined_filter()
+        case_sensitive = self._scan_kwargs.get("case_sensitive", True)
+        limit = self._scan_kwargs.get("limit")
+
+        get_read_task = partial(
+            _get_read_task,
+            table_io=table_io,
+            table_metadata=table_metadata,
+            row_filter=row_filter,
+            case_sensitive=case_sensitive,
+            limit=limit,
+            schema=projected_schema,
+            column_rename_map=self.get_column_renames(),
+        )
+
+        read_tasks = []
+        # Chunk the plan files based on the requested parallelism
+        for chunk_tasks in IcebergDatasource._distribute_tasks_into_equal_chunks(
+            plan_files, parallelism
+        ):
+            unique_deletes: Set[DataFile] = set(
+                itertools.chain.from_iterable(
+                    [task.delete_files for task in chunk_tasks]
+                )
+            )
+            # Get a rough estimate of the number of deletes by just looking at
+            # position deletes. Equality deletes are harder to estimate, as they
+            # can delete multiple rows.
+            position_delete_count = sum(
+                delete.record_count
+                for delete in unique_deletes
+                if delete.content == DataFileContent.POSITION_DELETES
+            )
+            metadata = BlockMetadata(
+                num_rows=sum(task.file.record_count for task in chunk_tasks)
+                - position_delete_count,
+                size_bytes=sum(task.length for task in chunk_tasks),
+                input_files=[task.file.file_path for task in chunk_tasks],
+                exec_stats=None,
+            )
+            read_tasks.append(
+                ReadTask(
+                    read_fn=lambda tasks=chunk_tasks: get_read_task(tasks),
+                    metadata=metadata,
+                    schema=pya_schema,
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+
+        return read_tasks
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/image_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/image_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac561fbaaa9f40bc9e9562eb75186bdea23433d3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/image_datasink.py
@@ -0,0 +1,24 @@
+import io
+from typing import Any, Dict
+
+import pyarrow
+
+from ray.data.datasource.file_datasink import RowBasedFileDatasink
+
+
+class ImageDatasink(RowBasedFileDatasink):
+    def __init__(
+        self, path: str, column: str, file_format: str, **file_datasink_kwargs
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        self.column = column
+        self.file_format = file_format
+
+    def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
+        from PIL import Image
+
+        image = Image.fromarray(row[self.column])
+        buffer = io.BytesIO()
+        image.save(buffer, format=self.file_format)
+        file.write(buffer.getvalue())
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/image_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/image_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..18477bc52b1a69dfd69ba1b1738970a6a800fbab
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/image_datasource.py
@@ -0,0 +1,174 @@
+import io
+import logging
+import time
+from typing import TYPE_CHECKING, Iterator, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+from ray.data.datasource.file_meta_provider import DefaultFileMetadataProvider
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+# The default size multiplier for reading image data source.
+# This essentially is using image on-disk file size to estimate
+# in-memory data size.
+IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT = 1
+
+# The lower bound value to estimate image encoding ratio.
+IMAGE_ENCODING_RATIO_ESTIMATE_LOWER_BOUND = 0.5
+
+
+class ImageDatasource(FileBasedDatasource):
+    """A datasource that lets you read images."""
+
+    _WRITE_FILE_PER_ROW = True
+    _FILE_EXTENSIONS = ["png", "jpg", "jpeg", "tif", "tiff", "bmp", "gif"]
+    # Use 8 threads per task to read image files.
+    _NUM_THREADS_PER_TASK = 8
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        size: Optional[Tuple[int, int]] = None,
+        mode: Optional[str] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="PIL", package="Pillow")
+
+        if size is not None and len(size) != 2:
+            raise ValueError(
+                "Expected `size` to contain two integers for height and width, "
+                f"but got {len(size)} integers instead."
+            )
+
+        if size is not None and (size[0] < 0 or size[1] < 0):
+            raise ValueError(
+                f"Expected `size` to contain positive integers, but got {size} instead."
+            )
+
+        self.size = size
+        self.mode = mode
+
+        meta_provider = file_based_datasource_kwargs.get("meta_provider", None)
+        if isinstance(meta_provider, ImageFileMetadataProvider):
+            self._encoding_ratio = self._estimate_files_encoding_ratio()
+            meta_provider._set_encoding_ratio(self._encoding_ratio)
+        else:
+            self._encoding_ratio = IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+    def _read_stream(
+        self,
+        f: "pyarrow.NativeFile",
+        path: str,
+    ) -> Iterator[Block]:
+        from PIL import Image, UnidentifiedImageError
+
+        data = f.readall()
+
+        try:
+            image = Image.open(io.BytesIO(data))
+        except UnidentifiedImageError as e:
+            raise ValueError(f"PIL couldn't load image file at path '{path}'.") from e
+
+        if self.size is not None and image.size != tuple(reversed(self.size)):
+            height, width = self.size
+            image = image.resize((width, height), resample=Image.BILINEAR)
+        if self.mode is not None and image.mode != self.mode:
+            image = image.convert(self.mode)
+
+        builder = DelegatingBlockBuilder()
+        array = np.asarray(image)
+        item = {"image": array}
+        builder.add(item)
+        block = builder.build()
+
+        yield block
+
+    def _rows_per_file(self):
+        return 1
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        total_size = 0
+        for file_size in self._file_sizes():
+            # NOTE: check if file size is not None, because some metadata provider
+            # such as FastFileMetadataProvider does not provide file size information.
+            if file_size is not None:
+                total_size += file_size
+        return total_size * self._encoding_ratio
+
+    def _estimate_files_encoding_ratio(self) -> float:
+        """Return an estimate of the image files encoding ratio."""
+        start_time = time.perf_counter()
+        # Filter out empty file to avoid noise.
+        non_empty_path_and_size = list(
+            filter(lambda p: p[1] > 0, zip(self._paths(), self._file_sizes()))
+        )
+        num_files = len(non_empty_path_and_size)
+        if num_files == 0:
+            logger.warning(
+                "All input image files are empty. "
+                "Use on-disk file size to estimate images in-memory size."
+            )
+            return IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+        if self.size is not None and self.mode is not None:
+            # Use image size and mode to calculate data size for all images,
+            # because all images are homogeneous with same size after resizing.
+            # Resizing is enforced when reading every image in `ImageDatasource`
+            # when `size` argument is provided.
+            if self.mode in ["1", "L", "P"]:
+                dimension = 1
+            elif self.mode in ["RGB", "YCbCr", "LAB", "HSV"]:
+                dimension = 3
+            elif self.mode in ["RGBA", "CMYK", "I", "F"]:
+                dimension = 4
+            else:
+                logger.warning(f"Found unknown image mode: {self.mode}.")
+                return IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+            height, width = self.size
+            single_image_size = height * width * dimension
+            total_estimated_size = single_image_size * num_files
+            total_file_size = sum(p[1] for p in non_empty_path_and_size)
+            ratio = total_estimated_size / total_file_size
+        else:
+            # TODO(chengsu): sample images to estimate data size
+            ratio = IMAGE_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+        sampling_duration = time.perf_counter() - start_time
+        if sampling_duration > 5:
+            logger.warning(
+                "Image input size estimation took "
+                f"{round(sampling_duration, 2)} seconds."
+            )
+        logger.debug(f"Estimated image encoding ratio from sampling is {ratio}.")
+        return max(ratio, IMAGE_ENCODING_RATIO_ESTIMATE_LOWER_BOUND)
+
+
+class ImageFileMetadataProvider(DefaultFileMetadataProvider):
+    def _set_encoding_ratio(self, encoding_ratio: int):
+        """Set image file encoding ratio, to provide accurate size in bytes metadata."""
+        self._encoding_ratio = encoding_ratio
+
+    def _get_block_metadata(
+        self,
+        paths: List[str],
+        *,
+        rows_per_file: Optional[int],
+        file_sizes: List[Optional[int]],
+    ) -> BlockMetadata:
+        metadata = super()._get_block_metadata(
+            paths, rows_per_file=rows_per_file, file_sizes=file_sizes
+        )
+        if metadata.size_bytes is not None:
+            metadata.size_bytes = int(metadata.size_bytes * self._encoding_ratio)
+        return metadata
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/json_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/json_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..c81cea387616a2bf48859c0401c2d1a77e6cd77c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/json_datasink.py
@@ -0,0 +1,36 @@
+from typing import Any, Callable, Dict, Optional
+
+import pyarrow
+
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_based_datasource import _resolve_kwargs
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class JSONDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        pandas_json_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        pandas_json_args: Optional[Dict[str, Any]] = None,
+        file_format: str = "json",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        if pandas_json_args_fn is None:
+            pandas_json_args_fn = lambda: {}  # noqa: E731
+
+        if pandas_json_args is None:
+            pandas_json_args = {}
+
+        self.pandas_json_args_fn = pandas_json_args_fn
+        self.pandas_json_args = pandas_json_args
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        writer_args = _resolve_kwargs(self.pandas_json_args_fn, **self.pandas_json_args)
+        orient = writer_args.pop("orient", "records")
+        lines = writer_args.pop("lines", True)
+
+        block.to_pandas().to_json(file, orient=orient, lines=lines, **writer_args)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/json_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/json_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..7dc1d9c6a85a98f6ce3dcf1bca3d10b883c0c9f8
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/json_datasource.py
@@ -0,0 +1,285 @@
+import io
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import pandas as pd
+
+from ray.air.util.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data._internal.pandas_block import PandasBlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+logger = logging.getLogger(__name__)
+
+JSON_FILE_EXTENSIONS = [
+    "json",
+    "jsonl",
+    # gzip-compressed files
+    "json.gz",
+    "jsonl.gz",
+    # Brotli-compressed fi;es
+    "json.br",
+    "jsonl.br",
+    # Zstandard-compressed files
+    "json.zst",
+    "jsonl.zst",
+    # lz4-compressed files
+    "json.lz4",
+    "jsonl.lz4",
+]
+
+
+class ArrowJSONDatasource(FileBasedDatasource):
+    """JSON datasource, for reading and writing JSON and JSONL files."""
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        *,
+        arrow_json_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        from pyarrow import json
+
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if arrow_json_args is None:
+            arrow_json_args = {}
+
+        self.read_options = arrow_json_args.pop(
+            "read_options", json.ReadOptions(use_threads=False)
+        )
+        self.arrow_json_args = arrow_json_args
+
+    def _read_with_pyarrow_read_json(self, buffer: "pyarrow.lib.Buffer"):
+        """Read with PyArrow JSON reader, trying to auto-increase the
+        read block size in the case of the read object
+        straddling block boundaries."""
+        import pyarrow as pa
+        import pyarrow.json as pajson
+
+        # When reading large files, the default block size configured in PyArrow can be
+        # too small, resulting in the following error: `pyarrow.lib.ArrowInvalid:
+        # straddling object straddles two block boundaries (try to increase block
+        # size?)`. More information on this issue can be found here:
+        # https://github.com/apache/arrow/issues/25674
+        # The read will be retried with geometrically increasing block size
+        # until the size reaches `DataContext.get_current().target_max_block_size`.
+        # The initial block size will start at the PyArrow default block size
+        # or it can be manually set through the `read_options` parameter as follows.
+        # >>> import pyarrow.json as pajson
+        # >>> block_size = 10 << 20 # Set block size to 10MB
+        # >>> ray.data.read_json(  # doctest: +SKIP
+        # ...     "s3://anonymous@ray-example-data/log.json",
+        # ...     read_options=pajson.ReadOptions(block_size=block_size)
+        # ... )
+
+        init_block_size = self.read_options.block_size
+        max_block_size = DataContext.get_current().target_max_block_size
+        while True:
+            try:
+                yield pajson.read_json(
+                    io.BytesIO(buffer),
+                    read_options=self.read_options,
+                    **self.arrow_json_args,
+                )
+                self.read_options.block_size = init_block_size
+                break
+            except pa.ArrowInvalid as e:
+                if "straddling object straddles two block boundaries" in str(e):
+                    if (
+                        max_block_size is None
+                        or self.read_options.block_size < max_block_size
+                    ):
+                        # Increase the block size in case it was too small.
+                        logger.debug(
+                            f"JSONDatasource read failed with "
+                            f"block_size={self.read_options.block_size}. Retrying with "
+                            f"block_size={self.read_options.block_size * 2}."
+                        )
+                        self.read_options.block_size *= 2
+                    else:
+                        raise pa.ArrowInvalid(
+                            f"{e} - Auto-increasing block size to "
+                            f"{self.read_options.block_size} bytes failed. "
+                            f"Please try manually increasing the block size through "
+                            f"the `read_options` parameter to a larger size. "
+                            f"For example: `read_json(..., read_options="
+                            f"pyarrow.json.ReadOptions(block_size=10 << 25))`"
+                            f"More information on this issue can be found here: "
+                            f"https://github.com/apache/arrow/issues/25674"
+                        )
+                else:
+                    # unrelated error, simply reraise
+                    raise e
+
+    def _read_with_python_json(self, buffer: "pyarrow.lib.Buffer"):
+        """Fallback method to read JSON files with Python's native json.load(),
+        in case the default pyarrow json reader fails."""
+        import json
+
+        import pyarrow as pa
+
+        # Check if the buffer is empty
+        if buffer.size == 0:
+            return
+
+        parsed_json = json.load(io.BytesIO(buffer))
+        try:
+            yield pa.Table.from_pylist(parsed_json)
+        except AttributeError as e:
+            # For PyArrow < 7.0.0, `pa.Table.from_pylist()` is not available.
+            # Construct a dict from the list and call
+            # `pa.Table.from_pydict()` instead.
+            assert "no attribute 'from_pylist'" in str(e), str(e)
+            from collections import defaultdict
+
+            dct = defaultdict(list)
+            for row in parsed_json:
+                for k, v in row.items():
+                    dct[k].append(v)
+            yield pyarrow_table_from_pydict(dct)
+
+    # TODO(ekl) The PyArrow JSON reader doesn't support streaming reads.
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        import pyarrow as pa
+
+        buffer: pa.lib.Buffer = f.read_buffer()
+
+        try:
+            yield from self._read_with_pyarrow_read_json(buffer)
+        except pa.ArrowInvalid as e:
+            # If read with PyArrow fails, try falling back to native json.load().
+            logger.warning(
+                f"Error reading with pyarrow.json.read_json(). "
+                f"Falling back to native json.load(), which may be slower. "
+                f"PyArrow error was:\n{e}"
+            )
+            yield from self._read_with_python_json(buffer)
+
+
+class PandasJSONDatasource(FileBasedDatasource):
+
+    # Buffer size in bytes for reading files. Default is 1MB.
+    #
+    # pandas reads data in small chunks (~8 KiB), which leads to many costly
+    # small read requests when accessing cloud storage. To reduce overhead and
+    # improve performance, we wrap the file in a larger buffered reader that
+    # reads bigger blocks at once.
+    _BUFFER_SIZE = 1024**2
+
+    # In the case of zipped json files, we cannot infer the chunk_size.
+    _DEFAULT_CHUNK_SIZE = 10000
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        target_output_size_bytes: int,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self._target_output_size_bytes = target_output_size_bytes
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        chunksize = self._estimate_chunksize(f)
+
+        stream = StrictBufferedReader(f, buffer_size=self._BUFFER_SIZE)
+        if chunksize is None:
+            # When chunksize=None, pandas returns DataFrame directly (no context manager)
+            df = pd.read_json(stream, chunksize=chunksize, lines=True)
+            yield _cast_range_index_to_string(df)
+        else:
+            # When chunksize is a number, pandas returns JsonReader (supports context manager)
+            with pd.read_json(stream, chunksize=chunksize, lines=True) as reader:
+                for df in reader:
+                    yield _cast_range_index_to_string(df)
+
+    def _estimate_chunksize(self, f: "pyarrow.NativeFile") -> Optional[int]:
+        """Estimate the chunksize by sampling the first row.
+
+        This is necessary to avoid OOMs while reading the file.
+        """
+
+        if not f.seekable():
+            return self._DEFAULT_CHUNK_SIZE
+        assert f.tell() == 0, "File pointer must be at the beginning"
+
+        if self._target_output_size_bytes is None:
+            return None
+
+        stream = StrictBufferedReader(f, buffer_size=self._BUFFER_SIZE)
+        with pd.read_json(stream, chunksize=1, lines=True) as reader:
+            try:
+                df = _cast_range_index_to_string(next(reader))
+            except StopIteration:
+                return 1
+
+        block_accessor = PandasBlockAccessor.for_block(df)
+        if block_accessor.num_rows() == 0:
+            chunksize = 1
+        else:
+            bytes_per_row = block_accessor.size_bytes() / block_accessor.num_rows()
+            chunksize = max(round(self._target_output_size_bytes / bytes_per_row), 1)
+
+        # Reset file pointer to the beginning.
+        f.seek(0)
+
+        return chunksize
+
+    def _open_input_source(
+        self,
+        filesystem: "pyarrow.fs.FileSystem",
+        path: str,
+        **open_args,
+    ) -> "pyarrow.NativeFile":
+
+        compression = self.resolve_compression(path, open_args)
+
+        if compression is None:
+            # We use a seekable file to estimate chunksize.
+            return filesystem.open_input_file(path)
+
+        return super()._open_input_source(filesystem, path, **open_args)
+
+
+def _cast_range_index_to_string(df: pd.DataFrame):
+    # NOTE: PandasBlockAccessor doesn't support RangeIndex, so we need to convert
+    # to string.
+    if isinstance(df.columns, pd.RangeIndex):
+        df.columns = df.columns.astype(str)
+    return df
+
+
+class StrictBufferedReader(io.RawIOBase):
+    """Wrapper that prevents premature file closure and ensures full-buffered reads.
+
+    This is necessary for two reasons:
+    1. The datasource reads the file twice -- first to sample and determine the chunk size,
+       and again to load the actual data. Since pandas assumes ownership of the file and
+       may close it, we prevent that by explicitly detaching the underlying file before
+       closing the buffer.
+
+    2. pandas wraps the file in a TextIOWrapper to decode bytes into text. TextIOWrapper
+       prefers calling read1(), which doesn't prefetch for random-access files
+       (e.g., from PyArrow). This wrapper forces all reads through the full buffer to
+       avoid inefficient small-range S3 GETs.
+    """
+
+    def __init__(self, file: io.RawIOBase, buffer_size: int):
+        self._file = io.BufferedReader(file, buffer_size=buffer_size)
+
+    def read(self, size=-1, /):
+        return self._file.read(size)
+
+    def readable(self) -> bool:
+        return True
+
+    def close(self):
+        if not self.closed:
+            self._file.detach()
+            self._file.close()
+            super().close()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/lance_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/lance_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..39551929a901109e43f6ddeb09e5880468e60d78
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/lance_datasink.py
@@ -0,0 +1,238 @@
+import pickle
+from itertools import chain
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import pyarrow as pa
+
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from lance.fragment import FragmentMetadata
+
+
+def _write_fragment(
+    stream: Iterable[Union["pa.Table", "pd.DataFrame"]],
+    uri: str,
+    *,
+    schema: Optional["pa.Schema"] = None,
+    max_rows_per_file: int = 64 * 1024 * 1024,
+    max_bytes_per_file: Optional[int] = None,
+    max_rows_per_group: int = 1024,  # Only useful for v1 writer.
+    data_storage_version: Optional[str] = None,
+    storage_options: Optional[Dict[str, Any]] = None,
+) -> List[Tuple["FragmentMetadata", "pa.Schema"]]:
+    import pandas as pd
+    from lance.fragment import DEFAULT_MAX_BYTES_PER_FILE, write_fragments
+
+    if schema is None:
+        first = next(stream)
+        if isinstance(first, pd.DataFrame):
+            schema = pa.Schema.from_pandas(first).remove_metadata()
+        else:
+            schema = first.schema
+        if len(schema.names) == 0:
+            # Empty table.
+            schema = None
+
+        stream = chain([first], stream)
+
+    def record_batch_converter():
+        for block in stream:
+            tbl = BlockAccessor.for_block(block).to_arrow()
+            yield from tbl.to_batches()
+
+    max_bytes_per_file = (
+        DEFAULT_MAX_BYTES_PER_FILE if max_bytes_per_file is None else max_bytes_per_file
+    )
+
+    reader = pa.RecordBatchReader.from_batches(schema, record_batch_converter())
+    fragments = write_fragments(
+        reader,
+        uri,
+        schema=schema,
+        max_rows_per_file=max_rows_per_file,
+        max_rows_per_group=max_rows_per_group,
+        max_bytes_per_file=max_bytes_per_file,
+        data_storage_version=data_storage_version,
+        storage_options=storage_options,
+    )
+    return [(fragment, schema) for fragment in fragments]
+
+
+class _BaseLanceDatasink(Datasink):
+    """Base class for Lance Datasink."""
+
+    def __init__(
+        self,
+        uri: str,
+        schema: Optional[pa.Schema] = None,
+        mode: Literal["create", "append", "overwrite"] = "create",
+        storage_options: Optional[Dict[str, Any]] = None,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.uri = uri
+        self.schema = schema
+        self.mode = mode
+
+        self.read_version: Optional[int] = None
+        self.storage_options = storage_options
+
+    @property
+    def supports_distributed_writes(self) -> bool:
+        return True
+
+    def on_write_start(self):
+        _check_import(self, module="lance", package="pylance")
+
+        import lance
+
+        if self.mode == "append":
+            ds = lance.LanceDataset(self.uri, storage_options=self.storage_options)
+            self.read_version = ds.version
+            if self.schema is None:
+                self.schema = ds.schema
+
+    def on_write_complete(
+        self,
+        write_results: List[List[Tuple[str, str]]],
+    ):
+        import warnings
+
+        import lance
+
+        if not write_results:
+            warnings.warn(
+                "write_results is empty.",
+                DeprecationWarning,
+            )
+            return
+        if hasattr(write_results, "write_returns"):
+            write_results = write_results.write_returns
+
+        if len(write_results) == 0:
+            warnings.warn(
+                "write results is empty. please check ray version or internal error",
+                DeprecationWarning,
+            )
+            return
+
+        fragments = []
+        schema = None
+        for batch in write_results:
+            for fragment_str, schema_str in batch:
+                fragment = pickle.loads(fragment_str)
+                fragments.append(fragment)
+                schema = pickle.loads(schema_str)
+        # Check weather writer has fragments or not.
+        # Skip commit when there are no fragments.
+        if not schema:
+            return
+        if self.mode in {"create", "overwrite"}:
+            op = lance.LanceOperation.Overwrite(schema, fragments)
+        elif self.mode == "append":
+            op = lance.LanceOperation.Append(fragments)
+        lance.LanceDataset.commit(
+            self.uri,
+            op,
+            read_version=self.read_version,
+            storage_options=self.storage_options,
+        )
+
+
+class LanceDatasink(_BaseLanceDatasink):
+    """Lance Ray Datasink.
+
+    Write a Ray dataset to lance.
+
+    If we expect to write larger-than-memory files,
+    we can use `LanceFragmentWriter` and `LanceCommitter`.
+
+    Args:
+        uri : the base URI of the dataset.
+        schema : pyarrow.Schema, optional.
+            The schema of the dataset.
+        mode : str, optional
+            The write mode. Default is 'append'.
+            Choices are 'append', 'create', 'overwrite'.
+        min_rows_per_file : int, optional
+            The minimum number of rows per file. Default is 1024 * 1024.
+        max_rows_per_file : int, optional
+            The maximum number of rows per file. Default is 64 * 1024 * 1024.
+        data_storage_version: optional, str, default None
+            The version of the data storage format to use. Newer versions are more
+            efficient but require newer versions of lance to read.  The default is
+            "legacy" which will use the legacy v1 version.  See the user guide
+            for more details.
+        storage_options : Dict[str, Any], optional
+            The storage options for the writer. Default is None.
+    """
+
+    NAME = "Lance"
+
+    def __init__(
+        self,
+        uri: str,
+        schema: Optional[pa.Schema] = None,
+        mode: Literal["create", "append", "overwrite"] = "create",
+        min_rows_per_file: int = 1024 * 1024,
+        max_rows_per_file: int = 64 * 1024 * 1024,
+        data_storage_version: Optional[str] = None,
+        storage_options: Optional[Dict[str, Any]] = None,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(
+            uri,
+            schema=schema,
+            mode=mode,
+            storage_options=storage_options,
+            *args,
+            **kwargs,
+        )
+
+        self.min_rows_per_file = min_rows_per_file
+        self.max_rows_per_file = max_rows_per_file
+        self.data_storage_version = data_storage_version
+        # if mode is append, read_version is read from existing dataset.
+        self.read_version: Optional[int] = None
+
+    @property
+    def min_rows_per_write(self) -> int:
+        return self.min_rows_per_file
+
+    def get_name(self) -> str:
+        return self.NAME
+
+    def write(
+        self,
+        blocks: Iterable[Union[pa.Table, "pd.DataFrame"]],
+        _ctx,
+    ):
+        fragments_and_schema = _write_fragment(
+            blocks,
+            self.uri,
+            schema=self.schema,
+            max_rows_per_file=self.max_rows_per_file,
+            data_storage_version=self.data_storage_version,
+            storage_options=self.storage_options,
+        )
+        return [
+            (pickle.dumps(fragment), pickle.dumps(schema))
+            for fragment, schema in fragments_and_schema
+        ]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/lance_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/lance_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..389d21a37d11ed9fb131658f44622312ee61135e
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/lance_datasource.py
@@ -0,0 +1,136 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional
+
+import numpy as np
+
+from ray._common.retry import call_with_retry
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+
+class LanceDatasource(Datasource):
+    """Lance datasource, for reading Lance dataset."""
+
+    # Errors to retry when reading Lance fragments.
+    READ_FRAGMENTS_ERRORS_TO_RETRY = ["LanceError(IO)"]
+    # Maximum number of attempts to read Lance fragments.
+    READ_FRAGMENTS_MAX_ATTEMPTS = 10
+    # Maximum backoff seconds between attempts to read Lance fragments.
+    READ_FRAGMENTS_RETRY_MAX_BACKOFF_SECONDS = 32
+
+    def __init__(
+        self,
+        uri: str,
+        columns: Optional[List[str]] = None,
+        filter: Optional[str] = None,
+        storage_options: Optional[Dict[str, str]] = None,
+        scanner_options: Optional[Dict[str, Any]] = None,
+    ):
+        _check_import(self, module="lance", package="pylance")
+
+        import lance
+
+        self.uri = uri
+        self.scanner_options = scanner_options or {}
+        if columns is not None:
+            self.scanner_options["columns"] = columns
+        if filter is not None:
+            self.scanner_options["filter"] = filter
+        self.storage_options = storage_options
+        self.lance_ds = lance.dataset(uri=uri, storage_options=storage_options)
+
+        match = []
+        match.extend(self.READ_FRAGMENTS_ERRORS_TO_RETRY)
+        match.extend(DataContext.get_current().retried_io_errors)
+        self._retry_params = {
+            "description": "read lance fragments",
+            "match": match,
+            "max_attempts": self.READ_FRAGMENTS_MAX_ATTEMPTS,
+            "max_backoff_s": self.READ_FRAGMENTS_RETRY_MAX_BACKOFF_SECONDS,
+        }
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        read_tasks = []
+        ds_fragments = self.scanner_options.get("fragments")
+        if ds_fragments is None:
+            ds_fragments = self.lance_ds.get_fragments()
+
+        for fragments in np.array_split(ds_fragments, parallelism):
+            if len(fragments) <= 0:
+                continue
+
+            fragment_ids = [f.metadata.id for f in fragments]
+            num_rows = sum(f.count_rows() for f in fragments)
+            input_files = [
+                data_file.path() for f in fragments for data_file in f.data_files()
+            ]
+
+            # TODO(chengsu): Take column projection into consideration for schema.
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=None,
+                input_files=input_files,
+                exec_stats=None,
+            )
+            scanner_options = self.scanner_options
+            lance_ds = self.lance_ds
+            retry_params = self._retry_params
+
+            read_task = ReadTask(
+                lambda f=fragment_ids: _read_fragments_with_retry(
+                    f,
+                    lance_ds,
+                    scanner_options,
+                    retry_params,
+                ),
+                metadata,
+                schema=fragments[0].schema,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO(chengsu): Add memory size estimation to improve auto-tune of parallelism.
+        return None
+
+
+def _read_fragments_with_retry(
+    fragment_ids,
+    lance_ds,
+    scanner_options,
+    retry_params,
+) -> Iterator["pyarrow.Table"]:
+    return call_with_retry(
+        lambda: _read_fragments(fragment_ids, lance_ds, scanner_options),
+        **retry_params,
+    )
+
+
+def _read_fragments(
+    fragment_ids,
+    lance_ds,
+    scanner_options,
+) -> Iterator["pyarrow.Table"]:
+    """Read Lance fragments in batches.
+
+    NOTE: Use fragment ids, instead of fragments as parameter, because pickling
+    LanceFragment is expensive.
+    """
+    import pyarrow
+
+    fragments = [lance_ds.get_fragment(id) for id in fragment_ids]
+    scanner_options["fragments"] = fragments
+    scanner = lance_ds.scanner(**scanner_options)
+    for batch in scanner.to_reader():
+        yield pyarrow.Table.from_batches([batch])
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/mcap_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/mcap_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4ae4fdf90a5661823ecf6cca8dd7305904ffed6
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/mcap_datasource.py
@@ -0,0 +1,258 @@
+"""MCAP (Message Capture) datasource for Ray Data.
+
+MCAP is a standardized format for storing timestamped messages from robotics and
+autonomous systems, commonly used for sensor data, control commands, and other
+time-series data.
+"""
+
+import json
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow
+    from mcap.reader import Channel, Message, Schema
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TimeRange:
+    """Time range for filtering MCAP messages.
+
+    Attributes:
+        start_time: Start time in nanoseconds (inclusive).
+        end_time: End time in nanoseconds (exclusive).
+    """
+
+    start_time: int
+    end_time: int
+
+    def __post_init__(self):
+        """Validate time range after initialization."""
+        if self.start_time >= self.end_time:
+            raise ValueError(
+                f"start_time ({self.start_time}) must be less than "
+                f"end_time ({self.end_time})"
+            )
+        if self.start_time < 0 or self.end_time < 0:
+            raise ValueError(
+                f"time values must be non-negative, got start_time={self.start_time}, "
+                f"end_time={self.end_time}"
+            )
+
+
+@DeveloperAPI
+class MCAPDatasource(FileBasedDatasource):
+    """MCAP (Message Capture) datasource for Ray Data.
+
+    This datasource provides reading of MCAP files with predicate pushdown
+    optimization for filtering by topics, time ranges, and message types.
+
+    MCAP is a standardized format for storing timestamped messages from robotics and
+    autonomous systems, commonly used for sensor data, control commands, and other
+    time-series data.
+
+    Examples:
+        Basic usage:
+
+        >>> import ray  # doctest: +SKIP
+        >>> ds = ray.data.read_mcap("/path/to/data.mcap")  # doctest: +SKIP
+
+        With topic filtering and time range:
+
+        >>> from ray.data.datasource import TimeRange  # doctest: +SKIP
+        >>> ds = ray.data.read_mcap(  # doctest: +SKIP
+        ...     "/path/to/data.mcap",
+        ...     topics={"/camera/image_raw", "/lidar/points"},
+        ...     time_range=TimeRange(start_time=1000000000, end_time=2000000000)
+        ... )  # doctest: +SKIP
+
+        With multiple files and metadata:
+
+        >>> ds = ray.data.read_mcap(  # doctest: +SKIP
+        ...     ["file1.mcap", "file2.mcap"],
+        ...     topics={"/camera/image_raw", "/lidar/points"},
+        ...     message_types={"sensor_msgs/Image", "sensor_msgs/PointCloud2"},
+        ...     include_metadata=True
+        ... )  # doctest: +SKIP
+    """
+
+    _FILE_EXTENSIONS = ["mcap"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        topics: Optional[Union[List[str], Set[str]]] = None,
+        time_range: Optional[TimeRange] = None,
+        message_types: Optional[Union[List[str], Set[str]]] = None,
+        include_metadata: bool = True,
+        **file_based_datasource_kwargs,
+    ):
+        """Initialize MCAP datasource.
+
+        Args:
+            paths: Path or list of paths to MCAP files.
+            topics: Optional list/set of topic names to include. If specified,
+                only messages from these topics will be read.
+            time_range: Optional TimeRange for filtering messages by timestamp.
+                TimeRange contains start_time and end_time in nanoseconds, where
+                both values must be non-negative and start_time < end_time.
+            message_types: Optional list/set of message type names (schema names)
+                to include. Only messages with matching schema names will be read.
+            include_metadata: Whether to include MCAP metadata fields in the output.
+                Defaults to True. When True, includes schema, channel, and message
+                metadata.
+            **file_based_datasource_kwargs: Additional arguments for FileBasedDatasource.
+        """
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="mcap", package="mcap")
+
+        # Convert to sets for faster lookup
+        self._topics = set(topics) if topics else None
+        self._message_types = set(message_types) if message_types else None
+        self._time_range = time_range
+        self._include_metadata = include_metadata
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        """Read MCAP file and yield blocks of message data.
+
+        This method implements efficient MCAP reading with predicate pushdown.
+        It uses MCAP's built-in filtering capabilities for optimal performance
+        and applies additional filters when needed.
+
+        Args:
+            f: File-like object to read from. Must be seekable for MCAP reading.
+            path: Path to the MCAP file being processed.
+
+        Yields:
+            Block: Blocks of MCAP message data as pyarrow Tables.
+
+        Raises:
+            ValueError: If the MCAP file cannot be read or has invalid format.
+        """
+        from mcap.reader import make_reader
+
+        reader = make_reader(f)
+        # Note: MCAP summaries are optional and iter_messages works without them
+        # We don't need to validate the summary since it's not required
+
+        # Use MCAP's built-in filtering for topics and time range
+        messages = reader.iter_messages(
+            topics=list(self._topics) if self._topics else None,
+            start_time=self._time_range.start_time if self._time_range else None,
+            end_time=self._time_range.end_time if self._time_range else None,
+            log_time_order=True,
+            reverse=False,
+        )
+
+        builder = DelegatingBlockBuilder()
+
+        for schema, channel, message in messages:
+            # Apply filters that couldn't be pushed down to MCAP level
+            if not self._should_include_message(schema, channel, message):
+                continue
+
+            # Convert message to dictionary format
+            message_data = self._message_to_dict(schema, channel, message, path)
+            builder.add(message_data)
+
+        # Yield the block if we have any messages
+        if builder.num_rows() > 0:
+            yield builder.build()
+
+    def _should_include_message(
+        self, schema: "Schema", channel: "Channel", message: "Message"
+    ) -> bool:
+        """Check if a message should be included based on filters.
+
+        This method applies Python-level filtering that cannot be pushed down
+        to the MCAP library level. Topic filters are already handled by the
+        MCAP reader, so only message_types filtering is needed here.
+
+        Args:
+            schema: MCAP schema object containing message type information.
+            channel: MCAP channel object containing topic and metadata.
+            message: MCAP message object containing the actual data.
+
+        Returns:
+            True if the message should be included, False otherwise.
+        """
+        # Message type filter (cannot be pushed down to MCAP reader)
+        if self._message_types and schema and schema.name not in self._message_types:
+            return False
+
+        return True
+
+    def _message_to_dict(
+        self, schema: "Schema", channel: "Channel", message: "Message", path: str
+    ) -> Dict[str, Any]:
+        """Convert MCAP message to dictionary format.
+
+        This method converts MCAP message objects into a standardized dictionary
+        format suitable for Ray Data processing.
+
+        Args:
+            schema: MCAP schema object containing message type and encoding info.
+            channel: MCAP channel object containing topic and channel metadata.
+            message: MCAP message object containing the actual message data.
+            path: Path to the source file (for include_paths functionality).
+
+        Returns:
+            Dictionary containing message data in Ray Data format.
+        """
+        # Decode message data based on encoding
+        decoded_data = message.data
+        if channel.message_encoding == "json" and isinstance(message.data, bytes):
+            try:
+                decoded_data = json.loads(message.data.decode("utf-8"))
+            except (json.JSONDecodeError, UnicodeDecodeError):
+                # Keep raw bytes if decoding fails
+                decoded_data = message.data
+
+        # Core message data
+        message_data = {
+            "data": decoded_data,
+            "topic": channel.topic,
+            "log_time": message.log_time,
+            "publish_time": message.publish_time,
+            "sequence": message.sequence,
+        }
+
+        # Add metadata if requested
+        if self._include_metadata:
+            message_data.update(
+                {
+                    "channel_id": message.channel_id,
+                    "message_encoding": channel.message_encoding,
+                    "schema_name": schema.name if schema else None,
+                    "schema_encoding": schema.encoding if schema else None,
+                    "schema_data": schema.data if schema else None,
+                }
+            )
+
+        # Add file path if include_paths is enabled (from FileBasedDatasource)
+        if getattr(self, "include_paths", False):
+            message_data["path"] = path
+
+        return message_data
+
+    def get_name(self) -> str:
+        """Return a human-readable name for this datasource."""
+        return "MCAP"
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        """Whether this datasource supports distributed reads.
+
+        MCAP files can be read in parallel across multiple files.
+        """
+        return True
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/mongo_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/mongo_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..e3b61d17612244959ed803161cb1ec2634798441
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/mongo_datasink.py
@@ -0,0 +1,48 @@
+import logging
+from typing import Iterable
+
+from ray.data._internal.datasource.mongo_datasource import (
+    _validate_database_collection_exist,
+)
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDatasink(Datasink[None]):
+    def __init__(self, uri: str, database: str, collection: str) -> None:
+        _check_import(self, module="pymongo", package="pymongo")
+        _check_import(self, module="pymongoarrow", package="pymongoarrow")
+
+        self.uri = uri
+        self.database = database
+        self.collection = collection
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        import pymongo
+
+        _validate_database_collection_exist(
+            pymongo.MongoClient(self.uri), self.database, self.collection
+        )
+
+        def write_block(uri: str, database: str, collection: str, block: Block):
+            from pymongoarrow.api import write
+
+            block = BlockAccessor.for_block(block).to_arrow()
+            client = pymongo.MongoClient(uri)
+            write(client[database][collection], block)
+
+        builder = DelegatingBlockBuilder()
+        for block in blocks:
+            builder.add_block(block)
+        block = builder.build()
+
+        write_block(self.uri, self.database, self.collection, block)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/mongo_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/mongo_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..e392482114ae9fd142b84b175b508a9e136c8982
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/mongo_datasource.py
@@ -0,0 +1,142 @@
+import logging
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import pymongoarrow.api
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDatasource(Datasource):
+    """Datasource for reading from and writing to MongoDB."""
+
+    def __init__(
+        self,
+        uri: str,
+        database: str,
+        collection: str,
+        pipeline: Optional[List[Dict]] = None,
+        schema: Optional["pymongoarrow.api.Schema"] = None,
+        **mongo_args,
+    ):
+        self._uri = uri
+        self._database = database
+        self._collection = collection
+        self._pipeline = pipeline
+        self._schema = schema
+        self._mongo_args = mongo_args
+        # If pipeline is unspecified, read the entire collection.
+        if not pipeline:
+            self._pipeline = [{"$match": {"_id": {"$exists": "true"}}}]
+        # Initialize Mongo client lazily later when creating read tasks.
+        self._client = None
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO(jian): Add memory size estimation to improve auto-tune of parallelism.
+        return None
+
+    def _get_match_query(self, pipeline: List[Dict]) -> Dict:
+        if len(pipeline) == 0 or "$match" not in pipeline[0]:
+            return {}
+        return pipeline[0]["$match"]
+
+    def _get_or_create_client(self):
+        import pymongo
+
+        if self._client is None:
+            self._client = pymongo.MongoClient(self._uri)
+            _validate_database_collection_exist(
+                self._client, self._database, self._collection
+            )
+            self._avg_obj_size = self._client[self._database].command(
+                "collStats", self._collection
+            )["avgObjSize"]
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        from bson.objectid import ObjectId
+
+        self._get_or_create_client()
+        coll = self._client[self._database][self._collection]
+        match_query = self._get_match_query(self._pipeline)
+        partitions_ids = list(
+            coll.aggregate(
+                [
+                    {"$match": match_query},
+                    {"$bucketAuto": {"groupBy": "$_id", "buckets": parallelism}},
+                ],
+                allowDiskUse=True,
+            )
+        )
+
+        def make_block(
+            uri: str,
+            database: str,
+            collection: str,
+            pipeline: List[Dict],
+            min_id: ObjectId,
+            max_id: ObjectId,
+            right_closed: bool,
+            schema: "pymongoarrow.api.Schema",
+            kwargs: dict,
+        ) -> Block:
+            import pymongo
+            from pymongoarrow.api import aggregate_arrow_all
+
+            # A range query over the partition.
+            match = [
+                {
+                    "$match": {
+                        "_id": {
+                            "$gte": min_id,
+                            "$lte" if right_closed else "$lt": max_id,
+                        }
+                    }
+                }
+            ]
+            client = pymongo.MongoClient(uri)
+            return aggregate_arrow_all(
+                client[database][collection], match + pipeline, schema=schema, **kwargs
+            )
+
+        read_tasks: List[ReadTask] = []
+
+        for i, partition in enumerate(partitions_ids):
+            metadata = BlockMetadata(
+                num_rows=partition["count"],
+                size_bytes=partition["count"] * self._avg_obj_size,
+                input_files=None,
+                exec_stats=None,
+            )
+            make_block_args = (
+                self._uri,
+                self._database,
+                self._collection,
+                self._pipeline,
+                partition["_id"]["min"],
+                partition["_id"]["max"],
+                i == len(partitions_ids) - 1,
+                self._schema,
+                self._mongo_args,
+            )
+            read_task = ReadTask(
+                lambda args=make_block_args: [make_block(*args)],
+                metadata,
+                per_task_row_limit=per_task_row_limit,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+
+def _validate_database_collection_exist(client, database: str, collection: str):
+    db_names = client.list_database_names()
+    if database not in db_names:
+        raise ValueError(f"The destination database {database} doesn't exist.")
+    collection_names = client[database].list_collection_names()
+    if collection not in collection_names:
+        raise ValueError(f"The destination collection {collection} doesn't exist.")
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/numpy_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/numpy_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..5fc6f3d4df43d81b299cdc4da9a182798158a454
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/numpy_datasink.py
@@ -0,0 +1,23 @@
+import numpy as np
+import pyarrow
+
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class NumpyDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        column: str,
+        *,
+        file_format: str = "npy",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        self.column = column
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        value = block.to_numpy(self.column)
+        np.save(file, value)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/numpy_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/numpy_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f7d09d44910060be971cb9ce03d58f4a4d4c0b7
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/numpy_datasource.py
@@ -0,0 +1,41 @@
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+
+import numpy as np
+
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class NumpyDatasource(FileBasedDatasource):
+    """Numpy datasource, for reading and writing Numpy files."""
+
+    _COLUMN_NAME = "data"
+    _FILE_EXTENSIONS = ["npy"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        numpy_load_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if numpy_load_args is None:
+            numpy_load_args = {}
+
+        self.numpy_load_args = numpy_load_args
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        # TODO(ekl) Ideally numpy can read directly from the file, but it
+        # seems like it requires the file to be seekable.
+        buf = BytesIO()
+        data = f.readall()
+        buf.write(data)
+        buf.seek(0)
+        yield BlockAccessor.batch_to_block(
+            {"data": np.load(buf, allow_pickle=True, **self.numpy_load_args)}
+        )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_bulk_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_bulk_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..72d2015713ead6d1f16dcac07dbb27f2f5a0f421
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_bulk_datasource.py
@@ -0,0 +1,51 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+
+class ParquetBulkDatasource(FileBasedDatasource):
+    """Minimal Parquet datasource, for reading and writing Parquet files."""
+
+    _FILE_EXTENSIONS = ["parquet"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        read_table_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        if read_table_args is None:
+            read_table_args = {}
+
+        self.read_table_args = read_table_args
+
+    def get_name(self):
+        """Return a human-readable name for this datasource.
+        This will be used as the names of the read tasks.
+        Note: overrides the base `FileBasedDatasource` method.
+        """
+        return "ParquetBulk"
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        import pyarrow.parquet as pq
+
+        use_threads = self.read_table_args.pop("use_threads", False)
+        yield pq.read_table(f, use_threads=use_threads, **self.read_table_args)
+
+    def _open_input_source(
+        self,
+        filesystem: "pyarrow.fs.FileSystem",
+        path: str,
+        **open_args,
+    ) -> "pyarrow.NativeFile":
+        # Parquet requires `open_input_file` due to random access reads
+        return filesystem.open_input_file(path, **open_args)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..91639787672d70d8f3ff9b5c1d41eefa4d916182
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_datasink.py
@@ -0,0 +1,295 @@
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional
+
+from ray._common.retry import call_with_retry
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME
+from ray.data._internal.savemode import SaveMode
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.file_based_datasource import _resolve_kwargs
+from ray.data.datasource.file_datasink import _FileDatasink
+from ray.data.datasource.filename_provider import FilenameProvider
+
+if TYPE_CHECKING:
+    import pyarrow
+
+WRITE_FILE_MAX_ATTEMPTS = 10
+WRITE_FILE_RETRY_MAX_BACKOFF_SECONDS = 32
+
+# Map Ray Data's SaveMode to pyarrow's existing_data_behavior property which is exposed via the
+# `pyarrow.dataset.write_dataset` function.
+# Docs: https://arrow.apache.org/docs/python/generated/pyarrow.dataset.write_dataset.html
+EXISTING_DATA_BEHAVIOR_MAP = {
+    SaveMode.APPEND: "overwrite_or_ignore",
+    SaveMode.OVERWRITE: "overwrite_or_ignore",  # delete_matching is not a suitable choice for parallel writes.
+    SaveMode.IGNORE: "overwrite_or_ignore",
+    SaveMode.ERROR: "error",
+}
+
+FILE_FORMAT = "parquet"
+
+# These args are part of https://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html#pyarrow.fs.FileSystem.open_output_stream
+# and are not supported by ParquetDatasink.
+UNSUPPORTED_OPEN_STREAM_ARGS = {"path", "buffer", "metadata"}
+
+# https://arrow.apache.org/docs/python/generated/pyarrow.dataset.write_dataset.html
+ARROW_DEFAULT_MAX_ROWS_PER_GROUP = 1024 * 1024
+
+logger = logging.getLogger(__name__)
+
+
+def choose_row_group_limits(
+    row_group_size: Optional[int],
+    min_rows_per_file: Optional[int],
+    max_rows_per_file: Optional[int],
+) -> tuple[Optional[int], Optional[int], Optional[int]]:
+    """
+    Configure `min_rows_per_group`, `max_rows_per_group`, `max_rows_per_file` parameters of Pyarrow's `write_dataset` API based on Ray Data's configuration
+
+    Returns
+    -------
+    (min_rows_per_group, max_rows_per_group, max_rows_per_file)
+    """
+
+    if (
+        row_group_size is None
+        and min_rows_per_file is None
+        and max_rows_per_file is None
+    ):
+        return None, None, None
+
+    elif row_group_size is None:
+        # No explicit row group size provided. We are defaulting to
+        # either the caller's min_rows_per_file or max_rows_per_file limits
+        # or Arrow's defaults
+        min_rows_per_group, max_rows_per_group, max_rows_per_file = (
+            min_rows_per_file,
+            max_rows_per_file,
+            max_rows_per_file,
+        )
+
+        # If min_rows_per_group is provided and max_rows_per_group is not,
+        # and min_rows_per_group is greater than Arrow's default max_rows_per_group,
+        # we set max_rows_per_group to min_rows_per_group to avoid creating too many row groups.
+        if (
+            min_rows_per_group is not None
+            and max_rows_per_group is None
+            and min_rows_per_group > ARROW_DEFAULT_MAX_ROWS_PER_GROUP
+        ):
+            max_rows_per_group, max_rows_per_file = (
+                min_rows_per_group,
+                min_rows_per_group,
+            )
+
+        return min_rows_per_group, max_rows_per_group, max_rows_per_file
+
+    elif row_group_size is not None and (
+        min_rows_per_file is None or max_rows_per_file is None
+    ):
+        return row_group_size, row_group_size, max_rows_per_file
+
+    else:
+        # Clamp the requested `row_group_size` so that it is
+        # * no smaller than `min_rows_per_file` (`lower`)
+        # * no larger than `max_rows_per_file` (or Arrow's default cap) (`upper`)
+        # This keeps each row-group within the per-file limits while staying
+        # as close as possible to the requested size.
+        clamped_group_size = max(
+            min_rows_per_file, min(row_group_size, max_rows_per_file)
+        )
+        return clamped_group_size, clamped_group_size, max_rows_per_file
+
+
+class ParquetDatasink(_FileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        partition_cols: Optional[List[str]] = None,
+        arrow_parquet_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        arrow_parquet_args: Optional[Dict[str, Any]] = None,
+        min_rows_per_file: Optional[int] = None,
+        max_rows_per_file: Optional[int] = None,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        try_create_dir: bool = True,
+        open_stream_args: Optional[Dict[str, Any]] = None,
+        filename_provider: Optional[FilenameProvider] = None,
+        dataset_uuid: Optional[str] = None,
+        mode: SaveMode = SaveMode.APPEND,
+    ):
+        if arrow_parquet_args_fn is None:
+            arrow_parquet_args_fn = lambda: {}  # noqa: E731
+
+        if arrow_parquet_args is None:
+            arrow_parquet_args = {}
+
+        self.arrow_parquet_args_fn = arrow_parquet_args_fn
+        self.arrow_parquet_args = arrow_parquet_args
+        self.min_rows_per_file = min_rows_per_file
+        self.max_rows_per_file = max_rows_per_file
+        self.partition_cols = partition_cols
+
+        if self.min_rows_per_file is not None and self.max_rows_per_file is not None:
+            assert (
+                self.min_rows_per_file <= self.max_rows_per_file
+            ), "min_rows_per_file must be less than or equal to max_rows_per_file"
+
+        if open_stream_args is not None:
+            intersecting_keys = UNSUPPORTED_OPEN_STREAM_ARGS.intersection(
+                set(open_stream_args.keys())
+            )
+            if intersecting_keys:
+                logger.warning(
+                    "open_stream_args contains unsupported arguments: %s. These arguments "
+                    "are not supported by ParquetDatasink. They will be ignored.",
+                    intersecting_keys,
+                )
+
+            if "compression" in open_stream_args:
+                self.arrow_parquet_args["compression"] = open_stream_args["compression"]
+
+        super().__init__(
+            path,
+            filesystem=filesystem,
+            try_create_dir=try_create_dir,
+            open_stream_args=open_stream_args,
+            filename_provider=filename_provider,
+            dataset_uuid=dataset_uuid,
+            file_format=FILE_FORMAT,
+            mode=mode,
+        )
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        import pyarrow as pa
+
+        blocks = list(blocks)
+
+        if all(BlockAccessor.for_block(block).num_rows() == 0 for block in blocks):
+            return
+
+        blocks = [
+            block for block in blocks if BlockAccessor.for_block(block).num_rows() > 0
+        ]
+
+        filename = self.filename_provider.get_filename_for_block(
+            blocks[0], ctx.kwargs[WRITE_UUID_KWARG_NAME], ctx.task_idx, 0
+        )
+        write_kwargs = _resolve_kwargs(
+            self.arrow_parquet_args_fn, **self.arrow_parquet_args
+        )
+        user_schema = write_kwargs.pop("schema", None)
+
+        def write_blocks_to_path():
+            tables = [BlockAccessor.for_block(block).to_arrow() for block in blocks]
+            if user_schema is None:
+                output_schema = pa.unify_schemas([table.schema for table in tables])
+            else:
+                output_schema = user_schema
+
+            self._write_parquet_files(
+                tables,
+                filename,
+                output_schema,
+                ctx.kwargs[WRITE_UUID_KWARG_NAME],
+                write_kwargs,
+            )
+
+        logger.debug(f"Writing {filename} file to {self.path}.")
+
+        call_with_retry(
+            write_blocks_to_path,
+            description=f"write '{filename}' to '{self.path}'",
+            match=self._data_context.retried_io_errors,
+            max_attempts=WRITE_FILE_MAX_ATTEMPTS,
+            max_backoff_s=WRITE_FILE_RETRY_MAX_BACKOFF_SECONDS,
+        )
+
+    def _get_basename_template(self, filename: str, write_uuid: str) -> str:
+        # Check if write_uuid is present in filename, add if missing
+        if write_uuid not in filename and self.mode == SaveMode.APPEND:
+            raise ValueError(
+                f"Write UUID '{write_uuid}' missing from filename template '{filename}'. This could result in files being overwritten."
+                f"Modify your FileNameProvider implementation to include the `write_uuid` into the filename template or change your write mode to SaveMode.OVERWRITE. "
+            )
+        # Check if filename is already templatized
+        if "{i}" in filename:
+            # Filename is already templatized, but may need file extension
+            if FILE_FORMAT not in filename:
+                # Add file extension to templatized filename
+                basename_template = f"{filename}.{FILE_FORMAT}"
+            else:
+                # Already has extension, use as-is
+                basename_template = filename
+        elif FILE_FORMAT not in filename:
+            # No extension and not templatized, add extension and template
+            basename_template = f"{filename}-{{i}}.{FILE_FORMAT}"
+        else:
+            # TODO(@goutamvenkat-anyscale): Add a warning if you pass in a custom
+            # filename provider and it isn't templatized.
+            # Use pathlib.Path to properly handle filenames with dots
+            filename_path = Path(filename)
+            stem = filename_path.stem  # filename without extension
+            assert "." not in stem, "Filename should not contain a dot"
+            suffix = filename_path.suffix  # extension including the dot
+            basename_template = f"{stem}-{{i}}{suffix}"
+        return basename_template
+
+    def _write_parquet_files(
+        self,
+        tables: List["pyarrow.Table"],
+        filename: str,
+        output_schema: "pyarrow.Schema",
+        write_uuid: str,
+        write_kwargs: Dict[str, Any],
+    ) -> None:
+        import pyarrow.dataset as ds
+
+        # Make every incoming batch conform to the final schema *before* writing
+        for idx, table in enumerate(tables):
+            if output_schema and not table.schema.equals(output_schema):
+                table = table.cast(output_schema)
+            tables[idx] = table
+
+        row_group_size = write_kwargs.pop("row_group_size", None)
+
+        existing_data_behavior = EXISTING_DATA_BEHAVIOR_MAP.get(
+            self.mode, "overwrite_or_ignore"
+        )
+
+        (
+            min_rows_per_group,
+            max_rows_per_group,
+            max_rows_per_file,
+        ) = choose_row_group_limits(
+            row_group_size,
+            min_rows_per_file=self.min_rows_per_file,
+            max_rows_per_file=self.max_rows_per_file,
+        )
+
+        basename_template = self._get_basename_template(filename, write_uuid)
+
+        ds.write_dataset(
+            data=tables,
+            base_dir=self.path,
+            schema=output_schema,
+            basename_template=basename_template,
+            filesystem=self.filesystem,
+            partitioning=self.partition_cols,
+            format=FILE_FORMAT,
+            existing_data_behavior=existing_data_behavior,
+            partitioning_flavor="hive",
+            use_threads=True,
+            min_rows_per_group=min_rows_per_group,
+            max_rows_per_group=max_rows_per_group,
+            max_rows_per_file=max_rows_per_file,
+            file_options=ds.ParquetFileFormat().make_write_options(**write_kwargs),
+        )
+
+    @property
+    def min_rows_per_write(self) -> Optional[int]:
+        return self.min_rows_per_file
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..5bab245c3f4a2566893de7baec19e179a2888312
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/parquet_datasource.py
@@ -0,0 +1,988 @@
+import logging
+import math
+import os
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+from packaging.version import parse as parse_version
+
+import ray
+from ray._private.arrow_utils import get_pyarrow_version
+from ray.data._internal.arrow_block import (
+    _BATCH_SIZE_PRESERVING_STUB_COL_NAME,
+    ArrowBlockAccessor,
+)
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.util import (
+    RetryingPyFileSystem,
+    _check_pyarrow_version,
+    _is_local_scheme,
+    iterate_with_retry,
+)
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource import Datasource
+from ray.data.datasource.datasource import ReadTask
+from ray.data.datasource.file_based_datasource import FileShuffleConfig
+from ray.data.datasource.file_meta_provider import (
+    FileMetadataProvider,
+    _handle_read_os_error,
+    _list_files,
+)
+from ray.data.datasource.partitioning import (
+    PartitionDataType,
+    Partitioning,
+    PathPartitionFilter,
+    PathPartitionParser,
+)
+from ray.data.datasource.path_util import (
+    _resolve_paths_and_filesystem,
+)
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    import pyarrow
+    from pyarrow.dataset import ParquetFileFragment
+
+
+logger = logging.getLogger(__name__)
+
+
+MIN_PYARROW_TO_BATCHES_READAHEAD = parse_version("10.0.0")
+
+
+# The `num_cpus` for each metadata prefetching task.
+# Default to 0.5 instead of 1 because it is cheaper than normal read task.
+NUM_CPUS_FOR_META_FETCH_TASK = 0.5
+
+# The number of rows to read per batch. This is sized to generate 10MiB batches
+# for rows about 1KiB in size.
+DEFAULT_PARQUET_READER_ROW_BATCH_SIZE = 10_000
+FILE_READING_RETRY = 8
+
+# The default size multiplier for reading Parquet data source in Arrow.
+# Parquet data format is encoded with various encoding techniques (such as
+# dictionary, RLE, delta), so Arrow in-memory representation uses much more memory
+# compared to Parquet encoded representation. Parquet file statistics only record
+# encoded (i.e. uncompressed) data size information.
+#
+# To estimate real-time in-memory data size, Datasets will try to estimate the
+# correct inflation ratio from Parquet to Arrow, using this constant as the default
+# value for safety. See https://github.com/ray-project/ray/pull/26516 for more context.
+PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT = 5
+
+# The lower bound size to estimate Parquet encoding ratio.
+PARQUET_ENCODING_RATIO_ESTIMATE_LOWER_BOUND = 1
+
+# The percentage of files (1% by default) to be sampled from the dataset to estimate
+# Parquet encoding ratio.
+PARQUET_ENCODING_RATIO_ESTIMATE_SAMPLING_RATIO = 0.01
+
+# The minimal and maximal number of file samples to take from the dataset to estimate
+# Parquet encoding ratio.
+# This is to restrict `PARQUET_ENCODING_RATIO_ESTIMATE_SAMPLING_RATIO` within the
+# proper boundary.
+PARQUET_ENCODING_RATIO_ESTIMATE_MIN_NUM_SAMPLES = 2
+PARQUET_ENCODING_RATIO_ESTIMATE_MAX_NUM_SAMPLES = 10
+
+# The number of rows to read from each file for sampling. Try to keep it low to avoid
+# reading too much data into memory.
+PARQUET_ENCODING_RATIO_ESTIMATE_NUM_ROWS = 1024
+
+
+class _ParquetFragment:
+    """This wrapper class is created to avoid utilizing `ParquetFileFragment` original
+    serialization protocol that actually does network RPCs during serialization
+    (to fetch actual parquet metadata)"""
+
+    def __init__(self, f: "ParquetFileFragment", file_size: int):
+        self._fragment = f
+        self._file_size = file_size
+
+    @property
+    def file_size(self) -> int:
+        return self._file_size
+
+    @property
+    def original(self) -> "ParquetFileFragment":
+        return self._fragment
+
+    def __reduce__(self):
+        return _ParquetFragment.make_fragment, (
+            self._fragment.format,
+            self._fragment.path,
+            self._fragment.filesystem,
+            self._fragment.partition_expression,
+            self._file_size,
+        )
+
+    @staticmethod
+    def make_fragment(format, path, filesystem, partition_expression, file_size):
+        fragment = format.make_fragment(path, filesystem, partition_expression)
+        return _ParquetFragment(fragment, file_size)
+
+
+def check_for_legacy_tensor_type(schema):
+    """Check for the legacy tensor extension type and raise an error if found.
+
+    Ray Data uses an extension type to represent tensors in Arrow tables. Previously,
+    the extension type extended `PyExtensionType`. However, this base type can expose
+    users to arbitrary code execution. To prevent this, we don't load the type by
+    default.
+    """
+    import pyarrow as pa
+
+    for name, type in zip(schema.names, schema.types):
+        if isinstance(type, pa.UnknownExtensionType) and isinstance(
+            type, pa.PyExtensionType
+        ):
+            raise RuntimeError(
+                f"Ray Data couldn't infer the type of column '{name}'. This might mean "
+                "you're trying to read data written with an older version of Ray. "
+                "Reading data written with older versions of Ray might expose you to "
+                "arbitrary code execution. To try reading the data anyway, set "
+                "`RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE=1` on *all* nodes."
+                "To learn more, see https://github.com/ray-project/ray/issues/41314."
+            )
+
+
+class ParquetDatasource(Datasource):
+    """Parquet datasource, for reading and writing Parquet files.
+
+    The primary difference from ParquetBulkDatasource is that this uses
+    PyArrow's `ParquetDataset` abstraction for dataset reads, and thus offers
+    automatic Arrow dataset schema inference and row count collection at the
+    cost of some potential performance and/or compatibility penalties.
+    """
+
+    _FILE_EXTENSIONS = ["parquet"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        *,
+        columns: Optional[List[str]] = None,
+        dataset_kwargs: Optional[Dict[str, Any]] = None,
+        to_batch_kwargs: Optional[Dict[str, Any]] = None,
+        _block_udf: Optional[Callable[[Block], Block]] = None,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        schema: Optional[Union[type, "pyarrow.lib.Schema"]] = None,
+        meta_provider: Optional[FileMetadataProvider] = None,
+        partition_filter: PathPartitionFilter = None,
+        partitioning: Optional[Partitioning] = Partitioning("hive"),
+        shuffle: Union[Literal["files"], None] = None,
+        include_paths: bool = False,
+        file_extensions: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        _check_pyarrow_version()
+
+        self._supports_distributed_reads = not _is_local_scheme(paths)
+        if not self._supports_distributed_reads and ray.util.client.ray.is_connected():
+            raise ValueError(
+                "Because you're using Ray Client, read tasks scheduled on the Ray "
+                "cluster can't access your local files. To fix this issue, store "
+                "files in cloud storage or a distributed filesystem like NFS."
+            )
+
+        self._local_scheduling = None
+        if not self._supports_distributed_reads:
+            from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+            self._local_scheduling = NodeAffinitySchedulingStrategy(
+                ray.get_runtime_context().get_node_id(), soft=False
+            )
+        # Need this property for lineage tracking
+        self._source_paths = paths
+        paths, self._filesystem = _resolve_paths_and_filesystem(paths, filesystem)
+        filesystem = RetryingPyFileSystem.wrap(
+            self._filesystem,
+            retryable_errors=DataContext.get_current().retried_io_errors,
+        )
+
+        listed_files = _list_files(
+            paths,
+            filesystem,
+            partition_filter=partition_filter,
+            file_extensions=file_extensions,
+        )
+
+        if listed_files:
+            paths, file_sizes = zip(*listed_files)
+        else:
+            paths, file_sizes = [], []
+
+        if dataset_kwargs is not None:
+            logger.warning(
+                "Please note that `ParquetDatasource.__init__`s `dataset_kwargs` "
+                "is a deprecated parameter and will be removed in the future."
+            )
+        else:
+            dataset_kwargs = {}
+
+        if "partitioning" in dataset_kwargs:
+            raise ValueError(
+                "The 'partitioning' parameter isn't supported in 'dataset_kwargs'. "
+                "Use the top-level 'partitioning' parameter instead."
+            )
+
+        # This datasource manually adds partition data at the Ray Data-level. To avoid
+        # duplicating the partition data, we disable PyArrow's partitioning.
+        dataset_kwargs["partitioning"] = None
+
+        # NOTE: ParquetDataset only accepts list of paths, hence we need to convert
+        #       it to a list
+        pq_ds = get_parquet_dataset(list(paths), filesystem, dataset_kwargs)
+
+        # Users can pass both data columns and partition columns in the 'columns'
+        # argument. To prevent PyArrow from complaining about missing columns, we
+        # separate the partition columns from the data columns. When we read the
+        # fragments, we pass the data columns to PyArrow and add the partition
+        # columns manually.
+        data_columns, partition_columns = None, None
+        if columns is not None:
+            data_columns, partition_columns = _infer_data_and_partition_columns(
+                columns, pq_ds.fragments[0], partitioning
+            )
+
+        if to_batch_kwargs is None:
+            to_batch_kwargs = {}
+
+        # NOTE: Store the custom serialized `ParquetFileFragment` to avoid unexpected
+        # network calls when `_ParquetDatasourceReader` is serialized. See
+        # `_SerializedFragment()` implementation for more details.
+        self._pq_fragments = [
+            _ParquetFragment(fragment, file_size)
+            for fragment, file_size in zip(pq_ds.fragments, file_sizes)
+        ]
+        self._pq_paths = [p.path for p in pq_ds.fragments]
+        self._block_udf = _block_udf
+        self._to_batches_kwargs = to_batch_kwargs
+        # Store as projection_map (identity mapping if columns specified, None otherwise)
+        # Note: Empty list [] means no columns, None means all columns
+        if data_columns is None:
+            self._projection_map = None
+        else:
+            self._projection_map = {col: col for col in data_columns}
+        self._partition_columns = partition_columns
+        self._read_schema = schema
+        self._file_schema = pq_ds.schema
+        self._partition_schema = _get_partition_columns_schema(
+            partitioning, self._pq_paths
+        )
+        self._file_metadata_shuffler = None
+        self._include_paths = include_paths
+        self._partitioning = partitioning
+        if shuffle == "files":
+            self._file_metadata_shuffler = np.random.default_rng()
+        elif isinstance(shuffle, FileShuffleConfig):
+            self._file_metadata_shuffler = np.random.default_rng(shuffle.seed)
+
+        # Sample small number of parquet files to estimate
+        #   - Encoding ratio: ratio of file size on disk to approximate expected
+        #     size of the corresponding block in memory
+        #   - Default batch-size: number of rows to be read from a file at a time,
+        #     used to limit amount of memory pressure
+        sampled_fragments = _sample_fragments(
+            self._pq_fragments,
+        )
+
+        sampled_file_infos = _fetch_file_infos(
+            sampled_fragments,
+            columns=self._get_data_columns(),
+            schema=schema,
+            local_scheduling=self._local_scheduling,
+        )
+
+        self._encoding_ratio = _estimate_files_encoding_ratio(
+            sampled_fragments,
+            sampled_file_infos,
+        )
+
+        self._default_batch_size = _estimate_reader_batch_size(
+            sampled_file_infos, DataContext.get_current().target_max_block_size
+        )
+
+    def estimate_inmemory_data_size(self) -> int:
+        # In case of empty projections no data will be read
+        if self._projection_map == {}:
+            return 0
+
+        return self._estimate_in_mem_size(self._pq_fragments)
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        # NOTE: We override the base class FileBasedDatasource.get_read_tasks()
+        # method in order to leverage pyarrow's ParquetDataset abstraction,
+        # which simplifies partitioning logic. We still use
+        # FileBasedDatasource's write side, however.
+        if self._file_metadata_shuffler is not None:
+            files_metadata = list(zip(self._pq_fragments, self._pq_paths))
+            shuffled_files_metadata = [
+                files_metadata[i]
+                for i in self._file_metadata_shuffler.permutation(len(files_metadata))
+            ]
+            pq_fragments, pq_paths = list(map(list, zip(*shuffled_files_metadata)))
+        else:
+            pq_fragments, pq_paths = (
+                self._pq_fragments,
+                self._pq_paths,
+            )
+
+        # Derive expected target schema of the blocks being read
+        target_schema = _derive_schema(
+            self._read_schema,
+            file_schema=self._file_schema,
+            partition_schema=self._partition_schema,
+            projected_columns=self.get_current_projection(),
+            _block_udf=self._block_udf,
+        )
+
+        read_tasks = []
+        filter_expr = (
+            self._predicate_expr.to_pyarrow()
+            if self._predicate_expr is not None
+            else None
+        )
+
+        for fragments, paths in zip(
+            np.array_split(pq_fragments, parallelism),
+            np.array_split(pq_paths, parallelism),
+        ):
+            if len(fragments) <= 0:
+                continue
+
+            meta = BlockMetadata(
+                num_rows=None,
+                size_bytes=self._estimate_in_mem_size(fragments),
+                input_files=paths,
+                exec_stats=None,
+            )
+
+            (
+                block_udf,
+                to_batches_kwargs,
+                default_read_batch_size_rows,
+                data_columns,
+                data_columns_rename_map,
+                partition_columns,
+                read_schema,
+                include_paths,
+                partitioning,
+            ) = (
+                self._block_udf,
+                self._to_batches_kwargs,
+                self._default_batch_size,
+                self._get_data_columns(),
+                self.get_column_renames(),
+                self._partition_columns,
+                self._read_schema,
+                self._include_paths,
+                self._partitioning,
+            )
+
+            read_tasks.append(
+                ReadTask(
+                    lambda f=fragments: read_fragments(
+                        block_udf,
+                        to_batches_kwargs,
+                        default_read_batch_size_rows,
+                        data_columns,
+                        data_columns_rename_map,
+                        partition_columns,
+                        read_schema,
+                        f,
+                        include_paths,
+                        partitioning,
+                        filter_expr,
+                    ),
+                    meta,
+                    schema=target_schema,
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+
+        return read_tasks
+
+    def get_name(self):
+        """Return a human-readable name for this datasource.
+
+        This will be used as the names of the read tasks.
+        """
+        return "Parquet"
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        return self._supports_distributed_reads
+
+    def supports_projection_pushdown(self) -> bool:
+        return True
+
+    def supports_predicate_pushdown(self) -> bool:
+        return True
+
+    def get_current_projection(self) -> Optional[List[str]]:
+        """Override to include partition columns in addition to data columns."""
+        # NOTE: In case there's no projection both file and partition columns
+        #       will be none
+        data_columns = self._get_data_columns()
+        if data_columns is None and self._partition_columns is None:
+            return None
+
+        return (data_columns or []) + (self._partition_columns or [])
+
+    def _estimate_in_mem_size(self, fragments: List[_ParquetFragment]) -> int:
+        in_mem_size = sum([f.file_size for f in fragments]) * self._encoding_ratio
+
+        return round(in_mem_size)
+
+
+def read_fragments(
+    block_udf: Callable[[Block], Optional[Block]],
+    to_batches_kwargs: Dict[str, Any],
+    default_read_batch_size_rows: Optional[int],
+    data_columns: Optional[List[str]],
+    data_columns_rename_map: Optional[Dict[str, str]],
+    partition_columns: Optional[List[str]],
+    schema: Optional[Union[type, "pyarrow.lib.Schema"]],
+    fragments: List[_ParquetFragment],
+    include_paths: bool,
+    partitioning: Partitioning,
+    filter_expr: Optional["pyarrow.dataset.Expression"] = None,
+) -> Iterator["pyarrow.Table"]:
+    # This import is necessary to load the tensor extension type.
+    from ray.data.extensions.tensor_extension import ArrowTensorType  # noqa
+
+    # Ensure that we're reading at least one dataset fragment.
+    assert len(fragments) > 0
+
+    logger.debug(f"Reading {len(fragments)} parquet fragments")
+    for fragment in fragments:
+        # S3 can raise transient errors during iteration, and PyArrow doesn't expose a
+        # way to retry specific batches.
+        ctx = ray.data.DataContext.get_current()
+        for table in iterate_with_retry(
+            lambda: _read_batches_from(
+                fragment.original,
+                schema=schema,
+                data_columns=data_columns,
+                data_columns_rename_map=data_columns_rename_map,
+                partition_columns=partition_columns,
+                partitioning=partitioning,
+                include_path=include_paths,
+                filter_expr=filter_expr,
+                batch_size=default_read_batch_size_rows,
+                to_batches_kwargs=to_batches_kwargs,
+            ),
+            "reading batches",
+            match=ctx.retried_io_errors,
+        ):
+            # If the table is empty, drop it.
+            if table.num_rows > 0:
+                if block_udf is not None:
+                    yield block_udf(table)
+                else:
+                    yield table
+
+
+def _read_batches_from(
+    fragment: "ParquetFileFragment",
+    *,
+    schema: "pyarrow.Schema",
+    data_columns: Optional[List[str]],
+    data_columns_rename_map: Optional[Dict[str, str]],
+    partition_columns: Optional[List[str]],
+    partitioning: Partitioning,
+    filter_expr: Optional["pyarrow.dataset.Expression"] = None,
+    batch_size: Optional[int] = None,
+    include_path: bool = False,
+    use_threads: bool = False,
+    to_batches_kwargs: Optional[Dict[str, Any]] = None,
+) -> Iterable["pyarrow.Table"]:
+    """Get an iterable of batches from a parquet fragment."""
+
+    import pyarrow as pa
+
+    from ray.data.datasource.datasource import _DatasourceProjectionPushdownMixin
+
+    # Copy to avoid modifying passed in arg
+    to_batches_kwargs = dict(to_batches_kwargs or {})
+
+    # NOTE: Passed in kwargs overrides always take precedence
+    # TODO deprecate to_batches_kwargs
+    use_threads = to_batches_kwargs.pop("use_threads", use_threads)
+    # TODO: We should deprecate filter through the read_parquet API and only allow through dataset.filter()
+    filter_from_kwargs = to_batches_kwargs.pop("filter", None)
+    if filter_from_kwargs is not None:
+        filter_expr = (
+            filter_from_kwargs
+            if filter_expr is None
+            else filter_expr & filter_from_kwargs
+        )
+    # NOTE: Arrow's ``to_batches`` expects ``batch_size`` as an int
+    if batch_size is not None:
+        to_batches_kwargs.setdefault("batch_size", batch_size)
+
+    partition_col_values = _parse_partition_column_values(
+        fragment, partition_columns, partitioning
+    )
+
+    def _generate_tables() -> "pa.Table":
+        """Inner generator that yields tables without renaming."""
+        try:
+            for batch in fragment.to_batches(
+                columns=data_columns,
+                filter=filter_expr,
+                schema=schema,
+                use_threads=use_threads,
+                **to_batches_kwargs,
+            ):
+                table = pa.Table.from_batches([batch])
+
+                if include_path:
+                    table = ArrowBlockAccessor.for_block(table).fill_column(
+                        "path", fragment.path
+                    )
+
+                if partition_col_values:
+                    table = _add_partitions_to_table(partition_col_values, table)
+
+                # ``ParquetFileFragment.to_batches`` returns ``RecordBatch``,
+                # which could have empty projection (ie ``num_columns`` == 0)
+                # while having non-empty rows (ie ``num_rows`` > 0), which
+                # could occur when list of requested columns is empty.
+                #
+                # However, when ``RecordBatches`` are concatenated using
+                # ``pyarrow.concat_tables`` it will return a single ``Table``
+                # with 0 columns and therefore 0 rows (since ``Table``s number of
+                # rows is determined as the length of its columns).
+                #
+                # To avoid running into this pitfall, we introduce a stub column
+                # holding just nulls to maintain invariance of the number of rows.
+                #
+                # NOTE: There's no impact from this as the binary size of the
+                #       extra column is basically 0
+                if table.num_columns == 0 and table.num_rows > 0:
+                    table = table.append_column(
+                        _BATCH_SIZE_PRESERVING_STUB_COL_NAME, pa.nulls(table.num_rows)
+                    )
+
+                yield table
+
+        except pa.lib.ArrowInvalid as e:
+            error_message = str(e)
+            if (
+                "No match for FieldRef.Name" in error_message
+                and filter_expr is not None
+            ):
+                filename = os.path.basename(fragment.path)
+                file_columns = set(fragment.physical_schema.names)
+                raise RuntimeError(
+                    f"Filter expression: '{filter_expr}' failed on parquet "
+                    f"file: '{filename}' with columns: {file_columns}"
+                )
+            raise
+
+    # Apply renames to all tables from the generator
+    yield from _DatasourceProjectionPushdownMixin._apply_rename_to_tables(
+        _generate_tables(), data_columns_rename_map
+    )
+
+
+def _parse_partition_column_values(
+    fragment: "ParquetFileFragment",
+    partition_columns: Optional[List[str]],
+    partitioning: Partitioning,
+):
+    partitions = {}
+
+    if partitioning is not None:
+        parse = PathPartitionParser(partitioning)
+        partitions = parse(fragment.path)
+
+    # Filter out partitions that aren't in the user-specified columns list.
+    if partition_columns is not None:
+        partitions = {
+            field_name: value
+            for field_name, value in partitions.items()
+            if field_name in partition_columns
+        }
+
+    return partitions
+
+
+def _fetch_parquet_file_info(
+    fragment: _ParquetFragment,
+    *,
+    columns: Optional[List[str]],
+    schema: Optional["pyarrow.Schema"],
+) -> Optional["_ParquetFileInfo"]:
+    # If the fragment has no row groups, it's an empty or metadata-only file.
+    # Skip it by returning empty sample info.
+    #
+    # NOTE: Accessing `ParquetFileFragment.metadata` does fetch a parquet footer
+    #       from storage
+    metadata = fragment.original.metadata
+
+    if metadata.num_row_groups == 0:
+        return None
+
+    # Only sample the first row group.
+    row_group_fragment = fragment.original.subset(row_group_ids=[0])
+    batch_size = max(
+        min(
+            row_group_fragment.metadata.num_rows,
+            PARQUET_ENCODING_RATIO_ESTIMATE_NUM_ROWS,
+        ),
+        1,
+    )
+
+    to_batches_kwargs = {}
+
+    if get_pyarrow_version() >= MIN_PYARROW_TO_BATCHES_READAHEAD:
+        # Limit prefetching to just 1 batch
+        to_batches_kwargs["batch_readahead"] = 1
+
+    batches_iter = row_group_fragment.to_batches(
+        columns=columns,
+        schema=schema,
+        batch_size=batch_size,
+        **to_batches_kwargs,
+    )
+
+    avg_row_size: Optional[int] = None
+    # Use first batch non-empty batch to estimate the avg size of the
+    # row in-memory
+    for batch in batches_iter:
+        if batch.num_rows > 0:
+            avg_row_size = math.ceil(batch.nbytes / batch.num_rows)
+            break
+
+    return _ParquetFileInfo(
+        avg_row_in_mem_bytes=avg_row_size,
+        metadata=metadata,
+    )
+
+
+@dataclass
+class _ParquetFileInfo:
+    # Estimated avg byte size of a row (in-memory)
+    avg_row_in_mem_bytes: Optional[int]
+    # Corresponding file metadata
+    metadata: "pyarrow._parquet.FileMetaData"
+
+    def estimate_in_memory_bytes(self) -> Optional[int]:
+        if self.avg_row_in_mem_bytes is None:
+            return None
+
+        return self.avg_row_in_mem_bytes * self.metadata.num_rows
+
+
+def _estimate_files_encoding_ratio(
+    fragments: List[_ParquetFragment],
+    file_infos: List[_ParquetFileInfo],
+) -> float:
+    """Return an estimate of the Parquet files encoding ratio.
+
+    To avoid OOMs, it is safer to return an over-estimate than an underestimate.
+    """
+    if not DataContext.get_current().decoding_size_estimation:
+        return PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+    assert len(file_infos) == len(fragments)
+
+    # Estimate size of the rows in a file in memory
+    estimated_in_mem_size_arr = [
+        fi.estimate_in_memory_bytes() if fi is not None else None for fi in file_infos
+    ]
+
+    file_size_arr = [f.file_size for f in fragments]
+
+    estimated_encoding_ratios = [
+        float(in_mem_size) / file_size
+        for in_mem_size, file_size in zip(estimated_in_mem_size_arr, file_size_arr)
+        if file_size > 0 and in_mem_size is not None
+    ]
+
+    # Return default estimate of 5 if all sampled files turned out to be empty
+    if not estimated_encoding_ratios:
+        return PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT
+
+    estimated_ratio = np.mean(estimated_encoding_ratios)
+
+    logger.info(f"Estimated parquet encoding ratio is {estimated_ratio:.3f}.")
+
+    return max(estimated_ratio, PARQUET_ENCODING_RATIO_ESTIMATE_LOWER_BOUND)
+
+
+def _fetch_file_infos(
+    sampled_fragments: List[_ParquetFragment],
+    *,
+    columns: Optional[List[str]],
+    schema: Optional["pyarrow.Schema"],
+    local_scheduling: Optional[bool],
+) -> List[Optional[_ParquetFileInfo]]:
+    fetch_file_info = cached_remote_fn(_fetch_parquet_file_info)
+    futures = []
+
+    for fragment in sampled_fragments:
+        # Sample the first rows batch in i-th file.
+        # Use SPREAD scheduling strategy to avoid packing many sampling tasks on
+        # same machine to cause OOM issue, as sampling can be memory-intensive.
+        futures.append(
+            fetch_file_info.options(
+                scheduling_strategy=local_scheduling
+                or DataContext.get_current().scheduling_strategy,
+                # Retry in case of transient errors during sampling.
+                retry_exceptions=[OSError],
+            ).remote(
+                fragment,
+                columns=columns,
+                schema=schema,
+            )
+        )
+
+    sample_bar = ProgressBar("Parquet dataset sampling", len(futures), unit="file")
+    file_infos = sample_bar.fetch_until_complete(futures)
+    sample_bar.close()
+
+    return file_infos
+
+
+def _estimate_reader_batch_size(
+    file_infos: List[Optional[_ParquetFileInfo]], target_block_size: Optional[int]
+) -> Optional[int]:
+    if target_block_size is None:
+        return None
+
+    avg_num_rows_per_block = [
+        target_block_size / fi.avg_row_in_mem_bytes
+        for fi in file_infos
+        if (
+            fi is not None
+            and fi.avg_row_in_mem_bytes is not None
+            and fi.avg_row_in_mem_bytes > 0
+        )
+    ]
+
+    if not avg_num_rows_per_block:
+        return DEFAULT_PARQUET_READER_ROW_BATCH_SIZE
+
+    estimated_batch_size: int = max(math.ceil(np.mean(avg_num_rows_per_block)), 1)
+
+    logger.info(f"Estimated parquet reader batch size at {estimated_batch_size} rows")
+
+    return estimated_batch_size
+
+
+def get_parquet_dataset(paths, filesystem, dataset_kwargs):
+    import pyarrow.parquet as pq
+
+    # If you pass a list containing a single directory path to `ParquetDataset`, PyArrow
+    # errors with 'IsADirectoryError: Path ... points to a directory, but only file
+    # paths are supported'. To avoid this, we pass the directory path directly.
+    if len(paths) == 1:
+        paths = paths[0]
+
+    try:
+        dataset = pq.ParquetDataset(
+            paths,
+            **dataset_kwargs,
+            filesystem=filesystem,
+        )
+    except OSError as e:
+        _handle_read_os_error(e, paths)
+
+    return dataset
+
+
+def _sample_fragments(
+    fragments: List[_ParquetFragment],
+) -> List[_ParquetFragment]:
+    if not fragments:
+        return []
+
+    target_num_samples = math.ceil(
+        len(fragments) * PARQUET_ENCODING_RATIO_ESTIMATE_SAMPLING_RATIO
+    )
+
+    target_num_samples = max(
+        min(target_num_samples, PARQUET_ENCODING_RATIO_ESTIMATE_MAX_NUM_SAMPLES),
+        PARQUET_ENCODING_RATIO_ESTIMATE_MIN_NUM_SAMPLES,
+    )
+
+    # Make sure number of samples doesn't exceed total # of files
+    target_num_samples = min(target_num_samples, len(fragments))
+
+    # Evenly distributed to choose which file to sample, to avoid biased prediction
+    # if data is skewed.
+    pivots = np.linspace(0, len(fragments) - 1, target_num_samples).astype(int)
+
+    return [fragments[idx] for idx in pivots.tolist()]
+
+
+def _add_partitions_to_table(
+    partition_col_values: Dict[str, PartitionDataType], table: "pyarrow.Table"
+) -> "pyarrow.Table":
+
+    for partition_col, value in partition_col_values.items():
+        field_index = table.schema.get_field_index(partition_col)
+        if field_index == -1:
+            table = BlockAccessor.for_block(table).fill_column(partition_col, value)
+        elif log_once(f"duplicate_partition_field_{partition_col}"):
+            logger.warning(
+                f"The partition field '{partition_col}' also exists in the Parquet "
+                f"file. Ray Data will default to using the value in the Parquet file."
+            )
+
+    return table
+
+
+def _get_partition_columns_schema(
+    partitioning: Partitioning,
+    file_paths: List[str],
+) -> "pyarrow.Schema":
+    """Return a new schema with partition fields added.
+
+    This function infers the partition fields from the first file path in the dataset.
+    """
+    import pyarrow as pa
+
+    # If the dataset is empty, we can't infer the partitioning
+    if len(file_paths) == 0:
+        return pa.schema([])
+    # If the dataset isn't partitioned, there's no partition schema
+    elif partitioning is None:
+        return pa.schema([])
+
+    first_path = file_paths[0]
+
+    fields = []
+
+    parser = PathPartitionParser(partitioning)
+    partitions = parser(first_path)
+    for field_name in partitions:
+        if field_name in partitioning.field_types:
+            field_type = pa.from_numpy_dtype(partitioning.field_types[field_name])
+        else:
+            field_type = pa.string()
+
+        # Without this check, we would add the same partition field multiple times,
+        # which silently fails when asking for `pa.field()`.
+        fields.append(pa.field(field_name, field_type))
+
+    return pa.schema(fields)
+
+
+def _derive_schema(
+    read_schema: Optional["pyarrow.Schema"],
+    *,
+    file_schema: "pyarrow.Schema",
+    partition_schema: Optional["pyarrow.Schema"],
+    projected_columns: Optional[List[str]],
+    _block_udf,
+) -> "pyarrow.Schema":
+    """Derives target schema for read operation"""
+
+    import pyarrow as pa
+
+    # Use target read schema if provided
+    if read_schema is not None:
+        target_schema = read_schema
+    else:
+        file_schema_fields = list(file_schema)
+        partition_schema_fields = (
+            list(partition_schema) if partition_schema is not None else []
+        )
+
+        # Otherwise, fallback to file + partitioning schema by default
+        target_schema = pa.schema(
+            fields=(
+                file_schema_fields
+                + [
+                    f
+                    for f in partition_schema_fields
+                    # Ignore fields from partition schema overlapping with
+                    # file's schema
+                    if file_schema.get_field_index(f.name) == -1
+                ]
+            ),
+            metadata=file_schema.metadata,
+        )
+
+    # Project schema if necessary
+    if projected_columns is not None:
+        target_schema = pa.schema(
+            [target_schema.field(column) for column in projected_columns],
+            target_schema.metadata,
+        )
+
+    if _block_udf is not None:
+        # Try to infer dataset schema by passing dummy table through UDF.
+        dummy_table = target_schema.empty_table()
+        try:
+            target_schema = _block_udf(dummy_table).schema.with_metadata(
+                target_schema.metadata
+            )
+        except Exception:
+            logger.debug(
+                "Failed to infer schema of dataset by passing dummy table "
+                "through UDF due to the following exception:",
+                exc_info=True,
+            )
+
+    check_for_legacy_tensor_type(target_schema)
+
+    return target_schema
+
+
+def _infer_data_and_partition_columns(
+    user_specified_columns: List[str],
+    fragment: "ParquetFileFragment",
+    partitioning: Optional[Partitioning],
+) -> Tuple[List[str], List[str]]:
+    """Infer which columns are in the files and which columns are partition columns.
+
+    This function uses the schema and path of the first file to infer what columns
+    represent.
+
+    Args:
+        user_specified_columns: A list of column names that the user specified.
+        fragment: The first fragment in the dataset.
+        partitioning: The partitioning scheme used to partition the data.
+
+    Returns:
+        A tuple of lists of column names. The first list contains the columns that are
+        in the file, and the second list contains the columns that are partition
+        columns.
+    """
+    data_columns = [
+        column
+        for column in user_specified_columns
+        if column in fragment.physical_schema.names
+    ]
+    if partitioning is not None:
+        parse = PathPartitionParser(partitioning)
+        partitions = parse(fragment.path)
+        partition_columns = [
+            column for column in user_specified_columns if column in partitions
+        ]
+    else:
+        partition_columns = []
+
+    return data_columns, partition_columns
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/range_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/range_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..06cd98768c2cb50a2b0155ea0f423024e52bddfa
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/range_datasource.py
@@ -0,0 +1,145 @@
+import builtins
+import functools
+from typing import Iterable, List, Optional, Tuple
+
+import numpy as np
+
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource import Datasource, ReadTask
+
+
+class RangeDatasource(Datasource):
+    """An example datasource that generates ranges of numbers from [0..n)."""
+
+    def __init__(
+        self,
+        n: int,
+        block_format: str = "arrow",
+        tensor_shape: Tuple = (1,),
+        column_name: Optional[str] = None,
+    ):
+        self._n = int(n)
+        self._block_format = block_format
+        self._tensor_shape = tensor_shape
+        self._column_name = column_name
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        if self._block_format == "tensor":
+            element_size = int(np.prod(self._tensor_shape))
+        else:
+            element_size = 1
+        return 8 * self._n * element_size
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+    ) -> List[ReadTask]:
+        if self._n == 0:
+            return []
+
+        read_tasks: List[ReadTask] = []
+        n = self._n
+        block_format = self._block_format
+        tensor_shape = self._tensor_shape
+        block_size = max(1, n // parallelism)
+        # TODO(swang): This target block size may not match the driver's
+        # context if it was overridden. Set target max block size during
+        # optimizer stage to fix this.
+        ctx = DataContext.get_current()
+        if ctx.target_max_block_size is None:
+            # If target_max_block_size is ``None``, treat it as unlimited and
+            # avoid further splitting.
+            target_rows_per_block = n  # whole block in one shot
+        else:
+            row_size_bytes = self.estimate_inmemory_data_size() // self._n
+            row_size_bytes = max(row_size_bytes, 1)
+            target_rows_per_block = max(1, ctx.target_max_block_size // row_size_bytes)
+
+        # Example of a read task. In a real datasource, this would pull data
+        # from an external system instead of generating dummy data.
+        def make_block(start: int, count: int) -> Block:
+            if block_format == "arrow":
+                import pyarrow as pa
+
+                return pa.Table.from_arrays(
+                    [np.arange(start, start + count)],
+                    names=[self._column_name or "value"],
+                )
+            elif block_format == "tensor":
+                import pyarrow as pa
+
+                tensor = np.ones(tensor_shape, dtype=np.int64) * np.expand_dims(
+                    np.arange(start, start + count),
+                    tuple(range(1, 1 + len(tensor_shape))),
+                )
+                return BlockAccessor.batch_to_block(
+                    {self._column_name: tensor} if self._column_name else tensor
+                )
+            else:
+                return list(builtins.range(start, start + count))
+
+        def make_blocks(
+            start: int, count: int, target_rows_per_block: int
+        ) -> Iterable[Block]:
+            while count > 0:
+                num_rows = min(count, target_rows_per_block)
+                yield make_block(start, num_rows)
+                start += num_rows
+                count -= num_rows
+
+        if block_format == "tensor":
+            element_size = int(np.prod(tensor_shape))
+        else:
+            element_size = 1
+
+        i = 0
+        while i < n:
+            count = min(block_size, n - i)
+            meta = BlockMetadata(
+                num_rows=count,
+                size_bytes=8 * count * element_size,
+                input_files=None,
+                exec_stats=None,
+            )
+            read_tasks.append(
+                ReadTask(
+                    lambda i=i, count=count: make_blocks(
+                        i, count, target_rows_per_block
+                    ),
+                    meta,
+                    schema=self._schema,
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+            i += block_size
+
+        return read_tasks
+
+    @functools.cached_property
+    def _schema(self):
+        if self._n == 0:
+            return None
+
+        if self._block_format == "arrow":
+            _check_pyarrow_version()
+            import pyarrow as pa
+
+            schema = pa.Table.from_pydict({self._column_name or "value": [0]}).schema
+        elif self._block_format == "tensor":
+            _check_pyarrow_version()
+            import pyarrow as pa
+
+            tensor = np.ones(self._tensor_shape, dtype=np.int64) * np.expand_dims(
+                np.arange(0, 10), tuple(range(1, 1 + len(self._tensor_shape)))
+            )
+            schema = BlockAccessor.batch_to_block(
+                {self._column_name: tensor} if self._column_name else tensor
+            ).schema
+        elif self._block_format == "list":
+            schema = int
+        else:
+            raise ValueError("Unsupported block type", self._block_format)
+        return schema
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/sql_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/sql_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..f773a186c58c3facff499301a5b1bd5eff97f7c6
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/sql_datasink.py
@@ -0,0 +1,35 @@
+from typing import Callable, Iterable
+
+from ray.data._internal.datasource.sql_datasource import Connection, _connect
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+
+class SQLDatasink(Datasink[None]):
+
+    _MAX_ROWS_PER_WRITE = 128
+
+    def __init__(self, sql: str, connection_factory: Callable[[], Connection]):
+        self.sql = sql
+        self.connection_factory = connection_factory
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        with _connect(self.connection_factory) as cursor:
+            for block in blocks:
+                block_accessor = BlockAccessor.for_block(block)
+
+                values = []
+                for row in block_accessor.iter_rows(public_row_format=False):
+                    values.append(tuple(row.values()))
+                    assert len(values) <= self._MAX_ROWS_PER_WRITE, len(values)
+                    if len(values) == self._MAX_ROWS_PER_WRITE:
+                        cursor.executemany(self.sql, values)
+                        values = []
+
+                if values:
+                    cursor.executemany(self.sql, values)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/sql_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/sql_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ac48f0b343f2de82013312e233d1af10545f20b
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/sql_datasource.py
@@ -0,0 +1,191 @@
+import logging
+import math
+from contextlib import contextmanager
+from typing import Any, Callable, Iterable, Iterator, List, Optional
+
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+Connection = Any  # A Python DB API2-compliant `Connection` object.
+Cursor = Any  # A Python DB API2-compliant `Cursor` object.
+
+logger = logging.getLogger(__name__)
+
+
+def _cursor_to_block(cursor) -> Block:
+    import pyarrow as pa
+
+    rows = cursor.fetchall()
+    # Each `column_description` is a 7-element sequence. The first element is the column
+    # name. To learn more, read https://peps.python.org/pep-0249/#description.
+    columns = [column_description[0] for column_description in cursor.description]
+    pydict = {column: [row[i] for row in rows] for i, column in enumerate(columns)}
+    return pa.Table.from_pydict(pydict)
+
+
+def _check_connection_is_dbapi2_compliant(connection) -> None:
+    for attr in "close", "commit", "cursor":
+        if not hasattr(connection, attr):
+            raise ValueError(
+                "Your `connection_factory` created a `Connection` object without a "
+                f"{attr!r} method, but this method is required by the Python DB API2 "
+                "specification. Check that your database connector is DB API2-"
+                "compliant. To learn more, read https://peps.python.org/pep-0249/."
+            )
+
+
+def _check_cursor_is_dbapi2_compliant(cursor) -> None:
+    # These aren't all the methods required by the specification, but it's all the ones
+    # we care about.
+    for attr in "execute", "executemany", "fetchone", "fetchall", "description":
+        if not hasattr(cursor, attr):
+            raise ValueError(
+                "Your database connector created a `Cursor` object without a "
+                f"{attr!r} method, but this method is required by the Python DB API2 "
+                "specification. Check that your database connector is DB API2-"
+                "compliant. To learn more, read https://peps.python.org/pep-0249/."
+            )
+
+
+@contextmanager
+def _connect(connection_factory: Callable[[], Connection]) -> Iterator[Cursor]:
+    connection = connection_factory()
+    _check_connection_is_dbapi2_compliant(connection)
+
+    try:
+        cursor = connection.cursor()
+        _check_cursor_is_dbapi2_compliant(cursor)
+        yield cursor
+        connection.commit()
+    except Exception:
+        # `rollback` is optional since not all databases provide transaction support.
+        try:
+            connection.rollback()
+        except Exception as e:
+            # Each connector implements its own `NotSupportError` class, so we check
+            # the exception's name instead of using `isinstance`.
+            if (
+                isinstance(e, AttributeError)
+                or e.__class__.__name__ == "NotSupportedError"
+            ):
+                pass
+        raise
+    finally:
+        connection.close()
+
+
+class SQLDatasource(Datasource):
+    MIN_ROWS_PER_READ_TASK = 50
+
+    def __init__(
+        self,
+        sql: str,
+        connection_factory: Callable[[], Connection],
+        shard_hash_fn: str,
+        shard_keys: Optional[List[str]] = None,
+    ):
+        self.sql = sql
+        if shard_keys and len(shard_keys) > 1:
+            self.shard_keys = f"CONCAT({','.join(shard_keys)})"
+        elif shard_keys and len(shard_keys) == 1:
+            self.shard_keys = f"{shard_keys[0]}"
+        else:
+            self.shard_keys = None
+        self.shard_hash_fn = shard_hash_fn
+        self.connection_factory = connection_factory
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def supports_sharding(self, parallelism: int) -> bool:
+        """Check if database supports sharding with MOD/ABS/CONCAT operations.
+
+        Returns:
+            bool: True if sharding is supported, False otherwise.
+        """
+        if parallelism <= 1 or self.shard_keys is None:
+            return False
+
+        # Test if database supports required operations (MOD, ABS, MD5, CONCAT)
+        # by executing a sample query
+        hash_fn = self.shard_hash_fn
+        query = (
+            f"SELECT COUNT(1) FROM ({self.sql}) as T"
+            f" WHERE MOD(ABS({hash_fn}({self.shard_keys})), {parallelism}) = 0"
+        )
+        try:
+            with _connect(self.connection_factory) as cursor:
+                cursor.execute(query)
+            return True
+        except Exception as e:
+            logger.info(f"Database does not support sharding: {str(e)}.")
+            return False
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        def fallback_read_fn() -> Iterable[Block]:
+            """Read all data in a single block when sharding is not supported."""
+            with _connect(self.connection_factory) as cursor:
+                cursor.execute(self.sql)
+                return [_cursor_to_block(cursor)]
+
+        # Check if sharding is supported by the database first
+        # If not, fall back to reading all data in a single task without counting rows
+        if not self.supports_sharding(parallelism):
+            logger.info(
+                "Sharding is not supported. "
+                "Falling back to reading all data in a single task."
+            )
+            metadata = BlockMetadata(None, None, None, None)
+            return [ReadTask(fallback_read_fn, metadata)]
+
+        # Only perform the expensive COUNT(*) query if sharding is supported
+        num_rows_total = self._get_num_rows()
+
+        if num_rows_total == 0:
+            return []
+
+        parallelism = min(
+            parallelism, math.ceil(num_rows_total / self.MIN_ROWS_PER_READ_TASK)
+        )
+        num_rows_per_block = num_rows_total // parallelism
+        num_blocks_with_extra_row = num_rows_total % parallelism
+
+        tasks = []
+        for i in range(parallelism):
+            num_rows = num_rows_per_block
+            if i < num_blocks_with_extra_row:
+                num_rows += 1
+            read_fn = self._create_parallel_read_fn(i, parallelism)
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=None,
+                input_files=None,
+                exec_stats=None,
+            )
+            tasks.append(
+                ReadTask(read_fn, metadata, per_task_row_limit=per_task_row_limit)
+            )
+
+        return tasks
+
+    def _get_num_rows(self) -> int:
+        with _connect(self.connection_factory) as cursor:
+            cursor.execute(f"SELECT COUNT(*) FROM ({self.sql}) as T")
+            return cursor.fetchone()[0]
+
+    def _create_parallel_read_fn(self, task_id: int, parallelism: int):
+        hash_fn = self.shard_hash_fn
+        query = (
+            f"SELECT * FROM ({self.sql}) as T "
+            f"WHERE MOD(ABS({hash_fn}({self.shard_keys})), {parallelism}) = {task_id}"
+        )
+
+        def read_fn() -> Iterable[Block]:
+            with _connect(self.connection_factory) as cursor:
+                cursor.execute(query)
+                block = _cursor_to_block(cursor)
+                return [block]
+
+        return read_fn
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/text_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/text_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..42e844e597c0d6085b7c9db4b1b61d31fc3f410a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/text_datasource.py
@@ -0,0 +1,42 @@
+from typing import TYPE_CHECKING, Iterator, List
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class TextDatasource(FileBasedDatasource):
+    """Text datasource, for reading and writing text files."""
+
+    _COLUMN_NAME = "text"
+
+    def __init__(
+        self,
+        paths: List[str],
+        *,
+        drop_empty_lines: bool = False,
+        encoding: str = "utf-8",
+        **file_based_datasource_kwargs
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self.drop_empty_lines = drop_empty_lines
+        self.encoding = encoding
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        data = f.readall()
+
+        builder = DelegatingBlockBuilder()
+
+        lines = data.decode(self.encoding).splitlines()
+        for line in lines:
+            if self.drop_empty_lines and line.strip() == "":
+                continue
+            item = {self._COLUMN_NAME: line}
+            builder.add(item)
+
+        block = builder.build()
+        yield block
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/tfrecords_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/tfrecords_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba89f26b9d66ad3a5f79cb3ecb9a12d2fc1791d3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/tfrecords_datasink.py
@@ -0,0 +1,205 @@
+import struct
+from typing import TYPE_CHECKING, Dict, Iterable, Optional, Union
+
+import numpy as np
+
+from .tfrecords_datasource import _get_single_true_type
+from ray.data._internal.util import _check_import
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+if TYPE_CHECKING:
+    import pyarrow
+    import tensorflow as tf
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+
+class TFRecordDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        *,
+        tf_schema: Optional["schema_pb2.Schema"] = None,
+        file_format: str = "tar",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format=file_format, **file_datasink_kwargs)
+
+        _check_import(self, module="crc32c", package="crc32c")
+
+        self.tf_schema = tf_schema
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        arrow_table = block.to_arrow()
+
+        # It seems like TFRecords are typically row-based,
+        # https://www.tensorflow.org/tutorials/load_data/tfrecord#writing_a_tfrecord_file_2
+        # so we must iterate through the rows of the block,
+        # serialize to tf.train.Example proto, and write to file.
+
+        examples = _convert_arrow_table_to_examples(arrow_table, self.tf_schema)
+
+        # Write each example to the arrow file in the TFRecord format.
+        for example in examples:
+            _write_record(file, example)
+
+
+def _convert_arrow_table_to_examples(
+    arrow_table: "pyarrow.Table",
+    tf_schema: Optional["schema_pb2.Schema"] = None,
+) -> Iterable["tf.train.Example"]:
+    import tensorflow as tf
+
+    schema_dict = {}
+    # Convert user-specified schema into dict for convenient mapping
+    if tf_schema is not None:
+        for schema_feature in tf_schema.feature:
+            schema_dict[schema_feature.name] = schema_feature.type
+
+    # Serialize each row[i] of the block to a tf.train.Example and yield it.
+    for i in range(arrow_table.num_rows):
+        # First, convert row[i] to a dictionary.
+        features: Dict[str, "tf.train.Feature"] = {}
+        for name in arrow_table.column_names:
+            if tf_schema is not None and name not in schema_dict:
+                raise ValueError(
+                    f"Found extra unexpected feature {name} "
+                    f"not in specified schema: {tf_schema}"
+                )
+            schema_feature_type = schema_dict.get(name)
+            features[name] = _value_to_feature(
+                arrow_table[name][i],
+                schema_feature_type,
+            )
+
+        # Convert the dictionary to an Example proto.
+        proto = tf.train.Example(features=tf.train.Features(feature=features))
+
+        yield proto
+
+
+def _value_to_feature(
+    value: Union["pyarrow.Scalar", "pyarrow.Array"],
+    schema_feature_type: Optional["schema_pb2.FeatureType"] = None,
+) -> "tf.train.Feature":
+    import pyarrow as pa
+    import tensorflow as tf
+
+    if isinstance(value, pa.ListScalar):
+        # Use the underlying type of the ListScalar's value in
+        # determining the output feature's data type.
+        value_type = value.type.value_type
+        value = value.as_py()
+    else:
+        value_type = value.type
+        value = value.as_py()
+        if value is None:
+            value = []
+        else:
+            value = [value]
+
+    underlying_value_type = {
+        "bytes": pa.types.is_binary(value_type),
+        "string": pa.types.is_string(value_type),
+        "float": pa.types.is_floating(value_type),
+        "int": pa.types.is_integer(value_type),
+    }
+    assert sum(bool(value) for value in underlying_value_type.values()) <= 1
+
+    if schema_feature_type is not None:
+        try:
+            from tensorflow_metadata.proto.v0 import schema_pb2
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "To use TensorFlow schemas, please install "
+                "the tensorflow-metadata package."
+            )
+        specified_feature_type = {
+            "bytes": schema_feature_type == schema_pb2.FeatureType.BYTES
+            and not underlying_value_type["string"],
+            "string": schema_feature_type == schema_pb2.FeatureType.BYTES
+            and underlying_value_type["string"],
+            "float": schema_feature_type == schema_pb2.FeatureType.FLOAT,
+            "int": schema_feature_type == schema_pb2.FeatureType.INT,
+        }
+
+        und_type = _get_single_true_type(underlying_value_type)
+        spec_type = _get_single_true_type(specified_feature_type)
+        if und_type is not None and und_type != spec_type:
+            raise ValueError(
+                "Schema field type mismatch during write: specified type is "
+                f"{spec_type}, but underlying type is {und_type}",
+            )
+        # Override the underlying value type with the type in the user-specified schema.
+        underlying_value_type = specified_feature_type
+
+    if underlying_value_type["int"]:
+        return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
+    if underlying_value_type["float"]:
+        return tf.train.Feature(float_list=tf.train.FloatList(value=value))
+    if underlying_value_type["bytes"]:
+        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+    if underlying_value_type["string"]:
+        value = [v.encode() for v in value]  # casting to bytes
+        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+    if pa.types.is_null(value_type):
+        raise ValueError(
+            "Unable to infer type from partially missing column. "
+            "Try setting read parallelism = 1, or use an input data source which "
+            "explicitly specifies the schema."
+        )
+    raise ValueError(
+        f"Value is of type {value_type}, "
+        "which we cannot convert to a supported tf.train.Feature storage type "
+        "(bytes, float, or int)."
+    )
+
+
+# Adapted from https://github.com/vahidk/tfrecord/blob/74b2d24a838081356d993ec0e147eaf59ccd4c84/tfrecord/writer.py#L57-L72  # noqa: E501
+#
+# MIT License
+#
+# Copyright (c) 2020 Vahid Kazemi
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+def _write_record(
+    file: "pyarrow.NativeFile",
+    example: "tf.train.Example",
+) -> None:
+    record = example.SerializeToString()
+    length = len(record)
+    length_bytes = struct.pack("<Q", length)
+    file.write(length_bytes)
+    file.write(_masked_crc(length_bytes))
+    file.write(record)
+    file.write(_masked_crc(record))
+
+
+def _masked_crc(data: bytes) -> bytes:
+    """CRC checksum."""
+    import crc32c
+
+    mask = 0xA282EAD8
+    crc = crc32c.crc32(data)
+    masked = ((crc >> 15) | (crc << 17)) + mask
+    masked = np.uint32(masked & np.iinfo(np.uint32).max)
+    masked_bytes = struct.pack("<I", masked)
+    return masked_bytes
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/tfrecords_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/tfrecords_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..925f45657e1d2ce6207978095ea77d7ee7e7a34a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/tfrecords_datasource.py
@@ -0,0 +1,439 @@
+import logging
+import struct
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, Iterable, Iterator, List, Optional, Union
+
+import pyarrow
+
+from ray.air.util.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data.aggregate import AggregateFn
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import tensorflow as tf
+    from tensorflow_metadata.proto.v0 import schema_pb2
+
+    from ray.data.dataset import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+@dataclass
+class TFXReadOptions:
+    """
+    Specifies read options when reading TFRecord files with TFX.
+    """
+
+    # An int representing the number of consecutive elements of
+    # this dataset to combine in a single batch when tfx-bsl is used to read
+    # the tfrecord files.
+    batch_size: int = 2048
+
+    # Toggles the schema inference applied; applicable
+    # only if tfx-bsl is used and tf_schema argument is missing.
+    # Defaults to True.
+    auto_infer_schema: bool = True
+
+
+class TFRecordDatasource(FileBasedDatasource):
+    """TFRecord datasource, for reading and writing TFRecord files."""
+
+    _FILE_EXTENSIONS = ["tfrecords"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        tf_schema: Optional["schema_pb2.Schema"] = None,
+        tfx_read_options: Optional["TFXReadOptions"] = None,
+        **file_based_datasource_kwargs,
+    ):
+        """
+        Args:
+            tf_schema: Optional TensorFlow Schema which is used to explicitly set
+                the schema of the underlying Dataset.
+            tfx_read_options: Optional options for enabling reading tfrecords
+                using tfx-bsl.
+
+        """
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self._tf_schema = tf_schema
+        self._tfx_read_options = tfx_read_options
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        if self._tfx_read_options:
+            yield from self._tfx_read_stream(f, path)
+        else:
+            yield from self._default_read_stream(f, path)
+
+    def _default_read_stream(
+        self, f: "pyarrow.NativeFile", path: str
+    ) -> Iterator[Block]:
+        import tensorflow as tf
+        from google.protobuf.message import DecodeError
+
+        for record in _read_records(f, path):
+            example = tf.train.Example()
+            try:
+                example.ParseFromString(record)
+            except DecodeError as e:
+                raise ValueError(
+                    "`TFRecordDatasource` failed to parse `tf.train.Example` "
+                    f"record in '{path}'. This error can occur if your TFRecord "
+                    f"file contains a message type other than `tf.train.Example`: {e}"
+                )
+
+            yield pyarrow_table_from_pydict(
+                _convert_example_to_dict(example, self._tf_schema)
+            )
+
+    def _tfx_read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import tensorflow as tf
+        from tfx_bsl.cc.tfx_bsl_extension.coders import ExamplesToRecordBatchDecoder
+
+        full_path = self._resolve_full_path(path)
+
+        compression = (self._open_stream_args or {}).get("compression", None)
+
+        if compression:
+            compression = compression.upper()
+
+        tf_schema_string = (
+            self._tf_schema.SerializeToString() if self._tf_schema else None
+        )
+
+        decoder = ExamplesToRecordBatchDecoder(tf_schema_string)
+        exception_thrown = None
+        try:
+            for record in tf.data.TFRecordDataset(
+                full_path, compression_type=compression
+            ).batch(self._tfx_read_options.batch_size):
+                yield _cast_large_list_to_list(
+                    pyarrow.Table.from_batches([decoder.DecodeBatch(record.numpy())])
+                )
+        except Exception as error:
+            logger.exception(f"Failed to read TFRecord file {full_path}")
+            exception_thrown = error
+
+        # we need to do this hack were we raise an exception outside of the
+        # except block because tensorflow DataLossError is unpickable, and
+        # even if we raise a runtime error, ray keeps information about the
+        # original error, which makes it unpickable still.
+        if exception_thrown:
+            raise RuntimeError(f"Failed to read TFRecord file {full_path}.")
+
+    def _resolve_full_path(self, relative_path):
+        from ray.data._internal.util import RetryingPyFileSystem
+
+        filesystem = self._filesystem
+        if isinstance(filesystem, RetryingPyFileSystem):
+            filesystem = filesystem.unwrap()
+        if isinstance(filesystem, pyarrow.fs.S3FileSystem):
+            return f"s3://{relative_path}"
+        if isinstance(filesystem, pyarrow.fs.GcsFileSystem):
+            return f"gs://{relative_path}"
+        if isinstance(filesystem, pyarrow.fs.HadoopFileSystem):
+            return f"hdfs:///{relative_path}"
+        if isinstance(filesystem, pyarrow.fs.PyFileSystem):
+            protocol = filesystem.handler.fs.protocol
+            if isinstance(protocol, list) or isinstance(protocol, tuple):
+                protocol = protocol[0]
+            if protocol == "gcs":
+                protocol = "gs"
+            return f"{protocol}://{relative_path}"
+
+        return relative_path
+
+
+def _convert_example_to_dict(
+    example: "tf.train.Example",
+    tf_schema: Optional["schema_pb2.Schema"],
+) -> Dict[str, pyarrow.Array]:
+    record = {}
+    schema_dict = {}
+    # Convert user-specified schema into dict for convenient mapping
+    if tf_schema is not None:
+        for schema_feature in tf_schema.feature:
+            schema_dict[schema_feature.name] = schema_feature.type
+
+    for feature_name, feature in example.features.feature.items():
+        if tf_schema is not None and feature_name not in schema_dict:
+            raise ValueError(
+                f"Found extra unexpected feature {feature_name} "
+                f"not in specified schema: {tf_schema}"
+            )
+        schema_feature_type = schema_dict.get(feature_name)
+        record[feature_name] = _get_feature_value(feature, schema_feature_type)
+    return record
+
+
+def _get_single_true_type(dct) -> str:
+    """Utility function for getting the single key which has a `True` value in
+    a dict. Used to filter a dict of `{field_type: is_valid}` to get
+    the field type from a schema or data source."""
+    filtered_types = iter([_type for _type in dct if dct[_type]])
+    # In the case where there are no keys with a `True` value, return `None`
+    return next(filtered_types, None)
+
+
+def _get_feature_value(
+    feature: "tf.train.Feature",
+    schema_feature_type: Optional["schema_pb2.FeatureType"] = None,
+) -> pyarrow.Array:
+    import pyarrow as pa
+
+    underlying_feature_type = {
+        "bytes": feature.HasField("bytes_list"),
+        "float": feature.HasField("float_list"),
+        "int": feature.HasField("int64_list"),
+    }
+    # At most one of `bytes_list`, `float_list`, and `int64_list`
+    # should contain values. If none contain data, this indicates
+    # an empty feature value.
+    assert sum(bool(value) for value in underlying_feature_type.values()) <= 1
+
+    if schema_feature_type is not None:
+        try:
+            from tensorflow_metadata.proto.v0 import schema_pb2
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "To use TensorFlow schemas, please install "
+                "the tensorflow-metadata package."
+            )
+        # If a schema is specified, compare to the underlying type
+        specified_feature_type = {
+            "bytes": schema_feature_type == schema_pb2.FeatureType.BYTES,
+            "float": schema_feature_type == schema_pb2.FeatureType.FLOAT,
+            "int": schema_feature_type == schema_pb2.FeatureType.INT,
+        }
+        und_type = _get_single_true_type(underlying_feature_type)
+        spec_type = _get_single_true_type(specified_feature_type)
+        if und_type is not None and und_type != spec_type:
+            raise ValueError(
+                "Schema field type mismatch during read: specified type is "
+                f"{spec_type}, but underlying type is {und_type}",
+            )
+        # Override the underlying value type with the type in the user-specified schema.
+        underlying_feature_type = specified_feature_type
+
+    if underlying_feature_type["bytes"]:
+        value = feature.bytes_list.value
+        type_ = pa.binary()
+    elif underlying_feature_type["float"]:
+        value = feature.float_list.value
+        type_ = pa.float32()
+    elif underlying_feature_type["int"]:
+        value = feature.int64_list.value
+        type_ = pa.int64()
+    else:
+        value = []
+        type_ = pa.null()
+    value = list(value)
+    if len(value) == 1 and schema_feature_type is None:
+        # Use the value itself if the features contains a single value.
+        # This is to give better user experience when writing preprocessing UDF on
+        # these single-value lists.
+        value = value[0]
+    else:
+        # If the feature value is empty and no type is specified in the user-provided
+        # schema, set the type to null for now to allow pyarrow to construct a valid
+        # Array; later, infer the type from other records which have non-empty values
+        # for the feature.
+        if len(value) == 0:
+            type_ = pa.null()
+        type_ = pa.list_(type_)
+    return pa.array([value], type=type_)
+
+
+# Adapted from https://github.com/vahidk/tfrecord/blob/74b2d24a838081356d993ec0e147eaf59ccd4c84/tfrecord/reader.py#L16-L96  # noqa: E501
+#
+# MIT License
+#
+# Copyright (c) 2020 Vahid Kazemi
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+def _read_records(
+    file: "pyarrow.NativeFile",
+    path: str,
+) -> Iterable[memoryview]:
+    """
+    Read records from TFRecord file.
+
+    A TFRecord file contains a sequence of records. The file can only be read
+    sequentially. Each record is stored in the following formats:
+        uint64 length
+        uint32 masked_crc32_of_length
+        byte   data[length]
+        uint32 masked_crc32_of_data
+
+    See https://www.tensorflow.org/tutorials/load_data/tfrecord#tfrecords_format_details
+    for more details.
+    """
+    length_bytes = bytearray(8)
+    crc_bytes = bytearray(4)
+    datum_bytes = bytearray(1024 * 1024)
+    row_count = 0
+    while True:
+        try:
+            # Read "length" field.
+            num_length_bytes_read = file.readinto(length_bytes)
+            if num_length_bytes_read == 0:
+                break
+            elif num_length_bytes_read != 8:
+                raise ValueError(
+                    "Failed to read the length of record data. Expected 8 bytes but "
+                    "got {num_length_bytes_read} bytes."
+                )
+
+            # Read "masked_crc32_of_length" field.
+            num_length_crc_bytes_read = file.readinto(crc_bytes)
+            if num_length_crc_bytes_read != 4:
+                raise ValueError(
+                    "Failed to read the length of CRC-32C hashes. Expected 4 bytes "
+                    "but got {num_length_crc_bytes_read} bytes."
+                )
+
+            # Read "data[length]" field.
+            (data_length,) = struct.unpack("<Q", length_bytes)
+            if data_length > len(datum_bytes):
+                datum_bytes = datum_bytes.zfill(int(data_length * 1.5))
+            datum_bytes_view = memoryview(datum_bytes)[:data_length]
+            num_datum_bytes_read = file.readinto(datum_bytes_view)
+            if num_datum_bytes_read != data_length:
+                raise ValueError(
+                    f"Failed to read the record. Exepcted {data_length} bytes but got "
+                    f"{num_datum_bytes_read} bytes."
+                )
+
+            # Read "masked_crc32_of_data" field.
+            # TODO(chengsu): ideally we should check CRC-32C against the actual data.
+            num_crc_bytes_read = file.readinto(crc_bytes)
+            if num_crc_bytes_read != 4:
+                raise ValueError(
+                    "Failed to read the CRC-32C hashes. Expected 4 bytes but got "
+                    f"{num_crc_bytes_read} bytes."
+                )
+
+            # Return the data.
+            yield datum_bytes_view
+
+            row_count += 1
+            data_length = None
+        except Exception as e:
+            error_message = (
+                f"Failed to read TFRecord file {path}. Please ensure that the "
+                f"TFRecord file has correct format. Already read {row_count} rows."
+            )
+            if data_length is not None:
+                error_message += f" Byte size of current record data is {data_length}."
+            raise RuntimeError(error_message) from e
+
+
+def _cast_large_list_to_list(batch: pyarrow.Table):
+    """
+    This function transform pyarrow.large_list into list and pyarrow.large_binary into
+    pyarrow.binary so that all types resulting from the tfrecord_datasource are usable
+    with dataset.to_tf().
+    """
+    old_schema = batch.schema
+    fields = {}
+
+    for column_name in old_schema.names:
+        field_type = old_schema.field(column_name).type
+        if type(field_type) is pyarrow.lib.LargeListType:
+            value_type = field_type.value_type
+
+            if value_type == pyarrow.large_binary():
+                value_type = pyarrow.binary()
+
+            fields[column_name] = pyarrow.list_(value_type)
+        elif field_type == pyarrow.large_binary():
+            fields[column_name] = pyarrow.binary()
+        else:
+            fields[column_name] = old_schema.field(column_name)
+
+    new_schema = pyarrow.schema(fields)
+    return batch.cast(new_schema)
+
+
+def _infer_schema_and_transform(dataset: "Dataset"):
+    list_sizes = dataset.aggregate(_MaxListSize(dataset.schema().names))
+
+    return dataset.map_batches(
+        _unwrap_single_value_lists,
+        fn_kwargs={"col_lengths": list_sizes["max_list_size"]},
+        batch_format="pyarrow",
+    )
+
+
+def _unwrap_single_value_lists(batch: pyarrow.Table, col_lengths: Dict[str, int]):
+    """
+    This function will transfrom the dataset converting list types that always
+    contain single values to thery underlying data type
+    (i.e. pyarrow.int64() and pyarrow.float64())
+    """
+    columns = {}
+
+    for col in col_lengths:
+        value_type = batch[col].type.value_type
+
+        if col_lengths[col] == 1:
+            if batch[col]:
+                columns[col] = pyarrow.array(
+                    [x.as_py()[0] if x.as_py() else None for x in batch[col]],
+                    type=value_type,
+                )
+        else:
+            columns[col] = batch[col]
+
+    return pyarrow.table(columns)
+
+
+class _MaxListSize(AggregateFn):
+    def __init__(self, columns: List[str]):
+        self._columns = columns
+        super().__init__(
+            init=self._init,
+            merge=self._merge,
+            accumulate_row=self._accumulate_row,
+            finalize=lambda a: a,
+            name="max_list_size",
+        )
+
+    def _init(self, k: str):
+        return {col: 0 for col in self._columns}
+
+    def _merge(self, acc1: Dict[str, int], acc2: Dict[str, int]):
+        merged = {}
+        for col in self._columns:
+            merged[col] = max(acc1[col], acc2[col])
+
+        return merged
+
+    def _accumulate_row(self, acc: Dict[str, int], row: "pd.Series"):
+        for k in row:
+            value = row[k]
+            if value:
+                acc[k] = max(len(value), acc[k])
+
+        return acc
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/torch_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/torch_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..69145f995b1dc657637f3f9ab7fbe55d24c87603
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/torch_datasource.py
@@ -0,0 +1,77 @@
+from typing import TYPE_CHECKING, Optional
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data.block import BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import torch
+
+
+TORCH_DATASOURCE_READER_BATCH_SIZE = 32
+
+
+class TorchDatasource(Datasource):
+    """Torch datasource, for reading from `Torch
+    datasets <https://pytorch.org/docs/stable/data.html/>`_.
+    This datasource implements a streaming read using a single read task.
+    """
+
+    def __init__(
+        self,
+        dataset: "torch.utils.data.Dataset",
+    ):
+        self._dataset = dataset
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ):
+        assert parallelism == 1
+
+        meta = BlockMetadata(
+            # Note: avoid len(self._dataset) because it will trigger
+            # iterating through IterableDataset, which can cause OOM.
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        read_task = ReadTask(
+            lambda subset=self._dataset: _read_subset(
+                subset,
+            ),
+            metadata=meta,
+            per_task_row_limit=per_task_row_limit,
+        )
+
+        return [read_task]
+
+    def estimate_inmemory_data_size(self):
+        return None
+
+
+def _read_subset(subset: "torch.utils.data.Subset"):
+    batch = []
+
+    # Get items from dataset based on its type
+    if hasattr(subset, "__iter__"):
+        # IterableDataset: Use the iterator directly
+        items = subset
+    else:
+        # Map-style dataset: Respect __len__
+        items = (subset[i] for i in range(len(subset)))
+
+    # Process items in batches
+    for item in items:
+        batch.append(item)
+        if len(batch) == TORCH_DATASOURCE_READER_BATCH_SIZE:
+            builder = DelegatingBlockBuilder()
+            builder.add_batch({"item": batch})
+            yield builder.build()
+            batch.clear()
+
+    # Handle any remaining items
+    if len(batch) > 0:
+        builder = DelegatingBlockBuilder()
+        builder.add_batch({"item": batch})
+        yield builder.build()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/uc_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/uc_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b81dddfa5b8e55d58838c20b0dd413482e3a0d4
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/uc_datasource.py
@@ -0,0 +1,195 @@
+import atexit
+import os
+import tempfile
+from typing import Any, Callable, Dict, Optional
+
+import requests
+
+import ray
+
+_FILE_FORMAT_TO_RAY_READER = {
+    "delta": "read_delta",
+    "parquet": "read_parquet",
+}
+
+
+class UnityCatalogConnector:
+    """
+    Load a Unity Catalog table or files into a Ray Dataset, handling cloud credentials automatically.
+
+    Currently only supports Databricks-managed Unity Catalog
+
+    Supported formats: delta, parquet.
+    Supports AWS, Azure, and GCP with automatic credential handoff.
+    """
+
+    def __init__(
+        self,
+        *,
+        base_url: str,
+        token: str,
+        table_full_name: str,
+        region: Optional[str] = None,
+        data_format: Optional[str] = "delta",
+        operation: str = "READ",
+        ray_init_kwargs: Optional[Dict] = None,
+        reader_kwargs: Optional[Dict] = None,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.token = token
+        self.table_full_name = table_full_name
+        self.data_format = data_format.lower() if data_format else None
+        self.region = region
+        self.operation = operation
+        self.ray_init_kwargs = ray_init_kwargs or {}
+        self.reader_kwargs = reader_kwargs or {}
+        self._gcp_temp_file = None
+
+    def _get_table_info(self) -> dict:
+        url = f"{self.base_url}/api/2.1/unity-catalog/tables/{self.table_full_name}"
+        headers = {"Authorization": f"Bearer {self.token}"}
+        resp = requests.get(url, headers=headers)
+        resp.raise_for_status()
+        data = resp.json()
+        self._table_info = data
+        self._table_id = data["table_id"]
+        return data
+
+    def _get_creds(self):
+        url = f"{self.base_url}/api/2.1/unity-catalog/temporary-table-credentials"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.token}",
+        }
+        payload = {"table_id": self._table_id, "operation": self.operation}
+        resp = requests.post(url, json=payload, headers=headers)
+        resp.raise_for_status()
+        self._creds_response = resp.json()
+        self._table_url = self._creds_response["url"]
+
+    def _set_env(self):
+        env_vars = {}
+        creds = self._creds_response
+
+        if "aws_temp_credentials" in creds:
+            aws = creds["aws_temp_credentials"]
+            env_vars["AWS_ACCESS_KEY_ID"] = aws["access_key_id"]
+            env_vars["AWS_SECRET_ACCESS_KEY"] = aws["secret_access_key"]
+            env_vars["AWS_SESSION_TOKEN"] = aws["session_token"]
+            if self.region:
+                env_vars["AWS_REGION"] = self.region
+                env_vars["AWS_DEFAULT_REGION"] = self.region
+        elif "azuresasuri" in creds:
+            env_vars["AZURE_STORAGE_SAS_TOKEN"] = creds["azuresasuri"]
+        elif "gcp_service_account" in creds:
+            gcp_json = creds["gcp_service_account"]
+            temp_file = tempfile.NamedTemporaryFile(
+                mode="w",
+                prefix="gcp_sa_",
+                suffix=".json",
+                delete=False,
+            )
+            temp_file.write(gcp_json)
+            temp_file.close()
+            env_vars["GOOGLE_APPLICATION_CREDENTIALS"] = temp_file.name
+            self._gcp_temp_file = temp_file.name
+            atexit.register(self._cleanup_gcp_temp_file, temp_file.name)
+        else:
+            raise ValueError(
+                "No known credential type found in Databricks UC response."
+            )
+
+        for k, v in env_vars.items():
+            os.environ[k] = v
+        self._runtime_env = {"env_vars": env_vars}
+
+    @staticmethod
+    def _cleanup_gcp_temp_file(temp_file_path: str):
+        """Clean up temporary GCP service account file."""
+        if temp_file_path and os.path.exists(temp_file_path):
+            try:
+                os.unlink(temp_file_path)
+            except OSError:
+                pass
+
+    def _infer_data_format(self) -> str:
+        if self.data_format:
+            return self.data_format
+
+        info = self._table_info or self._get_table_info()
+        if "data_source_format" in info and info["data_source_format"]:
+            fmt = info["data_source_format"].lower()
+            return fmt
+
+        storage_loc = info.get("storage_location") or getattr(self, "_table_url", None)
+        if storage_loc:
+            ext = os.path.splitext(storage_loc)[-1].replace(".", "").lower()
+            if ext in _FILE_FORMAT_TO_RAY_READER:
+                return ext
+
+        raise ValueError("Could not infer data format from table metadata.")
+
+    def _get_ray_reader(self, data_format: str) -> Callable[..., Any]:
+        fmt = data_format.lower()
+        if fmt in _FILE_FORMAT_TO_RAY_READER:
+            reader_func = getattr(ray.data, _FILE_FORMAT_TO_RAY_READER[fmt], None)
+            if reader_func:
+                return reader_func
+        raise ValueError(f"Unsupported data format: {fmt}")
+
+    def _read_delta_with_credentials(self):
+        """Read Delta table with proper PyArrow filesystem for session tokens."""
+        import pyarrow.fs as pafs
+
+        creds = self._creds_response
+        reader_kwargs = self.reader_kwargs.copy()
+
+        # For AWS, create PyArrow S3FileSystem with session tokens
+        if "aws_temp_credentials" in creds:
+            if not self.region:
+                raise ValueError(
+                    "The 'region' parameter is required for AWS S3 access. "
+                    "Please specify the AWS region (e.g., region='us-west-2')."
+                )
+            aws = creds["aws_temp_credentials"]
+            filesystem = pafs.S3FileSystem(
+                access_key=aws["access_key_id"],
+                secret_key=aws["secret_access_key"],
+                session_token=aws["session_token"],
+                region=self.region,
+            )
+            reader_kwargs["filesystem"] = filesystem
+
+        # Call ray.data.read_delta with proper error handling
+        try:
+            return ray.data.read_delta(self._table_url, **reader_kwargs)
+        except Exception as e:
+            error_msg = str(e)
+            if (
+                "DeletionVectors" in error_msg
+                or "Unsupported reader features" in error_msg
+            ):
+                raise RuntimeError(
+                    f"Delta table uses Deletion Vectors, which requires deltalake>=0.10.0. "
+                    f"Error: {error_msg}\n"
+                    f"Solution: pip install --upgrade 'deltalake>=0.10.0'"
+                ) from e
+            raise
+
+    def read(self):
+        self._get_table_info()
+        self._get_creds()
+        self._set_env()
+
+        data_format = self._infer_data_format()
+
+        if not ray.is_initialized():
+            ray.init(runtime_env=self._runtime_env, **self.ray_init_kwargs)
+
+        # Use special Delta reader for proper filesystem handling
+        if data_format == "delta":
+            return self._read_delta_with_credentials()
+
+        # Use standard reader for other formats
+        reader = self._get_ray_reader(data_format)
+        return reader(self._table_url, **self.reader_kwargs)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/video_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/video_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..07bc4e863f26c1140eda353b4296eefd60620b3d
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/video_datasource.py
@@ -0,0 +1,70 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+logger = logging.getLogger(__name__)
+
+
+class VideoDatasource(FileBasedDatasource):
+    _FILE_EXTENSIONS = [
+        "mp4",
+        "mkv",
+        "mov",
+        "avi",
+        "wmv",
+        "flv",
+        "webm",
+        "m4v",
+        "3gp",
+        "mpeg",
+        "mpg",
+        "ts",
+        "ogv",
+        "rm",
+        "rmvb",
+        "vob",
+        "asf",
+        "f4v",
+        "m2ts",
+        "mts",
+        "divx",
+        "xvid",
+        "mxf",
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        include_timestamps=False,
+        decord_load_args: Optional[Dict[str, Any]] = None,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="decord", package="decord")
+
+        self.include_timestamps = include_timestamps
+        if decord_load_args is None:
+            self.decord_load_args = {}
+        else:
+            self.decord_load_args = decord_load_args
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str):
+        from decord import VideoReader
+
+        reader = VideoReader(f, **self.decord_load_args)
+
+        for frame_index, frame in enumerate(reader):
+            item = {"frame": frame.asnumpy(), "frame_index": frame_index}
+            if self.include_timestamps is True:
+                item["frame_timestamp"] = reader.get_frame_timestamp(frame_index)
+
+            builder = DelegatingBlockBuilder()
+            builder.add(item)
+            yield builder.build()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/webdataset_datasink.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/webdataset_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..57e44600c63cf931ba11a32f65a80af2615155af
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/webdataset_datasink.py
@@ -0,0 +1,53 @@
+import io
+import tarfile
+import time
+import uuid
+from typing import Optional, Union
+
+import pyarrow
+
+from ray.data._internal.datasource.webdataset_datasource import (
+    _apply_list,
+    _default_encoder,
+    _make_iterable,
+)
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_datasink import BlockBasedFileDatasink
+
+
+class WebDatasetDatasink(BlockBasedFileDatasink):
+    def __init__(
+        self,
+        path: str,
+        encoder: Optional[Union[bool, str, callable, list]] = True,
+        *,
+        file_format: str = "tar",
+        **file_datasink_kwargs,
+    ):
+        super().__init__(path, file_format="tar", **file_datasink_kwargs)
+
+        self.encoder = encoder
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        stream = tarfile.open(fileobj=file, mode="w|")
+        samples = _make_iterable(block)
+        for sample in samples:
+            if not isinstance(sample, dict):
+                sample = sample.as_pydict()
+            if self.encoder is not None:
+                sample = _apply_list(self.encoder, sample, default=_default_encoder)
+            if "__key__" not in sample:
+                sample["__key__"] = uuid.uuid4().hex
+            key = sample["__key__"]
+            for k, v in sample.items():
+                if v is None or k.startswith("__"):
+                    continue
+                assert isinstance(v, bytes) or isinstance(v, str)
+                if not isinstance(v, bytes):
+                    v = v.encode("utf-8")
+                ti = tarfile.TarInfo(f"{key}.{k}")
+                ti.size = len(v)
+                ti.mtime = time.time()
+                ti.mode, ti.uname, ti.gname = 0o644, "data", "data"
+                stream.addfile(ti, io.BytesIO(v))
+        stream.close()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/datasource/webdataset_datasource.py b/lib/python3.12/site-packages/ray/data/_internal/datasource/webdataset_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..0289b5ac4838b122acad9ff6dde3c9efc9cd9c12
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/datasource/webdataset_datasource.py
@@ -0,0 +1,385 @@
+# Copyright NVIDIA Corporation 2023
+# SPDX-License-Identifier: Apache-2.0
+
+import fnmatch
+import io
+import json
+import re
+import tarfile
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+from ray.data._internal.util import iterate_with_retry
+from ray.data.block import BlockAccessor
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+def _base_plus_ext(path: str):
+    """Split off all file extensions.
+
+    Returns base, allext.
+
+    Args:
+        path: path with extensions
+
+    Returns:
+        str: path with all extensions removed
+    """
+    match = re.match(r"^((?:.*/|)[^.]+)[.]([^/]*)$", path)
+    if not match:
+        return None, None
+    return match.group(1), match.group(2)
+
+
+def _valid_sample(sample: Dict[str, Any]):
+    """Check whether a sample is valid.
+
+    Args:
+        sample: sample to be checked
+    """
+    return (
+        sample is not None
+        and isinstance(sample, dict)
+        and len(list(sample.keys())) > 0
+        and not sample.get("__bad__", False)
+    )
+
+
+def _apply_list(
+    f: Union[Callable, List[Callable]], sample: Dict[str, Any], default: Callable = None
+):
+    """Apply a list of functions to a sample.
+
+    Args:
+        f: function or list of functions
+        sample: sample to be modified
+        default: default function to be applied to all keys.
+            Defaults to None.
+
+    Returns:
+        modified sample
+    """
+    if f is None:
+        return sample
+    if not isinstance(f, list):
+        f = [f]
+    for g in f:
+        if default is not None and not callable(g):
+            g = partial(default, format=g)
+        sample = g(sample)
+    return sample
+
+
+def _check_suffix(suffix: str, suffixes: Union[list, callable]):
+    """Check whether a suffix is valid.
+
+    Suffixes can be either None (=accept everything), a callable,
+    or a list of patterns. If the pattern contains */? it is treated
+    as a glob pattern, otherwise it is treated as a literal.
+
+    Args:
+        suffix: suffix to be checked
+        suffixes: list of valid suffixes
+    """
+    if suffixes is None:
+        return True
+    if callable(suffixes):
+        return suffixes(suffix)
+    for pattern in suffixes:
+        if "*" in pattern or "?" in pattern:
+            if fnmatch.fnmatch("." + suffix, pattern):
+                return True
+        elif suffix == pattern or "." + suffix == pattern:
+            return True
+    return False
+
+
+def _tar_file_iterator(
+    fileobj: Any,
+    fileselect: Optional[Union[bool, callable, list]] = None,
+    filerename: Optional[Union[bool, callable, list]] = None,
+    verbose_open: bool = False,
+    meta: dict = None,
+):
+    """Iterate over tar file, yielding filename, content pairs for the given tar stream.
+
+    Args:
+        fileobj: file object
+        fileselect: patterns or function selecting
+            files to be selected
+        meta: metadata to be added to each sample
+    """
+    meta = meta or {}
+    stream = tarfile.open(fileobj=fileobj, mode="r|*")
+    if verbose_open:
+        print(f"start {meta}")
+    for tarinfo in stream:
+        fname = tarinfo.name
+        if not tarinfo.isreg() or fname is None:
+            continue
+        data = stream.extractfile(tarinfo).read()
+        fname = _apply_list(filerename, fname)
+        assert isinstance(fname, str)
+        if not _check_suffix(fname, fileselect):
+            continue
+        result = dict(fname=fname, data=data)
+        yield result
+    if verbose_open:
+        print(f"done {meta}")
+
+
+def _group_by_keys(
+    data: List[Dict[str, Any]],
+    keys: callable = _base_plus_ext,
+    suffixes: Optional[Union[list, callable]] = None,
+    meta: dict = None,
+):
+    """Return function over iterator that groups key, value pairs into samples.
+
+    Args:
+        data: iterator over key, value pairs
+        keys: function that returns key, suffix for a given key
+        suffixes: list of suffixes to be included in the sample
+        meta: metadata to be added to each sample
+    """
+    meta = meta or {}
+    current_sample = None
+    for filesample in data:
+        assert isinstance(filesample, dict)
+        fname, value = filesample["fname"], filesample["data"]
+        prefix, suffix = keys(fname)
+        if prefix is None:
+            continue
+        if current_sample is None or prefix != current_sample["__key__"]:
+            if _valid_sample(current_sample):
+                current_sample.update(meta)
+                yield current_sample
+            current_sample = dict(__key__=prefix)
+            if "__url__" in filesample:
+                current_sample["__url__"] = filesample["__url__"]
+        if suffix in current_sample:
+            raise ValueError(
+                f"{fname}: duplicate file name in tar file "
+                + f"{suffix} {current_sample.keys()}, tar is {meta['__url__']}"
+            )
+        if suffixes is None or _check_suffix(suffix, suffixes):
+            current_sample[suffix] = value
+    if _valid_sample(current_sample):
+        current_sample.update(meta)
+        yield current_sample
+
+
+def _default_decoder(sample: Dict[str, Any], format: Optional[Union[bool, str]] = True):
+    """A default decoder for webdataset.
+
+    This handles common file extensions: .txt, .cls, .cls2,
+        .jpg, .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl.
+    These are the most common extensions used in webdataset.
+    For other extensions, users can provide their own decoder.
+
+    Args:
+        sample: sample, modified in place
+    """
+    sample = dict(sample)
+    for key, value in sample.items():
+        extension = key.split(".")[-1]
+        if key.startswith("__"):
+            continue
+        elif extension in ["txt", "text"]:
+            sample[key] = value.decode("utf-8")
+        elif extension in ["cls", "cls2"]:
+            sample[key] = int(value.decode("utf-8"))
+        elif extension in ["jpg", "png", "ppm", "pgm", "pbm", "pnm"]:
+            import numpy as np
+            import PIL.Image
+
+            if format == "PIL":
+                sample[key] = PIL.Image.open(io.BytesIO(value))
+            else:
+                sample[key] = np.asarray(PIL.Image.open(io.BytesIO(value)))
+        elif extension == "json":
+            sample[key] = json.loads(value)
+        elif extension == "npy":
+            import numpy as np
+
+            sample[key] = np.load(io.BytesIO(value))
+        elif extension == "mp":
+            import msgpack
+
+            sample[key] = msgpack.unpackb(value, raw=False)
+        elif extension in ["pt", "pth"]:
+            import torch
+
+            sample[key] = torch.load(io.BytesIO(value))
+        elif extension in ["pickle", "pkl"]:
+            import pickle
+
+            sample[key] = pickle.loads(value)
+    return sample
+
+
+extension_to_format = {"jpg": "jpeg"}
+
+
+def _default_encoder(sample: Dict[str, Any], format: Optional[Union[str, bool]] = True):
+    """A default encoder for webdataset.
+
+    This handles common file extensions: .txt, .cls, .cls2, .jpg,
+        .png, .json, .npy, .mp, .pt, .pth, .pickle, .pkl
+    These are the most common extensions used in webdataset.
+    For other extensions, users can provide their own encoder.
+
+    Args:
+        sample (Dict[str, Any]): sample
+    """
+    sample = dict(sample)
+    for key, value in sample.items():
+        extension = key.split(".")[-1]
+        if key.startswith("__"):
+            continue
+        elif extension in ["txt"]:
+            sample[key] = value.encode("utf-8")
+        elif extension in ["cls", "cls2"]:
+            sample[key] = str(value).encode("utf-8")
+        elif extension in ["jpg", "jpeg", "png", "ppm", "pgm", "pbm", "pnm"]:
+            import numpy as np
+            import PIL.Image
+
+            if isinstance(value, np.ndarray):
+                value = PIL.Image.fromarray(value)
+            assert isinstance(value, PIL.Image.Image)
+            stream = io.BytesIO()
+            value.save(
+                stream, format=extension_to_format.get(extension.lower(), extension)
+            )
+            sample[key] = stream.getvalue()
+        elif extension == "json":
+            sample[key] = json.dumps(value).encode("utf-8")
+        elif extension == "npy":
+            import numpy as np
+
+            stream = io.BytesIO()
+            np.save(stream, value)
+            sample[key] = stream.getvalue()
+        elif extension == "mp":
+            import msgpack
+
+            sample[key] = msgpack.dumps(value)
+        elif extension in ["pt", "pth"]:
+            import torch
+
+            stream = io.BytesIO()
+            torch.save(value, stream)
+            sample[key] = stream.getvalue()
+        elif extension in ["pickle", "pkl"]:
+            import pickle
+
+            stream = io.BytesIO()
+            pickle.dump(value, stream)
+            sample[key] = stream.getvalue()
+    return sample
+
+
+def _make_iterable(block: BlockAccessor):
+    """Make a block iterable.
+
+    This is a placeholder for dealing with more complex blocks.
+
+    Args:
+        block: Ray Dataset block
+
+    Returns:
+        Iterable[Dict[str,Any]]: Iterable of samples
+    """
+    return block.iter_rows(public_row_format=False)
+
+
+class WebDatasetDatasource(FileBasedDatasource):
+    """A Datasource for WebDataset datasets (tar format with naming conventions)."""
+
+    _FILE_EXTENSIONS = ["tar"]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        decoder: Optional[Union[bool, str, callable, list]] = True,
+        fileselect: Optional[Union[bool, callable, list]] = None,
+        filerename: Optional[Union[bool, callable, list]] = None,
+        suffixes: Optional[Union[bool, callable, list]] = None,
+        verbose_open: bool = False,
+        expand_json: bool = False,
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        self.decoder = decoder
+        self.fileselect = fileselect
+        self.filerename = filerename
+        self.suffixes = suffixes
+        self.verbose_open = verbose_open
+        self.expand_json = expand_json
+
+    def _read_stream(self, stream: "pyarrow.NativeFile", path: str):
+        """Read and decode samples from a stream.
+
+        Note that fileselect selects files during reading, while suffixes
+        selects files during the grouping step.
+
+        Args:
+            stream: File descriptor to read from.
+            path: Path to the data.
+            decoder: decoder or list of decoders to be applied to samples
+            fileselect: Predicate for skipping files in tar decoder.
+                Defaults to lambda_:False.
+            suffixes: List of suffixes to be extracted. Defaults to None.
+            verbose_open: Print message when opening files. Defaults to False.
+
+        Yields:
+            List[Dict[str, Any]]: List of sample (list of length 1).
+        """
+
+        import pandas as pd
+
+        def get_tar_file_iterator():
+            return _tar_file_iterator(
+                stream,
+                fileselect=self.fileselect,
+                filerename=self.filerename,
+                verbose_open=self.verbose_open,
+            )
+
+        # S3 can raise transient errors during iteration
+        files = iterate_with_retry(
+            get_tar_file_iterator,
+            "iterate tar file",
+            match=self._data_context.retried_io_errors,
+        )
+
+        samples = _group_by_keys(files, meta=dict(__url__=path), suffixes=self.suffixes)
+        for sample in samples:
+            if self.decoder is not None:
+                sample = _apply_list(self.decoder, sample, default=_default_decoder)
+            if self.expand_json:
+                if isinstance(sample["json"], bytes):
+                    parsed_json = json.loads(sample["json"].decode("utf-8"))
+                elif isinstance(sample["json"], str):
+                    parsed_json = json.loads(sample["json"])
+                elif isinstance(sample["json"], dict):
+                    parsed_json = sample["json"]
+                else:
+                    raise TypeError(
+                        f"Unsupported data type" f" {type(sample['json'])} for sample"
+                    )
+                for k, v in parsed_json.items():
+                    if k not in sample:
+                        sample[k] = []
+                    sample[k].append(v)
+            yield pd.DataFrame(
+                {
+                    k: v if isinstance(v, list) and len(v) == 1 else [v]
+                    for k, v in sample.items()
+                }
+            )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..49ea853c77cdaeec1ffcdb150d7f0427710e2fab
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/__init__.py
@@ -0,0 +1,6 @@
+from .ranker import DefaultRanker, Ranker
+
+
+def create_ranker() -> Ranker:
+    """Create a ranker instance based on environment and configuration."""
+    return DefaultRanker()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..682587e7b26746763e075e3be1c6243a1587fd74
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/autoscaling_requester.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/autoscaling_requester.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..627759b1963705589773ad78e719af61970b5361
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/autoscaling_requester.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/dataset_state.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/dataset_state.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5a2efc4b792c255850048383a45d5308b1fc939f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/dataset_state.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/execution_callback.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/execution_callback.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7860c0911f642d315e9f3a526395d4a7358b17fb
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/execution_callback.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/legacy_compat.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/legacy_compat.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..55a0466421ea21e15a2b29f4ab4c21aa4056feaf
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/legacy_compat.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/progress_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/progress_manager.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b2ecce1bd927317a2efa8635b4237110bf9aba5f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/progress_manager.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/ranker.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/ranker.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1316fb4011f8f9072cebaf587be51e4732b41e7d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/ranker.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/resource_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/resource_manager.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2f7d11799b71969f19137c25f42be913b78fee85
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/resource_manager.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/streaming_executor.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/streaming_executor.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6e9d85feaf316e9605520d499c7dc5a23aa706e5
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/streaming_executor.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/streaming_executor_state.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/streaming_executor_state.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..af33428502dd9881ee0a5814dbe8b181ff02ff5e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/streaming_executor_state.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/util.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1f2f93fd4dc52536ca3e46ff13233e7c41e2cbf5
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/__pycache__/util.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/autoscaling_requester.py b/lib/python3.12/site-packages/ray/data/_internal/execution/autoscaling_requester.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ef1be6b598bbf676426e93acb32005fbb57f493
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/autoscaling_requester.py
@@ -0,0 +1,131 @@
+import math
+import threading
+import time
+from typing import Dict, List
+
+import ray
+from ray.data.context import DataContext
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+# Resource requests are considered stale after this number of seconds, and
+# will be purged.
+RESOURCE_REQUEST_TIMEOUT = 60
+PURGE_INTERVAL = RESOURCE_REQUEST_TIMEOUT * 2
+
+# When the autoscaling is driven by memory pressure and there are abundant
+# CPUs to support incremental CPUs needed to launch more tasks, we'll translate
+# memory pressure into an artificial request of CPUs. The amount of CPUs we'll
+# request is ARTIFICIAL_CPU_SCALING_FACTOR * ray.cluster_resources()["CPU"].
+ARTIFICIAL_CPU_SCALING_FACTOR = 1.2
+
+
+@ray.remote(num_cpus=0, max_restarts=-1, max_task_retries=-1)
+class AutoscalingRequester:
+    """Actor to make resource requests to autoscaler for the datasets.
+
+    The resource requests are set to timeout after RESOURCE_REQUEST_TIMEOUT seconds.
+    For those live requests, we keep track of the last request made for each execution,
+    which overrides all previous requests it made; then sum the requested amounts
+    across all executions as the final request to the autoscaler.
+    """
+
+    def __init__(self):
+        # execution_id -> (List[Dict], expiration timestamp)
+        self._resource_requests = {}
+        # TTL for requests.
+        self._timeout = RESOURCE_REQUEST_TIMEOUT
+
+        self._self_handle = ray.get_runtime_context().current_actor
+
+        # Start a thread to purge expired requests periodically.
+        def purge_thread_run():
+            while True:
+                time.sleep(PURGE_INTERVAL)
+                # Call purge_expired_requests() as an actor task,
+                # so we don't need to handle multi-threading.
+                ray.get(self._self_handle.purge_expired_requests.remote())
+
+        self._purge_thread = threading.Thread(target=purge_thread_run, daemon=True)
+        self._purge_thread.start()
+
+    def purge_expired_requests(self):
+        self._purge()
+        ray.autoscaler.sdk.request_resources(bundles=self._aggregate_requests())
+
+    def request_resources(self, req: List[Dict], execution_id: str):
+        # Purge expired requests before making request to autoscaler.
+        self._purge()
+        # For the same execution_id, we track the latest resource request and
+        # the its expiration timestamp.
+        self._resource_requests[execution_id] = (
+            req,
+            time.time() + self._timeout,
+        )
+        # We aggregate the resource requests across all execution_id's to Ray
+        # autoscaler.
+        ray.autoscaler.sdk.request_resources(bundles=self._aggregate_requests())
+
+    def _purge(self):
+        # Purge requests that are stale.
+        now = time.time()
+        for k, (_, t) in list(self._resource_requests.items()):
+            if t < now:
+                self._resource_requests.pop(k)
+
+    def _aggregate_requests(self) -> List[Dict]:
+        req = []
+        for _, (r, _) in self._resource_requests.items():
+            req.extend(r)
+
+        def get_cpus(req):
+            num_cpus = 0
+            for r in req:
+                if "CPU" in r:
+                    num_cpus += r["CPU"]
+            return num_cpus
+
+        # Round up CPUs to exceed total cluster CPUs so it can actually upscale.
+        # This is to handle the issue where the autoscaling is driven by memory
+        # pressure (rather than CPUs) from streaming executor. In such case, simply
+        # asking for incremental CPUs (e.g. 1 CPU for each ready operator) may not
+        # actually be able to trigger autoscaling if existing CPUs in cluster can
+        # already satisfy the incremental CPUs request.
+        num_cpus = get_cpus(req)
+        if num_cpus > 0:
+            total = ray.cluster_resources()
+            if "CPU" in total and num_cpus <= total["CPU"]:
+                delta = (
+                    math.ceil(ARTIFICIAL_CPU_SCALING_FACTOR * total["CPU"]) - num_cpus
+                )
+                req.extend([{"CPU": 1}] * delta)
+
+        return req
+
+    def _test_set_timeout(self, ttl):
+        """Set the timeout. This is for test only"""
+        self._timeout = ttl
+
+
+# Creating/getting an actor from multiple threads is not safe.
+# https://github.com/ray-project/ray/issues/41324
+_autoscaling_requester_lock: threading.RLock = threading.RLock()
+
+
+def get_or_create_autoscaling_requester_actor():
+    ctx = DataContext.get_current()
+    scheduling_strategy = ctx.scheduling_strategy
+    # Pin the autoscaling requester actor to the local node so it fate-shares with the driver.
+    # Note: for Ray Client, the ray.get_runtime_context().get_node_id() should
+    # point to the head node.
+    scheduling_strategy = NodeAffinitySchedulingStrategy(
+        ray.get_runtime_context().get_node_id(),
+        soft=False,
+    )
+    with _autoscaling_requester_lock:
+        return AutoscalingRequester.options(
+            name="AutoscalingRequester",
+            namespace="AutoscalingRequester",
+            get_if_exists=True,
+            lifetime="detached",
+            scheduling_strategy=scheduling_strategy,
+        ).remote()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0aad671df1029ad1a489dc655facc05c94dab61
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__init__.py
@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING, List
+
+from .backpressure_policy import BackpressurePolicy
+from .concurrency_cap_backpressure_policy import ConcurrencyCapBackpressurePolicy
+from .downstream_capacity_backpressure_policy import (
+    DownstreamCapacityBackpressurePolicy,
+)
+from .resource_budget_backpressure_policy import ResourceBudgetBackpressurePolicy
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+# Default enabled backpressure policies and its config key.
+# Use `DataContext.set_config` to config it.
+ENABLED_BACKPRESSURE_POLICIES = [
+    ConcurrencyCapBackpressurePolicy,
+    ResourceBudgetBackpressurePolicy,
+    DownstreamCapacityBackpressurePolicy,
+]
+ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY = "backpressure_policies.enabled"
+
+
+def get_backpressure_policies(
+    data_context: DataContext,
+    topology: "Topology",
+    resource_manager: "ResourceManager",
+) -> List[BackpressurePolicy]:
+    policies = data_context.get_config(
+        ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY, ENABLED_BACKPRESSURE_POLICIES
+    )
+
+    return [policy(data_context, topology, resource_manager) for policy in policies]
+
+
+__all__ = [
+    "BackpressurePolicy",
+    "ConcurrencyCapBackpressurePolicy",
+    "DownstreamCapacityBackpressurePolicy",
+    "ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY",
+    "get_backpressure_policies",
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b1e993321cc8463b358618981732452f5ea5b2d3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/backpressure_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/backpressure_policy.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..61af66a2d7bd0678f79ce77b6db8092fea2ddaef
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/backpressure_policy.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/concurrency_cap_backpressure_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/concurrency_cap_backpressure_policy.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3c07bad257fe27d135b10a38bca4c7a704dfcb8e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/concurrency_cap_backpressure_policy.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/downstream_capacity_backpressure_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/downstream_capacity_backpressure_policy.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d0142f579df2d61a97884999680b6cbcbdaee922
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/downstream_capacity_backpressure_policy.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/resource_budget_backpressure_policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/resource_budget_backpressure_policy.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..79e2b5ff2cb38a419314424b20759b83e396b616
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/__pycache__/resource_budget_backpressure_policy.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/backpressure_policy.py b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/backpressure_policy.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ec31c474c41ad6be99f7a318248001c8a254516
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/backpressure_policy.py
@@ -0,0 +1,62 @@
+from abc import ABC
+from typing import TYPE_CHECKING, Optional
+
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+
+class BackpressurePolicy(ABC):
+    """Interface for back pressure policies."""
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        """Initialize the backpressure policy.
+
+        Args:
+            data_context: The data context.
+            topology: The execution topology.
+            resource_manager: The resource manager.
+        """
+        self._data_context = data_context
+        self._topology = topology
+        self._resource_manager = resource_manager
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Determine if we can add a new input to the operator. If returns False, the
+        operator will be backpressured and will not be able to run new tasks.
+        Used in `streaming_executor_state.py::select_operator_to_run()`.
+
+        Returns: True if we can add a new input to the operator, False otherwise.
+
+        Note, if multiple backpressure policies are enabled, the operator will be
+        backpressured if any of the policies returns False.
+        """
+        return True
+
+    def max_task_output_bytes_to_read(self, op: "PhysicalOperator") -> Optional[int]:
+        """Return the maximum bytes of pending task outputs can be read for
+        the given operator. None means no limit.
+
+        This is used for output backpressure to limit how much data an operator
+        can read from its running tasks.
+
+        Note, if multiple backpressure policies return non-None values for an operator,
+        the minimum of those values will be used as the limit.
+
+        Args:
+            op: The operator to get the limit for.
+
+        Returns:
+            The maximum bytes that can be read, or None if no limit.
+        """
+        return None
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/concurrency_cap_backpressure_policy.py b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/concurrency_cap_backpressure_policy.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7e10da61c731bc863136b96efc584a83624a068
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/concurrency_cap_backpressure_policy.py
@@ -0,0 +1,231 @@
+import logging
+import math
+from collections import defaultdict
+from typing import TYPE_CHECKING, Dict
+
+from .backpressure_policy import BackpressurePolicy
+from ray._private.ray_constants import env_float
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.task_pool_map_operator import (
+    TaskPoolMapOperator,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.operators.map_operator import MapOperator
+
+logger = logging.getLogger(__name__)
+
+
+class ConcurrencyCapBackpressurePolicy(BackpressurePolicy):
+    """A backpressure policy that caps the concurrency of each operator.
+    This policy dynamically limits the number of concurrent tasks per operator
+    based on the output queue growth rate.
+
+      - Maintain asymmetric EWMA of total enqueued output bytes as the
+        typical level: `level`.
+      - Maintain asymmetric EWMA of absolute residual vs the *previous* level as a
+        scale proxy: `dev = EWMA(|q - level_prev|)`.
+      - Define deadband: Deadband is the acceptable range of the output queue size
+        around the typical level where the queue size is expected to stay stable.
+        deadband [lower, upper] = [level - K_DEV*dev, level + K_DEV*dev].
+      - If q > upper -> target cap = running - BACKOFF_FACTOR  (back off)
+        If q < lower -> target cap = running + RAMPUP_FACTOR  (ramp up)
+        Else         -> target cap = running      (hold)
+      - Apply user-configured max concurrency cap, admit iff running < target cap.
+
+    NOTE: Only support setting concurrency cap for `TaskPoolMapOperator` for now.
+    TODO(chengsu): Consolidate with actor scaling logic of `ActorPoolMapOperator`.
+    """
+
+    # Smoothing factor for the asymmetric EWMA (slow fall, faster rise).
+    EWMA_ALPHA = env_float("RAY_DATA_CONCURRENCY_CAP_EWMA_ALPHA", 0.2)
+    # Deadband width in units of the EWMA absolute deviation estimate.
+    K_DEV = env_float("RAY_DATA_CONCURRENCY_CAP_K_DEV", 2.0)
+    # Factor to back off when the queue is too large.
+    BACKOFF_FACTOR = env_float("RAY_DATA_CONCURRENCY_CAP_BACKOFF_FACTOR", 1)
+    # Factor to ramp up when the queue is too small.
+    RAMPUP_FACTOR = env_float("RAY_DATA_CONCURRENCY_CAP_RAMPUP_FACTOR", 1)
+    # Threshold for per-Op object store budget (available) vs total usage (used)
+    # (available / used) ratio to enable dynamic output queue size backpressure.
+    OBJECT_STORE_USAGE_RATIO = env_float(
+        "RAY_DATA_CONCURRENCY_CAP_OBJECT_STORE_USAGE_RATIO", 0.1
+    )
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Configured per-operator caps (+inf if unset).
+        self._concurrency_caps: Dict["PhysicalOperator", float] = {}
+
+        # EWMA state for level
+        self._q_level_nbytes: Dict["PhysicalOperator", float] = defaultdict(float)
+
+        # EWMA state for dev
+        self._q_level_dev: Dict["PhysicalOperator", float] = defaultdict(float)
+
+        # Per-operator cached threshold (bootstrapped from first sample).
+        self._queue_level_thresholds: Dict["PhysicalOperator", int] = defaultdict(int)
+
+        # Last effective cap for change logs.
+        self._last_effective_caps: Dict["PhysicalOperator", int] = {}
+
+        # Initialize caps from operators (infinite if unset)
+        for op, _ in self._topology.items():
+            if (
+                isinstance(op, TaskPoolMapOperator)
+                and op.get_max_concurrency_limit() is not None
+            ):
+                self._concurrency_caps[op] = op.get_max_concurrency_limit()
+            else:
+                self._concurrency_caps[op] = float("inf")
+
+        # Whether to cap the concurrency of an operator based on its and downstream's queue size.
+        self.enable_dynamic_output_queue_size_backpressure = (
+            self._data_context.enable_dynamic_output_queue_size_backpressure
+        )
+
+        dynamic_output_queue_size_backpressure_configs = ""
+        if self.enable_dynamic_output_queue_size_backpressure:
+            dynamic_output_queue_size_backpressure_configs = (
+                f", EWMA_ALPHA={self.EWMA_ALPHA}, K_DEV={self.K_DEV}, "
+                f"BACKOFF_FACTOR={self.BACKOFF_FACTOR}, RAMPUP_FACTOR={self.RAMPUP_FACTOR}, "
+                f"OBJECT_STORE_USAGE_RATIO={self.OBJECT_STORE_USAGE_RATIO}"
+            )
+        logger.debug(
+            f"ConcurrencyCapBackpressurePolicy caps: {self._concurrency_caps}, "
+            f"enabled: {self.enable_dynamic_output_queue_size_backpressure}{dynamic_output_queue_size_backpressure_configs}"
+        )
+
+    def _update_ewma_asymmetric(self, prev_value: float, sample: float) -> float:
+        """
+        Update EWMA with asymmetric behavior: fast rise, slow fall.
+        Args:
+            prev_value: Previous EWMA value
+            sample: New sample value
+
+        Returns:
+            Updated EWMA value
+        """
+        if prev_value <= 0:
+            return sample
+
+        alpha_up = 1.0 - (1.0 - self.EWMA_ALPHA) ** 2  # fast rise
+        alpha = alpha_up if sample > prev_value else self.EWMA_ALPHA  # slow fall
+        return (1 - alpha) * prev_value + alpha * sample
+
+    def _update_level_and_dev(self, op: "PhysicalOperator", q_bytes: int) -> None:
+        """Update EWMA level and dev (residual w.r.t. previous level)."""
+        q = float(q_bytes)
+
+        level_prev = self._q_level_nbytes[op]
+        dev_prev = self._q_level_dev[op]
+
+        # Deviation vs the previous level
+        dev_sample = abs(q - level_prev) if level_prev > 0 else 0.0
+        dev = self._update_ewma_asymmetric(dev_prev, dev_sample)
+
+        # Now update the level itself
+        level = self._update_ewma_asymmetric(level_prev, q)
+
+        self._q_level_nbytes[op] = level
+        self._q_level_dev[op] = dev
+
+        # For visibility, store the integer center of the band
+        self._queue_level_thresholds[op] = max(1, int(level))
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Return whether `op` may accept another input now."""
+        num_tasks_running = op.metrics.num_tasks_running
+
+        # If not a MapOperator or feature disabled, just enforce configured cap.
+        if (
+            not isinstance(op, MapOperator)
+            or not self.enable_dynamic_output_queue_size_backpressure
+        ):
+            return num_tasks_running < self._concurrency_caps[op]
+
+        # For this Op, if the objectstore budget (available) to total usage (used)
+        # ratio is below threshold (10%), skip dynamic output queue size backpressure.
+        op_usage = self._resource_manager.get_op_usage(op)
+        op_budget = self._resource_manager.get_budget(op)
+        if (
+            op_usage is not None
+            and op_budget is not None
+            and op_budget.object_store_memory > 0
+            and op_usage.object_store_memory > 0
+        ):
+            if (
+                op_budget.object_store_memory / op_usage.object_store_memory
+                > self.OBJECT_STORE_USAGE_RATIO
+            ):
+                # If the objectstore budget (available) to total usage (used)
+                # ratio is above threshold (10%), skip dynamic output queue size
+                # backpressure, but still enforce the configured cap.
+                return num_tasks_running < self._concurrency_caps[op]
+
+        # Current total queued bytes (this op + downstream)
+        current_queue_size_bytes = (
+            self._resource_manager.get_op_internal_object_store_usage(op)
+            + self._resource_manager.get_op_outputs_object_store_usage_with_downstream(
+                op
+            )
+        )
+
+        # Update EWMA state (level & dev) and compute effective cap. Note that
+        # we don't update the EWMA state if the objectstore budget (available) vs total usage (used)
+        # ratio is above threshold (10%), because the level and dev adjusts quickly.
+        self._update_level_and_dev(op, current_queue_size_bytes)
+        effective_cap = self._effective_cap(
+            op, num_tasks_running, current_queue_size_bytes
+        )
+
+        last = self._last_effective_caps.get(op, None)
+        if last != effective_cap:
+            logger.debug(
+                f"Cap change {op.name}: {last if last is not None else 'None'} -> "
+                f"{effective_cap} (running={num_tasks_running}, queue={current_queue_size_bytes}, "
+                f"thr={self._queue_level_thresholds[op]})"
+            )
+            self._last_effective_caps[op] = effective_cap
+
+        return num_tasks_running < effective_cap
+
+    def _effective_cap(
+        self,
+        op: "PhysicalOperator",
+        num_tasks_running: int,
+        current_queue_size_bytes: int,
+    ) -> int:
+        """A simple controller around EWMA level.
+        Args:
+            op: The operator to compute the effective cap for.
+            num_tasks_running: The number of tasks currently running.
+            current_queue_size_bytes: Current total queued bytes for this operator + downstream.
+        Returns:
+            The effective cap.
+        """
+        cap_cfg = self._concurrency_caps[op]
+
+        level = float(self._q_level_nbytes[op])
+        dev = max(1.0, float(self._q_level_dev[op]))
+        upper = level + self.K_DEV * dev
+        lower = level - self.K_DEV * dev
+
+        if current_queue_size_bytes > upper:
+            # back off
+            target = num_tasks_running - self.BACKOFF_FACTOR
+        elif current_queue_size_bytes < lower:
+            # ramp up
+            target = num_tasks_running + self.RAMPUP_FACTOR
+        else:
+            # hold
+            target = num_tasks_running
+
+        # Clamp to [1, configured_cap]
+        target = max(1, target)
+        if not math.isinf(cap_cfg):
+            target = min(target, int(cap_cfg))
+        return int(target)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/downstream_capacity_backpressure_policy.py b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/downstream_capacity_backpressure_policy.py
new file mode 100644
index 0000000000000000000000000000000000000000..26c54a23e53fb753305226070cb4fc52fab9bcd2
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/downstream_capacity_backpressure_policy.py
@@ -0,0 +1,92 @@
+import logging
+from typing import TYPE_CHECKING
+
+from .backpressure_policy import BackpressurePolicy
+from ray.data._internal.execution.operators.actor_pool_map_operator import (
+    ActorPoolMapOperator,
+)
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+logger = logging.getLogger(__name__)
+
+
+class DownstreamCapacityBackpressurePolicy(BackpressurePolicy):
+    """Backpressure policy based on downstream processing capacity.
+
+    This policy triggers backpressure when the output bundles size exceeds both:
+    1. A ratio threshold multiplied by the number of running tasks in downstream operators
+    2. An absolute threshold for the output bundles size
+
+    The policy monitors actual downstream processing capacity by tracking the number
+    of currently running tasks rather than configured parallelism. This approach
+    ensures effective backpressure even when cluster resources are insufficient or
+    scaling is slow, preventing memory pressure and maintaining pipeline stability.
+
+    Key benefits:
+    - Prevents memory bloat from unprocessed output objects
+    - Adapts to actual cluster conditions and resource availability
+    - Maintains balanced throughput across pipeline operators
+    - Reduces object spilling and unnecessary rebuilds
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        super().__init__(data_context, topology, resource_manager)
+        self._backpressure_concurrency_ratio = (
+            self._data_context.downstream_capacity_backpressure_ratio
+        )
+        self._backpressure_max_queued_blocks = (
+            self._data_context.downstream_capacity_backpressure_max_queued_bundles
+        )
+        self._backpressure_disabled = (
+            self._backpressure_concurrency_ratio is None
+            or self._backpressure_max_queued_blocks is None
+        )
+
+    def _max_concurrent_tasks(self, op: "PhysicalOperator") -> int:
+        if isinstance(op, ActorPoolMapOperator):
+            return sum(
+                [
+                    actor_pool.max_concurrent_tasks()
+                    for actor_pool in op.get_autoscaling_actor_pools()
+                ]
+            )
+        return op.num_active_tasks()
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Determine if we can add input to the operator based on downstream capacity."""
+        if self._backpressure_disabled:
+            return True
+        for output_dependency in op.output_dependencies:
+            total_enqueued_blocks = self._topology[
+                output_dependency
+            ].total_enqueued_input_blocks()
+
+            avg_inputs_per_task = (
+                output_dependency.metrics.num_task_inputs_processed
+                / max(output_dependency.metrics.num_tasks_finished, 1)
+            )
+            outstanding_tasks = total_enqueued_blocks / max(avg_inputs_per_task, 1)
+            max_allowed_outstanding = (
+                self._max_concurrent_tasks(output_dependency)
+                * self._backpressure_concurrency_ratio
+            )
+
+            if (
+                total_enqueued_blocks > self._backpressure_max_queued_blocks
+                and outstanding_tasks > max_allowed_outstanding
+            ):
+                return False
+
+        return True
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/resource_budget_backpressure_policy.py b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/resource_budget_backpressure_policy.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d383e07d7d609c47046df320d37b17302cb07cc
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/backpressure_policy/resource_budget_backpressure_policy.py
@@ -0,0 +1,32 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+from .backpressure_policy import BackpressurePolicy
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+class ResourceBudgetBackpressurePolicy(BackpressurePolicy):
+    """A backpressure policy based on resource budgets in ResourceManager."""
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        if self._resource_manager._op_resource_allocator is not None:
+            return self._resource_manager._op_resource_allocator.can_submit_new_task(op)
+
+        return True
+
+    def max_task_output_bytes_to_read(self, op: "PhysicalOperator") -> Optional[int]:
+        """Determine maximum bytes to read based on the resource budgets.
+
+        Args:
+            op: The operator to get the limit for.
+
+        Returns:
+            The maximum bytes that can be read, or None if no limit.
+        """
+        return self._resource_manager.max_task_output_bytes_to_read(op)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa51465258e2838a6b55b928ad3c0d0361a90d38
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__init__.py
@@ -0,0 +1,9 @@
+from .bundle_queue import BundleQueue
+from .fifo_bundle_queue import FIFOBundleQueue
+
+
+def create_bundle_queue() -> BundleQueue:
+    return FIFOBundleQueue()
+
+
+__all__ = ["BundleQueue", "create_bundle_queue"]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..571914e471ee45045fe991e53072377eec650118
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/bundle_queue.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/bundle_queue.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0d2726811a5c5682f93962b30f9bbfc191555040
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/bundle_queue.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/fifo_bundle_queue.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/fifo_bundle_queue.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..53107ecc96d7f379462ee12ceee873e9f8d5c3ed
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/__pycache__/fifo_bundle_queue.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/bundle_queue.py b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/bundle_queue.py
new file mode 100644
index 0000000000000000000000000000000000000000..af437222a63b1bf1c9603dabf2e55b498ee1a180
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/bundle_queue.py
@@ -0,0 +1,75 @@
+import abc
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+class BundleQueue(abc.ABC):
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Return the number of bundles in the queue."""
+        ...
+
+    @abc.abstractmethod
+    def __contains__(self, bundle: "RefBundle") -> bool:
+        """Return whether the bundle is in the queue."""
+        ...
+
+    @abc.abstractmethod
+    def add(self, bundle: "RefBundle") -> None:
+        """Add a bundle to the queue."""
+        ...
+
+    @abc.abstractmethod
+    def get_next(self) -> "RefBundle":
+        """Remove and return the head of the queue.
+
+        Raises:
+            IndexError: If the queue is empty.
+
+        Returns:
+            A Refbundle if has_next() is True
+        """
+        ...
+
+    @abc.abstractmethod
+    def peek_next(self) -> Optional["RefBundle"]:
+        """Return the head of the queue without removing it.
+
+        If the queue is empty, return `None`.
+        """
+        ...
+
+    @abc.abstractmethod
+    def has_next(self) -> bool:
+        """Check if the queue has a valid bundle."""
+        ...
+
+    @abc.abstractmethod
+    def remove(self, bundle: "RefBundle"):
+        """Remove a bundle from the queue."""
+        ...
+
+    @abc.abstractmethod
+    def clear(self):
+        """Remove all bundles from the queue."""
+        ...
+
+    @abc.abstractmethod
+    def estimate_size_bytes(self) -> int:
+        """Return an estimate of the total size of objects in the queue."""
+        ...
+
+    @abc.abstractmethod
+    def num_blocks(self) -> int:
+        """Return the number of blocks in the queue."""
+        ...
+
+    @abc.abstractmethod
+    def is_empty(self):
+        """Return whether this queue and all of its internal data structures are empty.
+
+        This method is used for testing.
+        """
+        ...
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/fifo_bundle_queue.py b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/fifo_bundle_queue.py
new file mode 100644
index 0000000000000000000000000000000000000000..b17cffd449851673c1173bdc0760030aaa566b9c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/bundle_queue/fifo_bundle_queue.py
@@ -0,0 +1,139 @@
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from .bundle_queue import BundleQueue
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+@dataclass
+class _Node:
+    value: "RefBundle"
+    next: Optional["_Node"] = None
+    prev: Optional["_Node"] = None
+
+
+class FIFOBundleQueue(BundleQueue):
+    """A bundle queue that follows a first-in-first-out policy."""
+
+    def __init__(self):
+        # We manually implement a linked list because we need to remove elements
+        # efficiently, and Python's built-in data structures have O(n) removal time.
+        self._head: Optional[_Node] = None
+        self._tail: Optional[_Node] = None
+        # We use a dictionary to keep track of the nodes corresponding to each bundle.
+        # This allows us to remove a bundle from the queue in O(1) time. We need a list
+        # because a bundle can be added to the queue multiple times. Nodes in each list
+        # are insertion-ordered.
+        self._bundle_to_nodes: Dict["RefBundle", List[_Node]] = defaultdict(deque)
+
+        self._nbytes = 0
+        self._num_blocks = 0
+        self._num_bundles = 0
+
+    def __len__(self) -> int:
+        return self._num_bundles
+
+    def __contains__(self, bundle: "RefBundle") -> bool:
+        return bundle in self._bundle_to_nodes
+
+    def add(self, bundle: "RefBundle") -> None:
+        """Add a bundle to the end (right) of the queue."""
+        new_node = _Node(value=bundle, next=None, prev=self._tail)
+        # Case 1: The queue is empty.
+        if self._head is None:
+            assert self._tail is None
+            self._head = new_node
+            self._tail = new_node
+        # Case 2: The queue has at least one element.
+        else:
+            self._tail.next = new_node
+            self._tail = new_node
+
+        self._bundle_to_nodes[bundle].append(new_node)
+
+        self._nbytes += bundle.size_bytes()
+        self._num_blocks += len(bundle.block_refs)
+        self._num_bundles += 1
+
+    def get_next(self) -> "RefBundle":
+        """Return the first (left) bundle in the queue."""
+        # Case 1: The queue is empty.
+        if not self._head:
+            raise IndexError("You can't pop from an empty queue")
+
+        bundle = self._head.value
+        self.remove(bundle)
+
+        return bundle
+
+    def has_next(self) -> bool:
+        return self._num_bundles > 0
+
+    def peek_next(self) -> Optional["RefBundle"]:
+        """Return the first (left) bundle in the queue without removing it."""
+        if self._head is None:
+            return None
+
+        return self._head.value
+
+    def remove(self, bundle: "RefBundle"):
+        """Remove a bundle from the queue.
+
+        If there are multiple instances of the bundle in the queue, this method only
+        removes the first one.
+        """
+        # Case 1: The queue is empty.
+        if bundle not in self._bundle_to_nodes:
+            raise ValueError(f"The bundle {bundle} is not in the queue.")
+
+        node = self._bundle_to_nodes[bundle].popleft()
+        if not self._bundle_to_nodes[bundle]:
+            del self._bundle_to_nodes[bundle]
+
+        # Case 2: The bundle is the only element in the queue.
+        if self._head is self._tail:
+            self._head = None
+            self._tail = None
+        # Case 3: The bundle is the first element in the queue.
+        elif node is self._head:
+            self._head = node.next
+            self._head.prev = None
+        # Case 4: The bundle is the last element in the queue.
+        elif node is self._tail:
+            self._tail = node.prev
+            self._tail.next = None
+        # Case 5: The bundle is in the middle of the queue.
+        else:
+            node.prev.next = node.next
+            node.next.prev = node.prev
+
+        self._num_bundles -= 1
+        self._num_blocks -= len(bundle)
+        self._nbytes -= bundle.size_bytes()
+
+        assert self._nbytes >= 0, (
+            "Expected the total size of objects in the queue to be non-negative, but "
+            f"got {self._nbytes} bytes instead."
+        )
+
+        return node.value
+
+    def clear(self):
+        self._head = None
+        self._tail = None
+        self._bundle_to_nodes.clear()
+        self._num_bundles = 0
+        self._num_blocks = 0
+        self._nbytes = 0
+
+    def estimate_size_bytes(self) -> int:
+        return self._nbytes
+
+    def num_blocks(self) -> int:
+        return self._num_blocks
+
+    def is_empty(self):
+        return not self._bundle_to_nodes and self._head is None and self._tail is None
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4c04197e91c89a2a75602ce186d6935b8d631f9
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__init__.py
@@ -0,0 +1,7 @@
+from ray.data._internal.execution.callbacks.insert_issue_detectors import (
+    IssueDetectionExecutionCallback,
+)
+
+__all__ = [
+    "IssueDetectionExecutionCallback",
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..156157e9129c3f96603aaf1e65df6fedc66af282
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__pycache__/insert_issue_detectors.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__pycache__/insert_issue_detectors.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..05697ec0378eb4100fd27ca4707cfe5f57b61a90
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/__pycache__/insert_issue_detectors.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/insert_issue_detectors.py b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/insert_issue_detectors.py
new file mode 100644
index 0000000000000000000000000000000000000000..661a873fbcc1371bdc815068e3bc718f2c80e05a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/callbacks/insert_issue_detectors.py
@@ -0,0 +1,23 @@
+from typing import TYPE_CHECKING
+
+from ray.data._internal.execution.execution_callback import (
+    ExecutionCallback,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+from ray.data._internal.issue_detection.issue_detector_manager import (
+    IssueDetectorManager,
+)
+
+
+class IssueDetectionExecutionCallback(ExecutionCallback):
+    """ExecutionCallback that handles issue detection."""
+
+    def before_execution_starts(self, executor: "StreamingExecutor"):
+        # Initialize issue detector in StreamingExecutor
+        executor._issue_detector_manager = IssueDetectorManager(executor)
+
+    def on_execution_step(self, executor: "StreamingExecutor"):
+        # Invoke all issue detectors
+        executor._issue_detector_manager.invoke_detectors()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/dataset_state.py b/lib/python3.12/site-packages/ray/data/_internal/execution/dataset_state.py
new file mode 100644
index 0000000000000000000000000000000000000000..702963234baf5c0bab465c68186d708c45cf2f9d
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/dataset_state.py
@@ -0,0 +1,22 @@
+import enum
+
+
+class DatasetState(enum.IntEnum):
+    """Enum representing the possible states of a dataset during execution."""
+
+    UNKNOWN = 0
+    RUNNING = 1
+    FINISHED = 2
+    FAILED = 3
+    PENDING = 4
+
+    def __str__(self):
+        return self.name
+
+    @classmethod
+    def from_string(cls, text):
+        """Get enum by name."""
+        try:
+            return cls[text]  # This uses the name to lookup the enum
+        except KeyError:
+            return cls.UNKNOWN
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/execution_callback.py b/lib/python3.12/site-packages/ray/data/_internal/execution/execution_callback.py
new file mode 100644
index 0000000000000000000000000000000000000000..bc13746d004ce7a8d26e832dd14bc752d69507d7
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/execution_callback.py
@@ -0,0 +1,83 @@
+import importlib
+import os
+from typing import TYPE_CHECKING, List
+
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+
+EXECUTION_CALLBACKS_CONFIG_KEY = "execution_callbacks"
+EXECUTION_CALLBACKS_ENV_VAR = "RAY_DATA_EXECUTION_CALLBACKS"
+ENV_CALLBACKS_INITIALIZED_KEY = "_env_callbacks_initialized"
+
+
+class ExecutionCallback:
+    """Callback interface for execution events."""
+
+    def before_execution_starts(self, executor: "StreamingExecutor"):
+        """Called before the Dataset execution starts."""
+        ...
+
+    def on_execution_step(self, executor: "StreamingExecutor"):
+        """Called at each step of the Dataset execution loop."""
+        ...
+
+    def after_execution_succeeds(self, executor: "StreamingExecutor"):
+        """Called after the Dataset execution succeeds."""
+        ...
+
+    def after_execution_fails(self, executor: "StreamingExecutor", error: Exception):
+        """Called after the Dataset execution fails."""
+        ...
+
+
+def _initialize_env_callbacks(context: DataContext) -> None:
+    """Initialize callbacks from environment variable and add them to the context."""
+    callbacks_str = os.environ.get(EXECUTION_CALLBACKS_ENV_VAR, "")
+    if not callbacks_str:
+        return
+
+    for callback_path in callbacks_str.split(","):
+        callback_path = callback_path.strip()
+        if not callback_path:
+            continue
+
+        try:
+            module_path, class_name = callback_path.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            callback_cls = getattr(module, class_name)
+            callback = callback_cls()
+            add_execution_callback(callback, context)
+        except (ImportError, AttributeError, ValueError) as e:
+            raise ValueError(f"Failed to import callback from '{callback_path}': {e}")
+
+
+def get_execution_callbacks(context: DataContext) -> List[ExecutionCallback]:
+    """Get all ExecutionCallbacks from the DataContext."""
+    # Initialize environment callbacks if not already done for this context
+    if not context.get_config(ENV_CALLBACKS_INITIALIZED_KEY, False):
+        _initialize_env_callbacks(context)
+        context.set_config(ENV_CALLBACKS_INITIALIZED_KEY, True)
+
+    from ray.data._internal.execution.callbacks.insert_issue_detectors import (
+        IssueDetectionExecutionCallback,
+    )
+
+    return context.get_config(
+        EXECUTION_CALLBACKS_CONFIG_KEY, [IssueDetectionExecutionCallback()]
+    )
+
+
+def add_execution_callback(callback: ExecutionCallback, context: DataContext):
+    """Add an ExecutionCallback to the DataContext."""
+    execution_callbacks = context.get_config(EXECUTION_CALLBACKS_CONFIG_KEY, [])
+    execution_callbacks.append(callback)
+    context.set_config(EXECUTION_CALLBACKS_CONFIG_KEY, execution_callbacks)
+
+
+def remove_execution_callback(callback: ExecutionCallback, context: DataContext):
+    """Remove an ExecutionCallback from the DataContext."""
+    execution_callbacks = context.get_config(EXECUTION_CALLBACKS_CONFIG_KEY, [])
+    execution_callbacks.remove(callback)
+    context.set_config(EXECUTION_CALLBACKS_CONFIG_KEY, execution_callbacks)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..dab70136d559567f7656379c1e641e915ddb87b0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__init__.py
@@ -0,0 +1,21 @@
+from .common import NodeIdStr
+from .execution_options import ExecutionOptions, ExecutionResources
+from .executor import Executor, OutputIterator
+from .physical_operator import PhysicalOperator, ReportsExtraResourceUsage
+from .ref_bundle import BlockSlice, RefBundle
+from .task_context import TaskContext
+from .transform_fn import AllToAllTransformFn
+
+__all__ = [
+    "AllToAllTransformFn",
+    "ExecutionOptions",
+    "ExecutionResources",
+    "Executor",
+    "NodeIdStr",
+    "OutputIterator",
+    "PhysicalOperator",
+    "RefBundle",
+    "BlockSlice",
+    "ReportsExtraResourceUsage",
+    "TaskContext",
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..53072a8494f85e078ee618ba1c0d88d1f4541f8a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/common.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a9308e6b2c5c4667be937bb339046fa165bf5ebb
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/common.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/execution_options.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/execution_options.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..86c8fe91b4c1692ae5c6c0600cab5d3b0c4cb1c8
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/execution_options.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/executor.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/executor.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..68d4dea7299064e6578527d8507efcb736ace7b6
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/executor.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/op_runtime_metrics.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/op_runtime_metrics.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..92023be16fc54c90f72cfa70b3ca7394fbbc67d8
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/op_runtime_metrics.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/physical_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/physical_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c0838b563b3879fa421b1d749daf7f7bdd8bb6cd
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/physical_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/ref_bundle.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/ref_bundle.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1fba27a2709d96778511ec22884821a83258be8d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/ref_bundle.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/task_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/task_context.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b81d4e2cba8c253ec8f42840fd35bf7680c31b81
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/task_context.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/transform_fn.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/transform_fn.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e4e8e682d808b6f739f323935cad11ff62776ea7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/__pycache__/transform_fn.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/common.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..beca90c0c9ffe6c842421f8086d64fa486d0f812
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/common.py
@@ -0,0 +1,186 @@
+import bisect
+import json
+from typing import Dict, List, Tuple
+
+from ray.data._internal.util import GiB, KiB, MiB
+from ray.util.metrics import Histogram
+
+# Node id string returned by `ray.get_runtime_context().get_node_id()`.
+NodeIdStr = str
+
+# Used for time-based histograms (e.g. task completion time, block completion time)
+histogram_buckets_s = [
+    0.1,
+    0.25,
+    0.5,
+    1.0,
+    2.5,
+    5.0,
+    7.5,
+    10.0,
+    15.0,
+    20.0,
+    25.0,
+    50.0,
+    75.0,
+    100.0,
+    150.0,
+    500.0,
+    1000.0,
+    2500.0,
+    5000.0,
+]
+
+# Used for size-based histograms (e.g. block size in bytes)
+histogram_buckets_bytes = [
+    KiB,
+    8 * KiB,
+    64 * KiB,
+    128 * KiB,
+    256 * KiB,
+    512 * KiB,
+    MiB,
+    8 * MiB,
+    64 * MiB,
+    128 * MiB,
+    256 * MiB,
+    512 * MiB,
+    GiB,
+    4 * GiB,
+    16 * GiB,
+    64 * GiB,
+    128 * GiB,
+    256 * GiB,
+    512 * GiB,
+    1024 * GiB,
+    4096 * GiB,
+]
+
+# Used for row count-based histograms (e.g. block size in rows)
+histogram_bucket_rows = [
+    1,
+    5,
+    10,
+    25,
+    50,
+    100,
+    250,
+    500,
+    1_000,
+    2_500,
+    5_000,
+    10_000,
+    25_000,
+    50_000,
+    100_000,
+    250_000,
+    500_000,
+    1_000_000,
+    2_500_000,
+    5_000_000,
+    10_000_000,
+]
+
+
+class RuntimeMetricsHistogram:
+    """
+    Class that tracks a histogram of values.
+
+    Contains helper methods to record the values and apply those values to a `ray.util.metrics.Histogram` metric.
+    """
+
+    def __init__(self, boundaries: List[float]):
+        self._boundaries = boundaries
+        # Initialize bucket counts to 0 (+1 additional bucket to represent the +Inf bucket)
+        self._bucket_counts = [0 for _ in range(len(boundaries) + 1)]
+        self._memoized_avg = None
+
+    def observe(self, value: float, num_observations: int = 1):
+        self._bucket_counts[self._find_bucket_index(value)] += num_observations
+        self._memoized_avg = None
+
+    def export_to(
+        self,
+        metric: Histogram,
+        tags: Dict[str, str],
+    ):
+        """
+        This method calculates the difference between the current bucket counts and the previous bucket counts,
+        and applies those observations to the metric.
+
+        This method stores the previous_bucket_counts in the metric as `last_applied_bucket_counts_for_tags`.
+        """
+        if getattr(metric, "last_applied_bucket_counts_for_tags", None) is None:
+            metric.last_applied_bucket_counts_for_tags = {}
+        tags_key = json.dumps(tags, sort_keys=True)
+        previous_bucket_counts = metric.last_applied_bucket_counts_for_tags.get(
+            tags_key
+        )
+
+        for i in range(len(self._bucket_counts)):
+            # Pick a value between the boundaries so the sample falls into the right bucket.
+            # We need to calculate the mid point because choosing the exact boundary value
+            # seems to have unreliable behavior on which bucket it ends up in.
+            boundary_upper_bound = (
+                self._boundaries[i]
+                if i < len(self._bucket_counts) - 1
+                # Since choosing an exact boundary value is unreliable to if it'll
+                # end up in the upper or lower bucket, we add a small buffer to the
+                # last boundary. The amount of the value doesn't matter much
+                # since it's the last bucket and should go to infinity.
+                else self._boundaries[-1] + 100
+            )
+            boundary_lower_bound = self._boundaries[i - 1] if i > 0 else 0
+            bucket_value = (boundary_upper_bound + boundary_lower_bound) / 2
+
+            # Calculate how many observations to add to the metric
+            diff = (
+                self._bucket_counts[i] - previous_bucket_counts[i]
+                if previous_bucket_counts is not None
+                else self._bucket_counts[i]
+            )
+            for _ in range(diff):
+                metric.observe(bucket_value, tags)
+
+        metric.last_applied_bucket_counts_for_tags[
+            tags_key
+        ] = self._bucket_counts.copy()
+
+    def __repr__(self):
+        if self._memoized_avg is None:
+            self._memoized_avg = self._calculate_average_value()
+        total_samples, average = self._memoized_avg
+        return f"(samples: {total_samples}, avg: {average:.2f})"
+
+    def _calculate_average_value(self) -> Tuple[int, float]:
+        """
+        Calculate the average value of all samples.
+
+        Used to show a representative value for the histogram when
+        printing the histogram as a string.
+        """
+        total_samples = sum(self._bucket_counts)
+        if total_samples == 0:
+            return total_samples, 0
+
+        weighted_sum = 0.0
+        for i, count in enumerate(self._bucket_counts):
+            if count > 0:
+                # Calculate representative value for this bucket
+                if i == 0:
+                    # First bucket: 0 to first boundary
+                    bucket_value = self._boundaries[0] / 2
+                elif i == len(self._bucket_counts) - 1:
+                    # Last bucket: last boundary to +inf
+                    bucket_value = self._boundaries[-1] * 1.5
+                else:
+                    # Middle buckets: between boundaries
+                    bucket_value = (self._boundaries[i - 1] + self._boundaries[i]) / 2
+
+                weighted_sum += bucket_value * count
+
+        average = weighted_sum / total_samples
+        return total_samples, average
+
+    def _find_bucket_index(self, value: float):
+        return bisect.bisect_left(self._boundaries, value)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/execution_options.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/execution_options.py
new file mode 100644
index 0000000000000000000000000000000000000000..3edfa2dceda5ebc31bb6fc7b3c4a73c1297dcae3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/execution_options.py
@@ -0,0 +1,374 @@
+import math
+import os
+from typing import Any, Dict, List, Optional, Union
+
+from .common import NodeIdStr
+from ray.data._internal.execution.util import memory_string
+from ray.util.annotations import DeveloperAPI
+
+
+class ExecutionResources:
+    """Specifies resources usage or resource limits for execution.
+
+    By default this class represents resource usage. Use `for_limits` or
+    set `default_to_inf` to True to create an object that represents resource limits.
+    """
+
+    def __init__(
+        self,
+        cpu: Optional[float] = None,
+        gpu: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+        memory: Optional[float] = None,
+    ):
+        """Initializes ExecutionResources.
+        Args:
+            cpu: Amount of logical CPU slots.
+            gpu: Amount of logical GPU slots.
+            object_store_memory: Amount of object store memory.
+            memory: Amount of logical memory in bytes.
+        """
+
+        # NOTE: Ray Core allocates fractional resources in up to 5th decimal
+        #       digit, hence we round the values here up to it
+        self._cpu: Optional[float] = safe_round(cpu, 5)
+        self._gpu: Optional[float] = safe_round(gpu, 5)
+        self._object_store_memory: Optional[float] = safe_round(object_store_memory)
+        self._memory: Optional[float] = safe_round(memory)
+
+    @classmethod
+    def from_resource_dict(
+        cls,
+        resource_dict: Dict[str, float],
+    ):
+        """Create an ExecutionResources object from a resource dict."""
+        return ExecutionResources(
+            cpu=resource_dict.get("CPU", None) or resource_dict.get("num_cpus", None),
+            gpu=resource_dict.get("GPU", None) or resource_dict.get("num_gpus", None),
+            object_store_memory=resource_dict.get("object_store_memory", None),
+            memory=resource_dict.get("memory", None),
+        )
+
+    def to_resource_dict(self) -> Dict[str, float]:
+        """Convert this ExecutionResources object to a resource dict."""
+        return {
+            "CPU": self.cpu,
+            "GPU": self.gpu,
+            "object_store_memory": self.object_store_memory,
+            "memory": self.memory,
+        }
+
+    @classmethod
+    def for_limits(
+        cls,
+        cpu: Optional[float] = None,
+        gpu: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+        memory: Optional[float] = None,
+    ) -> "ExecutionResources":
+        """Create an ExecutionResources object that represents resource limits.
+        Args:
+            cpu: Amount of logical CPU slots.
+            gpu: Amount of logical GPU slots.
+            object_store_memory: Amount of object store memory.
+            memory: Amount of logical memory in bytes.
+        """
+        return ExecutionResources(
+            cpu=safe_or(cpu, float("inf")),
+            gpu=safe_or(gpu, float("inf")),
+            object_store_memory=safe_or(object_store_memory, float("inf")),
+            memory=safe_or(memory, float("inf")),
+        )
+
+    @property
+    def cpu(self) -> float:
+        return self._cpu or 0.0
+
+    @property
+    def gpu(self) -> float:
+        return self._gpu or 0.0
+
+    @property
+    def object_store_memory(self) -> float:
+        return self._object_store_memory or 0
+
+    @property
+    def memory(self) -> float:
+        return self._memory or 0
+
+    def __repr__(self):
+        return (
+            f"ExecutionResources(cpu={self.cpu}, gpu={self.gpu}, "
+            f"object_store_memory={self.object_store_memory_str()}, "
+            f"memory={self.memory_str()})"
+        )
+
+    def __eq__(self, other: "ExecutionResources") -> bool:
+        return (
+            self.cpu == other.cpu
+            and self.gpu == other.gpu
+            and self.object_store_memory == other.object_store_memory
+            and self.memory == other.memory
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.cpu,
+                self.gpu,
+                self.object_store_memory,
+                self.memory,
+            )
+        )
+
+    @classmethod
+    def zero(cls) -> "ExecutionResources":
+        """Returns an ExecutionResources object with zero resources."""
+        return ExecutionResources(0.0, 0.0, 0.0, 0.0)
+
+    @classmethod
+    def inf(cls) -> "ExecutionResources":
+        """Returns an ExecutionResources object with infinite resources."""
+        return ExecutionResources.for_limits()
+
+    def is_zero(self) -> bool:
+        """Returns True if all resources are zero."""
+        return (
+            self.cpu == 0.0
+            and self.gpu == 0.0
+            and self.object_store_memory == 0.0
+            and self.memory == 0.0
+        )
+
+    def is_non_negative(self) -> bool:
+        """Returns True if all resources are non-negative."""
+        return (
+            self.cpu >= 0
+            and self.gpu >= 0
+            and self.object_store_memory >= 0
+            and self.memory >= 0
+        )
+
+    def object_store_memory_str(self) -> str:
+        """Returns a human-readable string for the object store memory field."""
+        if self.object_store_memory == float("inf"):
+            return "inf"
+        return memory_string(self.object_store_memory)
+
+    def memory_str(self) -> str:
+        """Returns a human-readable string for the memory field."""
+        if self.memory == float("inf"):
+            return "inf"
+        return memory_string(self.memory)
+
+    def copy(
+        self,
+        cpu: Optional[float] = None,
+        gpu: Optional[float] = None,
+        memory: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+    ) -> "ExecutionResources":
+        """Returns a copy of this ExecutionResources object allowing to override
+        specific resources as necessary"""
+        return ExecutionResources(
+            cpu=safe_or(cpu, self.cpu),
+            gpu=safe_or(gpu, self.gpu),
+            object_store_memory=safe_or(object_store_memory, self.object_store_memory),
+            memory=safe_or(memory, self.memory),
+        )
+
+    def add(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Adds execution resources.
+
+        Returns:
+            A new ExecutionResource object with summed resources.
+        """
+        return ExecutionResources(
+            cpu=self.cpu + other.cpu,
+            gpu=self.gpu + other.gpu,
+            object_store_memory=self.object_store_memory + other.object_store_memory,
+            memory=self.memory + other.memory,
+        )
+
+    def subtract(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Subtracts execution resources.
+
+        Returns:
+            A new ExecutionResource object with subtracted resources.
+        """
+        return ExecutionResources(
+            cpu=self.cpu - other.cpu,
+            gpu=self.gpu - other.gpu,
+            object_store_memory=self.object_store_memory - other.object_store_memory,
+            memory=self.memory - other.memory,
+        )
+
+    def max(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the maximum for each resource type."""
+        return ExecutionResources(
+            cpu=max(self.cpu, other.cpu),
+            gpu=max(self.gpu, other.gpu),
+            object_store_memory=max(
+                self.object_store_memory, other.object_store_memory
+            ),
+            memory=max(self.memory, other.memory),
+        )
+
+    def min(self, other: "ExecutionResources") -> "ExecutionResources":
+        """Returns the minimum for each resource type."""
+        return ExecutionResources(
+            cpu=min(self.cpu, other.cpu),
+            gpu=min(self.gpu, other.gpu),
+            object_store_memory=min(
+                self.object_store_memory, other.object_store_memory
+            ),
+            memory=min(self.memory, other.memory),
+        )
+
+    def satisfies_limit(
+        self,
+        limit: "ExecutionResources",
+        *,
+        ignore_object_store_memory=False,
+    ) -> bool:
+        """Return if this resource struct meets the specified limits.
+
+        Note that None for a field means no limit.
+
+        Args:
+            limit: The resource limits to check against.
+            ignore_object_store_memory: If True, ignore the object store memory
+                limit when checking if this resource struct meets the limits.
+        """
+        return (
+            self.cpu <= limit.cpu
+            and self.gpu <= limit.gpu
+            and (
+                ignore_object_store_memory
+                or self.object_store_memory <= limit.object_store_memory
+            )
+            and self.memory <= limit.memory
+        )
+
+    def scale(self, f: float) -> "ExecutionResources":
+        """Return copy with all set values scaled by `f`."""
+        if f < 0:
+            raise ValueError("Scaling factor must be non-negative.")
+        if f == 0:
+            # Explicitly handle the zero case, because `0 * inf` is undefined.
+            return ExecutionResources.zero()
+
+        return ExecutionResources(
+            cpu=self.cpu * f,
+            gpu=self.gpu * f,
+            object_store_memory=self.object_store_memory * f,
+            memory=self.memory * f,
+        )
+
+
+@DeveloperAPI
+class ExecutionOptions:
+    """Common options for execution.
+
+    Some options may not be supported on all executors (e.g., resource limits).
+
+    Attributes:
+        resource_limits: Set a soft limit on the resource usage during execution.
+            Autodetected by default.
+        exclude_resources: Amount of resources to exclude from Ray Data.
+            Set this if you have other workloads running on the same cluster.
+            Note,
+            - If using Ray Data with Ray Train, training resources will be
+            automatically excluded.
+            - For each resource type, resource_limits and exclude_resources can
+            not be both set.
+        locality_with_output: Set this to prefer running tasks on the same node as the
+            output node (node driving the execution). It can also be set to a list of
+            node ids to spread the outputs across those nodes. Off by default.
+        preserve_order: Set this to preserve the ordering between blocks processed by
+            operators. Off by default.
+        actor_locality_enabled: Whether to enable locality-aware task dispatch to
+            actors (off by default). This parameter applies to both stateful map and
+            streaming_split operations.
+        verbose_progress: Whether to report progress individually per operator. By
+            default, only AllToAll operators and global progress is reported. This
+            option is useful for performance debugging. On by default.
+    """
+
+    def __init__(
+        self,
+        resource_limits: Optional[ExecutionResources] = None,
+        exclude_resources: Optional[ExecutionResources] = None,
+        locality_with_output: Union[bool, List[NodeIdStr]] = False,
+        preserve_order: bool = False,
+        actor_locality_enabled: bool = True,
+        verbose_progress: Optional[bool] = None,
+    ):
+        if resource_limits is None:
+            resource_limits = ExecutionResources.for_limits()
+        self.resource_limits = resource_limits
+        if exclude_resources is None:
+            exclude_resources = ExecutionResources.zero()
+        self.exclude_resources = exclude_resources
+        self.locality_with_output = locality_with_output
+        self.preserve_order = preserve_order
+        self.actor_locality_enabled = actor_locality_enabled
+        if verbose_progress is None:
+            verbose_progress = bool(
+                int(os.environ.get("RAY_DATA_VERBOSE_PROGRESS", "1"))
+            )
+        self.verbose_progress = verbose_progress
+
+    def __repr__(self) -> str:
+        return (
+            f"ExecutionOptions(resource_limits={self.resource_limits}, "
+            f"exclude_resources={self.exclude_resources}, "
+            f"locality_with_output={self.locality_with_output}, "
+            f"preserve_order={self.preserve_order}, "
+            f"actor_locality_enabled={self.actor_locality_enabled}, "
+            f"verbose_progress={self.verbose_progress})"
+        )
+
+    @property
+    def resource_limits(self) -> ExecutionResources:
+        return self._resource_limits
+
+    @resource_limits.setter
+    def resource_limits(self, value: ExecutionResources) -> None:
+        self._resource_limits = ExecutionResources.for_limits(
+            cpu=value._cpu,
+            gpu=value._gpu,
+            object_store_memory=value._object_store_memory,
+            memory=value._memory,
+        )
+
+    def is_resource_limits_default(self):
+        """Returns True if resource_limits is the default value."""
+        return self._resource_limits == ExecutionResources.for_limits()
+
+    def validate(self) -> None:
+        """Validate the options."""
+        for attr in ["cpu", "gpu", "object_store_memory"]:
+            if (
+                getattr(self.resource_limits, attr) != float("inf")
+                and getattr(self.exclude_resources, attr, 0) > 0
+            ):
+                raise ValueError(
+                    "resource_limits and exclude_resources cannot "
+                    f" both be set for {attr} resource."
+                )
+
+
+def safe_or(value: Optional[Any], alt: Any) -> Any:
+    return value if value is not None else alt
+
+
+def safe_round(
+    value: Optional[float], ndigits: Optional[int] = None
+) -> Optional[float]:
+    if value is None:
+        return None
+    elif math.isinf(value):
+        return value
+    else:
+        return round(value, ndigits)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/executor.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/executor.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ba2229c428448b031677f725bfdf51149c4edfe
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/executor.py
@@ -0,0 +1,93 @@
+from abc import ABC, abstractmethod
+from typing import ContextManager, Iterator, Optional
+
+from .execution_options import ExecutionOptions
+from .physical_operator import PhysicalOperator
+from .ref_bundle import RefBundle
+from ray.data._internal.stats import DatasetStats
+
+
+class OutputIterator(Iterator[RefBundle], ABC):
+    """Iterator used to access the output of an Executor execution.
+
+    This is a blocking iterator. Datasets guarantees that all its iterators are
+    thread-safe (i.e., multiple threads can block on them at the same time).
+    """
+
+    @abstractmethod
+    def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
+        """Can be used to pull outputs by a specified output index.
+
+        This is used to support the streaming_split() API, where the output of a
+        streaming execution is to be consumed by multiple processes.
+
+        Args:
+            output_split_idx: The output split index to get results for. This arg is
+                only allowed for iterators created by `Dataset.streaming_split()`.
+
+        Raises:
+            StopIteration: If there are no more outputs to return.
+        """
+        ...
+
+    def __next__(self) -> RefBundle:
+        return self.get_next()
+
+
+class Executor(ContextManager, ABC):
+    """Abstract class for executors, which implement physical operator execution.
+
+    Subclasses:
+        StreamingExecutor
+    """
+
+    def __init__(self, options: ExecutionOptions):
+        """Create the executor."""
+        options.validate()
+        self._options = options
+
+    @abstractmethod
+    def execute(
+        self, dag: PhysicalOperator, initial_stats: Optional[DatasetStats] = None
+    ) -> OutputIterator:
+        """Start execution.
+
+        Args:
+            dag: The operator graph to execute.
+            initial_stats: The DatasetStats to prepend to the stats returned by the
+                executor. These stats represent actions done to compute inputs.
+        """
+        ...
+
+    def shutdown(self, force: bool, exception: Optional[Exception] = None):
+        """Shutdown an executor, which may still be running.
+
+        This should interrupt execution and clean up any used resources.
+
+        Args:
+            force: Controls whether shutdown should forcefully terminate all execution
+                   activity (making sure that upon returning from this method all
+                   activities are stopped). When force=False, some activities could be
+                   terminated asynchronously (ie this method won't provide guarantee
+                   that they stop executing before returning from this method)
+            exception: The exception that causes the executor to shut down, or None if
+                the executor finishes successfully.
+        """
+        pass
+
+    @abstractmethod
+    def get_stats(self) -> DatasetStats:
+        """Return stats for the execution so far.
+
+        This is generally called after `execute` has completed, but may be called
+        while iterating over `execute` results for streaming execution.
+        """
+        ...
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback, /):
+        # NOTE: ``ContextManager`` semantic must guarantee that executor
+        #       fully shutdown upon returning from this method
+        self.shutdown(force=True, exception=exc_value)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/op_runtime_metrics.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/op_runtime_metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..cf1f881c826a5c0a49bda9bbc436a54bfc0c9ae3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/op_runtime_metrics.py
@@ -0,0 +1,966 @@
+import math
+import time
+from collections import defaultdict
+from dataclasses import Field, dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+import ray
+from ray.data._internal.execution.bundle_queue import create_bundle_queue
+from ray.data._internal.execution.interfaces.common import (
+    RuntimeMetricsHistogram,
+    histogram_bucket_rows,
+    histogram_buckets_bytes,
+    histogram_buckets_s,
+)
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.memory_tracing import trace_allocation
+from ray.data.block import BlockMetadata
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+
+
+# A metadata key used to mark a dataclass field as a metric.
+_IS_FIELD_METRIC_KEY = "__is_metric"
+# Metadata keys used to store information about a metric.
+_METRIC_FIELD_DESCRIPTION_KEY = "__metric_description"
+_METRIC_FIELD_METRICS_GROUP_KEY = "__metric_metrics_group"
+_METRIC_FIELD_METRICS_TYPE_KEY = "__metric_metrics_type"
+_METRIC_FIELD_METRICS_ARGS_KEY = "__metric_metrics_args"
+_METRIC_FIELD_IS_MAP_ONLY_KEY = "__metric_is_map_only"
+
+_METRICS: List["MetricDefinition"] = []
+
+NODE_UNKNOWN = "unknown"
+
+
+class MetricsGroup(Enum):
+    INPUTS = "inputs"
+    OUTPUTS = "outputs"
+    TASKS = "tasks"
+    OBJECT_STORE_MEMORY = "object_store_memory"
+    MISC = "misc"
+    ACTORS = "actors"
+
+
+class MetricsType(Enum):
+    Counter = 0
+    Gauge = 1
+    Histogram = 2
+
+
+@dataclass(frozen=True)
+class MetricDefinition:
+    """Metadata for a metric.
+
+    Args:
+        name: The name of the metric.
+        description: A human-readable description of the metric, also used as the chart
+            description on the Ray Data dashboard.
+        metrics_group: The group of the metric, used to organize metrics into groups in
+            'StatsActor' and on the Ray Data dashboard.
+        map_only: Whether the metric is only measured for 'MapOperators'.
+    """
+
+    name: str
+    description: str
+    metrics_group: str
+    metrics_type: MetricsType
+    metrics_args: Dict[str, Any]
+    # TODO: Let's refactor this parameter so it isn't tightly coupled with a specific
+    # operator type (MapOperator).
+    map_only: bool = False
+    internal_only: bool = False  # do not expose this metric to the user
+
+
+def metric_field(
+    *,
+    description: str,
+    metrics_group: str,
+    metrics_type: MetricsType = MetricsType.Gauge,
+    metrics_args: Dict[str, Any] = None,
+    map_only: bool = False,
+    internal_only: bool = False,  # do not expose this metric to the user
+    **field_kwargs,
+):
+    """A dataclass field that represents a metric."""
+    metadata = field_kwargs.get("metadata", {})
+
+    metadata[_IS_FIELD_METRIC_KEY] = True
+
+    metadata[_METRIC_FIELD_DESCRIPTION_KEY] = description
+    metadata[_METRIC_FIELD_METRICS_GROUP_KEY] = metrics_group
+    metadata[_METRIC_FIELD_METRICS_TYPE_KEY] = metrics_type
+    metadata[_METRIC_FIELD_METRICS_ARGS_KEY] = metrics_args or {}
+    metadata[_METRIC_FIELD_IS_MAP_ONLY_KEY] = map_only
+
+    return field(metadata=metadata, **field_kwargs)
+
+
+def metric_property(
+    *,
+    description: str,
+    metrics_group: str,
+    metrics_type: MetricsType = MetricsType.Gauge,
+    metrics_args: Dict[str, Any] = None,
+    map_only: bool = False,
+    internal_only: bool = False,  # do not expose this metric to the user
+):
+    """A property that represents a metric."""
+
+    def wrap(func):
+        metric = MetricDefinition(
+            name=func.__name__,
+            description=description,
+            metrics_group=metrics_group,
+            metrics_type=metrics_type,
+            metrics_args=(metrics_args or {}),
+            map_only=map_only,
+            internal_only=internal_only,
+        )
+
+        _METRICS.append(metric)
+
+        return property(func)
+
+    return wrap
+
+
+@dataclass
+class RunningTaskInfo:
+    inputs: RefBundle
+    num_outputs: int
+    bytes_outputs: int
+    num_rows_produced: int
+    start_time: float
+    cum_block_gen_time: float
+
+
+@dataclass
+class NodeMetrics:
+    num_tasks_finished: int = 0
+    bytes_outputs_of_finished_tasks: int = 0
+    blocks_outputs_of_finished_tasks: int = 0
+
+
+class OpRuntimesMetricsMeta(type):
+    def __init__(cls, name, bases, dict):
+        # NOTE: `Field.name` isn't set until the dataclass is created, so we can't
+        # create the metrics in `metric_field` directly.
+        super().__init__(name, bases, dict)
+
+        # Iterate over the attributes and methods of 'OpRuntimeMetrics'.
+        for name, value in dict.items():
+            # If an attribute is a dataclass field and has _IS_FIELD_METRIC_KEY in its
+            # metadata, then create a metric from the field metadata and add it to the
+            # list of metrics. See also the 'metric_field' function.
+            if isinstance(value, Field) and value.metadata.get(_IS_FIELD_METRIC_KEY):
+                metric = MetricDefinition(
+                    name=name,
+                    description=value.metadata[_METRIC_FIELD_DESCRIPTION_KEY],
+                    metrics_group=value.metadata[_METRIC_FIELD_METRICS_GROUP_KEY],
+                    metrics_type=value.metadata[_METRIC_FIELD_METRICS_TYPE_KEY],
+                    metrics_args=value.metadata[_METRIC_FIELD_METRICS_ARGS_KEY],
+                    map_only=value.metadata[_METRIC_FIELD_IS_MAP_ONLY_KEY],
+                )
+                _METRICS.append(metric)
+
+
+def node_id_from_block_metadata(meta: BlockMetadata) -> str:
+    if meta.exec_stats is not None and meta.exec_stats.node_id is not None:
+        node_id = meta.exec_stats.node_id
+    else:
+        node_id = NODE_UNKNOWN
+    return node_id
+
+
+class TaskDurationStats:
+    """
+    Tracks the running mean and variance incrementally with Welford's algorithm
+    by updating the current mean and a measure of total squared differences.
+    It allows stable updates of mean and variance in a single pass over the data
+    while reducing numerical instability often found in naive computations.
+
+    More on the algorithm: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
+    """
+
+    def __init__(self):
+        self._count = 0
+        self._mean = 0.0
+        self._m2 = 0.0  # Sum of (x - mean)^2
+
+    def add_duration(self, duration: float) -> None:
+        """Add a new sample (task duration in seconds)."""
+        self._count += 1
+        delta = duration - self._mean
+        self._mean += delta / self._count
+        delta2 = duration - self._mean
+        self._m2 += delta * delta2
+
+    def count(self) -> int:
+        return self._count
+
+    def mean(self) -> float:
+        return self._mean
+
+    def _variance(self) -> float:
+        """Return the current variance of the observed durations."""
+        # Variance is m2/(count-1) for sample variance
+        if self._count < 2:
+            return 0.0
+        return self._m2 / (self._count - 1)
+
+    def stddev(self) -> float:
+        """Return the current standard deviation of the observed durations."""
+        return math.sqrt(self._variance())
+
+
+@dataclass
+class OpRuntimeMetrics(metaclass=OpRuntimesMetricsMeta):
+    """Runtime metrics for a 'PhysicalOperator'.
+
+    Metrics are updated dynamically during the execution of the Dataset.
+    This class can be used for either observablity or scheduling purposes.
+
+    DO NOT modify the fields of this class directly. Instead, use the provided
+    callback methods.
+    """
+
+    # TODO(hchen): Fields tagged with "map_only" currently only work for MapOperator.
+    # We should make them work for all operators by unifying the task execution code.
+
+    # === Inputs-related metrics ===
+    num_inputs_received: int = metric_field(
+        default=0,
+        description="Number of input blocks received by operator.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    num_row_inputs_received: int = metric_field(
+        default=0,
+        description="Number of input rows received by operator.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    bytes_inputs_received: int = metric_field(
+        default=0,
+        description="Byte size of input blocks received by operator.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    num_task_inputs_processed: int = metric_field(
+        default=0,
+        description=(
+            "Number of input blocks that operator's tasks have finished processing."
+        ),
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    bytes_task_inputs_processed: int = metric_field(
+        default=0,
+        description=(
+            "Byte size of input blocks that operator's tasks have finished processing."
+        ),
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    bytes_inputs_of_submitted_tasks: int = metric_field(
+        default=0,
+        description="Byte size of input blocks passed to submitted tasks.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    rows_inputs_of_submitted_tasks: int = metric_field(
+        default=0,
+        description="Number of rows in the input blocks passed to submitted tasks.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+
+    # === Outputs-related metrics ===
+    num_task_outputs_generated: int = metric_field(
+        default=0,
+        description="Number of output blocks generated by tasks.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    bytes_task_outputs_generated: int = metric_field(
+        default=0,
+        description="Byte size of output blocks generated by tasks.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    rows_task_outputs_generated: int = metric_field(
+        default=0,
+        description="Number of output rows generated by tasks.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    row_outputs_taken: int = metric_field(
+        default=0,
+        description="Number of rows that are already taken by downstream operators.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    block_outputs_taken: int = metric_field(
+        default=0,
+        description="Number of blocks that are already taken by downstream operators.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    num_outputs_taken: int = metric_field(
+        default=0,
+        description=(
+            "Number of output blocks that are already taken by downstream operators."
+        ),
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    bytes_outputs_taken: int = metric_field(
+        default=0,
+        description=(
+            "Byte size of output blocks that are already taken by downstream operators."
+        ),
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    num_outputs_of_finished_tasks: int = metric_field(
+        default=0,
+        description="Number of generated output blocks that are from finished tasks.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    bytes_outputs_of_finished_tasks: int = metric_field(
+        default=0,
+        description=(
+            "Total byte size of generated output blocks produced by finished tasks."
+        ),
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    rows_outputs_of_finished_tasks: int = metric_field(
+        default=0,
+        description=("Number of rows generated by finished tasks."),
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    num_external_inqueue_blocks: int = metric_field(
+        default=0,
+        description="Number of blocks in the external inqueue",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    num_external_inqueue_bytes: int = metric_field(
+        default=0,
+        description="Byte size of blocks in the external inqueue",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    num_external_outqueue_blocks: int = metric_field(
+        default=0,
+        description="Number of blocks in the external outqueue",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    num_external_outqueue_bytes: int = metric_field(
+        default=0,
+        description="Byte size of blocks in the external outqueue",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+
+    # === Tasks-related metrics ===
+    num_tasks_submitted: int = metric_field(
+        default=0,
+        description="Number of submitted tasks.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    num_tasks_running: int = metric_field(
+        default=0,
+        description="Number of running tasks.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    num_tasks_have_outputs: int = metric_field(
+        default=0,
+        description="Number of tasks that already have output.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    num_tasks_finished: int = metric_field(
+        default=0,
+        description="Number of finished tasks.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    num_tasks_failed: int = metric_field(
+        default=0,
+        description="Number of failed tasks.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    block_generation_time: float = metric_field(
+        default=0,
+        description="Time spent generating blocks in tasks.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    task_submission_backpressure_time: float = metric_field(
+        default=0,
+        description="Time spent in task submission backpressure.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    task_output_backpressure_time: float = metric_field(
+        default=0,
+        description="Time spent in task output backpressure.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    task_completion_time_total_s: float = metric_field(
+        default=0,
+        description="Time spent running tasks to completion. This is a sum of all tasks' completion times.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    task_completion_time: RuntimeMetricsHistogram = metric_field(
+        default_factory=lambda: RuntimeMetricsHistogram(histogram_buckets_s),
+        description="Time spent per task running those tasks to completion.",
+        metrics_group=MetricsGroup.TASKS,
+        metrics_type=MetricsType.Histogram,
+        metrics_args={"boundaries": histogram_buckets_s},
+    )
+    block_completion_time: RuntimeMetricsHistogram = metric_field(
+        default_factory=lambda: RuntimeMetricsHistogram(histogram_buckets_s),
+        description="Time spent running a single block to completion. If multiple blocks are generated per task, this is approximated by assuming each block took an equal amount of time to process.",
+        metrics_group=MetricsGroup.TASKS,
+        metrics_type=MetricsType.Histogram,
+        metrics_args={"boundaries": histogram_buckets_s},
+    )
+    task_completion_time_excl_backpressure_s: float = metric_field(
+        default=0,
+        description="Time spent running tasks to completion without backpressure.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    block_size_bytes: RuntimeMetricsHistogram = metric_field(
+        default_factory=lambda: RuntimeMetricsHistogram(histogram_buckets_bytes),
+        description="Size of blocks generated by tasks.",
+        metrics_group=MetricsGroup.TASKS,
+        metrics_type=MetricsType.Histogram,
+        metrics_args={"boundaries": histogram_buckets_bytes},
+    )
+    block_size_rows: RuntimeMetricsHistogram = metric_field(
+        default_factory=lambda: RuntimeMetricsHistogram(histogram_bucket_rows),
+        description="Number of rows in blocks generated by tasks.",
+        metrics_group=MetricsGroup.TASKS,
+        metrics_type=MetricsType.Histogram,
+        metrics_args={"boundaries": histogram_bucket_rows},
+    )
+
+    # === Actor-related metrics ===
+    num_alive_actors: int = metric_field(
+        default=0,
+        description="Number of alive actors.",
+        metrics_group=MetricsGroup.ACTORS,
+    )
+    num_restarting_actors: int = metric_field(
+        default=0,
+        description="Number of restarting actors.",
+        metrics_group=MetricsGroup.ACTORS,
+    )
+    num_pending_actors: int = metric_field(
+        default=0,
+        description="Number of pending actors.",
+        metrics_group=MetricsGroup.ACTORS,
+    )
+
+    # === Object store memory metrics ===
+    obj_store_mem_internal_inqueue_blocks: int = metric_field(
+        default=0,
+        description="Number of blocks in operator's internal input queue.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    obj_store_mem_internal_outqueue_blocks: int = metric_field(
+        default=0,
+        description="Number of blocks in the operator's internal output queue.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    obj_store_mem_freed: int = metric_field(
+        default=0,
+        description="Byte size of freed memory in object store.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    obj_store_mem_spilled: int = metric_field(
+        default=0,
+        description="Byte size of spilled memory in object store.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    obj_store_mem_used: int = metric_field(
+        default=0,
+        description="Byte size of used memory in object store.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+
+    # === Miscellaneous metrics ===
+    # Use "metrics_group: "misc" in the metadata for new metrics in this section.
+
+    def __init__(self, op: "PhysicalOperator"):
+        from ray.data._internal.execution.operators.map_operator import MapOperator
+
+        self._op = op
+        self._is_map = isinstance(op, MapOperator)
+        self._running_tasks: Dict[int, RunningTaskInfo] = {}
+        self._extra_metrics: Dict[str, Any] = {}
+        # Start time of current pause due to task submission backpressure
+        self._task_submission_backpressure_start_time = -1
+        # Start time of current pause due to task output backpressure
+        self._task_output_backpressure_start_time = -1
+
+        self._internal_inqueue = create_bundle_queue()
+        self._internal_outqueue = create_bundle_queue()
+        self._pending_task_inputs = create_bundle_queue()
+        self._op_task_duration_stats = TaskDurationStats()
+
+        self._per_node_metrics: Dict[str, NodeMetrics] = defaultdict(NodeMetrics)
+        self._per_node_metrics_enabled: bool = op.data_context.enable_per_node_metrics
+
+        self._cum_max_uss_bytes: Optional[int] = None
+        self._issue_detector_hanging = 0
+        self._issue_detector_high_memory = 0
+
+        # Initialize the histogram metrics
+        self.task_completion_time = RuntimeMetricsHistogram(histogram_buckets_s)
+        self.block_completion_time = RuntimeMetricsHistogram(histogram_buckets_s)
+        self.block_size_bytes = RuntimeMetricsHistogram(histogram_buckets_bytes)
+        self.block_size_rows = RuntimeMetricsHistogram(histogram_bucket_rows)
+
+    @property
+    def extra_metrics(self) -> Dict[str, Any]:
+        """Return a dict of extra metrics."""
+        return self._extra_metrics
+
+    @classmethod
+    def get_metrics(self) -> List[MetricDefinition]:
+        return list(_METRICS)
+
+    def as_dict(
+        self,
+        skip_internal_metrics: bool = False,
+    ) -> Dict[str, Any]:
+        """
+        Return a dict representation of the metrics.
+
+        Args:
+            skip_internal_metrics: Whether to skip internal metrics.
+
+        Returns:
+            A dict representation of the metrics.
+        """
+
+        result = []
+        for metric in self.get_metrics():
+            if not self._is_map and metric.map_only:
+                continue
+            if skip_internal_metrics and metric.internal_only:
+                continue
+            value = getattr(self, metric.name)
+            result.append((metric.name, value))
+
+        # TODO: record resource usage in OpRuntimeMetrics,
+        # avoid calling self._op.current_processor_usage()
+        resource_usage = self._op.current_processor_usage()
+        result.extend(
+            [
+                ("cpu_usage", resource_usage.cpu or 0),
+                ("gpu_usage", resource_usage.gpu or 0),
+            ]
+        )
+        result.extend(self._extra_metrics.items())
+        return dict(result)
+
+    @metric_property(
+        description="Average number of blocks generated per task.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    def average_num_outputs_per_task(self) -> Optional[float]:
+        """Average number of output blocks per task, or None if no task has finished."""
+        if self.num_tasks_finished == 0:
+            return None
+        else:
+            return self.num_outputs_of_finished_tasks / self.num_tasks_finished
+
+    @metric_property(
+        description="Average number of blocks generated per task.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    def average_num_inputs_per_task(self) -> Optional[float]:
+        """Average number of input blocks per task, or None if no task has finished."""
+        if self.num_tasks_finished == 0:
+            return None
+        else:
+            return self.num_task_inputs_processed / self.num_tasks_finished
+
+    @metric_property(
+        description="Average number of output blocks per task per second.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    def num_output_blocks_per_task_s(self) -> Optional[float]:
+        """Average number of output blocks per task per second.
+
+        If the operator hasn't produced any output yet, this metric returns `None`.
+        """
+        if self.block_generation_time == 0:
+            return None
+        else:
+            return self.num_task_outputs_generated / self.block_generation_time
+
+    @metric_property(
+        description="Average task's completion time in seconds (including throttling).",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    def average_total_task_completion_time_s(self) -> Optional[float]:
+        """Average task's completion time in seconds (including throttling)"""
+        if self.num_tasks_finished == 0:
+            return None
+        else:
+            return self.task_completion_time_total_s / self.num_tasks_finished
+
+    @metric_property(
+        description="Average task's completion time in seconds (excluding throttling).",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    def average_task_completion_excl_backpressure_time_s(self) -> Optional[float]:
+        """Average task's completion time in seconds (excluding throttling)"""
+        if self.num_tasks_finished == 0:
+            return None
+        else:
+            return (
+                self.task_completion_time_excl_backpressure_s / self.num_tasks_finished
+            )
+
+    @metric_property(
+        description="Average size of task output in bytes.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    def average_bytes_per_output(self) -> Optional[float]:
+        """Average size in bytes of output blocks."""
+        if self.num_task_outputs_generated == 0:
+            return None
+        else:
+            return self.bytes_task_outputs_generated / self.num_task_outputs_generated
+
+    @metric_property(
+        description="Byte size of input blocks in the operator's internal input queue.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    def obj_store_mem_internal_inqueue(self) -> int:
+        return self._internal_inqueue.estimate_size_bytes()
+
+    @metric_property(
+        description=(
+            "Byte size of output blocks in the operator's internal output queue."
+        ),
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    def obj_store_mem_internal_outqueue(self) -> int:
+        return self._internal_outqueue.estimate_size_bytes()
+
+    @metric_property(
+        description="Byte size of input blocks used by pending tasks.",
+        metrics_group=MetricsGroup.OBJECT_STORE_MEMORY,
+    )
+    def obj_store_mem_pending_task_inputs(self) -> int:
+        return self._pending_task_inputs.estimate_size_bytes()
+
+    @property
+    def obj_store_mem_pending_task_outputs(self) -> Optional[float]:
+        """Estimated size in bytes of output blocks in Ray generator buffers.
+
+        If an estimate isn't available, this property returns ``None``.
+        """
+        per_task_output = self.obj_store_mem_max_pending_output_per_task
+        if per_task_output is None:
+            return None
+
+        # Ray Data launches multiple tasks per actor, but only one task runs at a
+        # time per actor. So, the number of actually running tasks is capped by the
+        # number of active actors.
+        from ray.data._internal.execution.operators.actor_pool_map_operator import (
+            ActorPoolMapOperator,
+        )
+
+        num_tasks_running = self.num_tasks_running
+        if isinstance(self._op, ActorPoolMapOperator):
+            num_tasks_running = min(
+                num_tasks_running, self._op._actor_pool.num_active_actors()
+            )
+
+        return num_tasks_running * per_task_output
+
+    @property
+    def obj_store_mem_max_pending_output_per_task(self) -> Optional[float]:
+        """Estimated size in bytes of output blocks in a task's generator buffer."""
+        context = self._op.data_context
+        if context._max_num_blocks_in_streaming_gen_buffer is None:
+            return None
+
+        bytes_per_output = self.average_bytes_per_output
+        # If we don’t have a sample yet and the limit is “unlimited”, we can’t
+        # estimate – just bail out.
+        if bytes_per_output is None:
+            if context.target_max_block_size is None:
+                return None
+            bytes_per_output = context.target_max_block_size
+
+        num_pending_outputs = context._max_num_blocks_in_streaming_gen_buffer
+        if self.average_num_outputs_per_task is not None:
+            num_pending_outputs = min(
+                num_pending_outputs, self.average_num_outputs_per_task
+            )
+        return bytes_per_output * num_pending_outputs
+
+    @metric_property(
+        description="Average size of task inputs in bytes.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    def average_bytes_inputs_per_task(self) -> Optional[float]:
+        """Average size in bytes of ref bundles passed to tasks, or ``None`` if no
+        tasks have been submitted."""
+        if self.num_tasks_submitted == 0:
+            return None
+        else:
+            return self.bytes_inputs_of_submitted_tasks / self.num_tasks_submitted
+
+    @metric_property(
+        description="Average number of rows passed in to the task.",
+        metrics_group=MetricsGroup.INPUTS,
+    )
+    def average_rows_inputs_per_task(self) -> Optional[float]:
+        """Average number of rows in input blocks per task,
+        or None if no task has been submitted."""
+        if self.num_tasks_submitted == 0:
+            return None
+        else:
+            return self.rows_inputs_of_submitted_tasks / self.num_tasks_submitted
+
+    @metric_property(
+        description="Average total output size of task in bytes.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    def average_bytes_outputs_per_task(self) -> Optional[float]:
+        """Average size in bytes of output blocks per task,
+        or None if no task has finished."""
+        if self.num_tasks_finished == 0:
+            return None
+        else:
+            return self.bytes_outputs_of_finished_tasks / self.num_tasks_finished
+
+    @metric_property(
+        description="Average number of rows produced per task.",
+        metrics_group=MetricsGroup.OUTPUTS,
+    )
+    def average_rows_outputs_per_task(self) -> Optional[float]:
+        """Average number of rows in output blocks per task,
+        or None if no task has finished."""
+        if self.num_tasks_finished == 0:
+            return None
+        else:
+            return self.rows_outputs_of_finished_tasks / self.num_tasks_finished
+
+    @metric_property(
+        description="Average USS usage of tasks.",
+        metrics_group=MetricsGroup.TASKS,
+    )
+    def average_max_uss_per_task(self) -> Optional[float]:
+        """Average max USS usage of tasks."""
+        if self._cum_max_uss_bytes is None:
+            return None
+        else:
+            assert self.num_task_outputs_generated > 0, self.num_task_outputs_generated
+            return self._cum_max_uss_bytes / self.num_task_outputs_generated
+
+    @metric_property(
+        description="Indicates if the operator is hanging.",
+        metrics_group=MetricsGroup.MISC,
+        internal_only=True,
+    )
+    def issue_detector_hanging(self) -> int:
+        return self._issue_detector_hanging
+
+    @metric_property(
+        description="Indicates if the operator is using high memory.",
+        metrics_group=MetricsGroup.MISC,
+        internal_only=True,
+    )
+    def issue_detector_high_memory(self) -> int:
+        return self._issue_detector_high_memory
+
+    def on_input_received(self, input: RefBundle):
+        """Callback when the operator receives a new input."""
+        self.num_inputs_received += 1
+        self.num_row_inputs_received += input.num_rows() or 0
+        self.bytes_inputs_received += input.size_bytes()
+
+    def on_input_queued(self, input: RefBundle):
+        """Callback when the operator queues an input."""
+        self.obj_store_mem_internal_inqueue_blocks += len(input.blocks)
+        self._internal_inqueue.add(input)
+
+    def on_input_dequeued(self, input: RefBundle):
+        """Callback when the operator dequeues an input."""
+        self.obj_store_mem_internal_inqueue_blocks -= len(input.blocks)
+        input_size = input.size_bytes()
+        self._internal_inqueue.remove(input)
+        assert self.obj_store_mem_internal_inqueue >= 0, (
+            self._op,
+            self.obj_store_mem_internal_inqueue,
+            input_size,
+        )
+
+    def on_output_queued(self, output: RefBundle):
+        """Callback when an output is queued by the operator."""
+        self.obj_store_mem_internal_outqueue_blocks += len(output.blocks)
+        self._internal_outqueue.add(output)
+
+    def on_output_dequeued(self, output: RefBundle):
+        """Callback when an output is dequeued by the operator."""
+        self.obj_store_mem_internal_outqueue_blocks -= len(output.blocks)
+        output_size = output.size_bytes()
+        self._internal_outqueue.remove(output)
+        assert self.obj_store_mem_internal_outqueue >= 0, (
+            self._op,
+            self.obj_store_mem_internal_outqueue,
+            output_size,
+        )
+
+    def on_toggle_task_submission_backpressure(self, in_backpressure):
+        if in_backpressure and self._task_submission_backpressure_start_time == -1:
+            # backpressure starting, start timer
+            self._task_submission_backpressure_start_time = time.perf_counter()
+        elif self._task_submission_backpressure_start_time != -1:
+            # backpressure stopping, stop timer
+            self.task_submission_backpressure_time += (
+                time.perf_counter() - self._task_submission_backpressure_start_time
+            )
+            self._task_submission_backpressure_start_time = -1
+
+    def on_toggle_task_output_backpressure(self, in_backpressure):
+        if in_backpressure and self._task_output_backpressure_start_time == -1:
+            # backpressure starting, start timer
+            self._task_output_backpressure_start_time = time.perf_counter()
+        elif self._task_output_backpressure_start_time != -1:
+            # backpressure stopping, stop timer
+            delta = time.perf_counter() - self._task_output_backpressure_start_time
+            self.task_output_backpressure_time += delta
+            self._task_output_backpressure_start_time = -1
+
+    def on_output_taken(self, output: RefBundle):
+        """Callback when an output is taken from the operator."""
+        self.num_outputs_taken += 1
+        self.block_outputs_taken += len(output)
+        self.row_outputs_taken += output.num_rows() or 0
+        self.bytes_outputs_taken += output.size_bytes()
+
+    def on_task_submitted(self, task_index: int, inputs: RefBundle):
+        """Callback when the operator submits a task."""
+        self.num_tasks_submitted += 1
+        self.num_tasks_running += 1
+        self.bytes_inputs_of_submitted_tasks += inputs.size_bytes()
+        self.rows_inputs_of_submitted_tasks += inputs.num_rows() or 0
+        self._pending_task_inputs.add(inputs)
+        self._running_tasks[task_index] = RunningTaskInfo(
+            inputs=inputs,
+            num_outputs=0,
+            bytes_outputs=0,
+            num_rows_produced=0,
+            start_time=time.perf_counter(),
+            cum_block_gen_time=0,
+        )
+
+    def on_task_output_generated(self, task_index: int, output: RefBundle):
+        """Callback when a new task generates an output."""
+        num_outputs = len(output)
+        output_bytes = output.size_bytes()
+        num_rows_produced = output.num_rows()
+
+        self.num_task_outputs_generated += num_outputs
+        self.bytes_task_outputs_generated += output_bytes
+        self.rows_task_outputs_generated += num_rows_produced
+        for block in output.metadata:
+            if block.size_bytes is not None:
+                self.block_size_bytes.observe(block.size_bytes)
+            if block.num_rows is not None:
+                self.block_size_rows.observe(block.num_rows)
+
+        task_info = self._running_tasks[task_index]
+        if task_info.num_outputs == 0:
+            self.num_tasks_have_outputs += 1
+
+        task_info.num_outputs += num_outputs
+        task_info.bytes_outputs += output_bytes
+        task_info.num_rows_produced += num_rows_produced
+
+        for block_ref, meta in output.blocks:
+            assert (
+                meta.exec_stats is not None and meta.exec_stats.wall_time_s is not None
+            )
+            self.block_generation_time += meta.exec_stats.wall_time_s
+            task_info.cum_block_gen_time += meta.exec_stats.wall_time_s
+            assert meta.num_rows is not None
+            trace_allocation(block_ref, "operator_output")
+            if meta.exec_stats.max_uss_bytes is not None:
+                if self._cum_max_uss_bytes is None:
+                    self._cum_max_uss_bytes = meta.exec_stats.max_uss_bytes
+                else:
+                    self._cum_max_uss_bytes += meta.exec_stats.max_uss_bytes
+
+        # Update per node metrics
+        if self._per_node_metrics_enabled:
+            for _, meta in output.blocks:
+                node_id = node_id_from_block_metadata(meta)
+                node_metrics = self._per_node_metrics[node_id]
+
+                node_metrics.bytes_outputs_of_finished_tasks += meta.size_bytes
+                node_metrics.blocks_outputs_of_finished_tasks += 1
+
+    def on_task_finished(self, task_index: int, exception: Optional[Exception]):
+        """Callback when a task is finished."""
+        self.num_tasks_running -= 1
+        self.num_tasks_finished += 1
+        if exception is not None:
+            self.num_tasks_failed += 1
+
+        task_info = self._running_tasks[task_index]
+
+        self.num_outputs_of_finished_tasks += task_info.num_outputs
+        self.bytes_outputs_of_finished_tasks += task_info.bytes_outputs
+        self.rows_outputs_of_finished_tasks += task_info.num_rows_produced
+
+        task_time_delta = time.perf_counter() - task_info.start_time
+        self.task_completion_time_total_s += task_time_delta
+        self.task_completion_time.observe(task_time_delta)
+
+        assert task_info.cum_block_gen_time is not None
+        if task_info.num_outputs > 0:
+            # Calculate the average block generation time per block
+            block_time_delta = task_info.cum_block_gen_time / task_info.num_outputs
+            self.block_completion_time.observe(
+                block_time_delta, num_observations=task_info.num_outputs
+            )
+
+        # NOTE: This is used for Issue Detection
+        self._op_task_duration_stats.add_duration(task_time_delta)
+
+        self.task_completion_time_excl_backpressure_s += task_info.cum_block_gen_time
+        inputs = self._running_tasks[task_index].inputs
+        self.num_task_inputs_processed += len(inputs)
+        total_input_size = inputs.size_bytes()
+        self.bytes_task_inputs_processed += total_input_size
+        input_size = inputs.size_bytes()
+        self._pending_task_inputs.remove(inputs)
+        assert self.obj_store_mem_pending_task_inputs >= 0, (
+            self._op,
+            self.obj_store_mem_pending_task_inputs,
+            input_size,
+        )
+
+        ctx = self._op.data_context
+        if ctx.enable_get_object_locations_for_metrics:
+            locations = ray.experimental.get_object_locations(inputs.block_refs)
+            for block, meta in inputs.blocks:
+                if locations[block].get("did_spill", False):
+                    assert meta.size_bytes is not None
+                    self.obj_store_mem_spilled += meta.size_bytes
+
+        self.obj_store_mem_freed += total_input_size
+
+        # Update per node metrics
+        if self._per_node_metrics_enabled:
+            node_ids = set()
+            for _, meta in inputs.blocks:
+                node_id = node_id_from_block_metadata(meta)
+                node_metrics = self._per_node_metrics[node_id]
+
+                # Stats to update once per node id or if node id is unknown
+                if node_id not in node_ids or node_id == NODE_UNKNOWN:
+                    node_metrics.num_tasks_finished += 1
+
+                # Keep track of node ids to ensure we don't double count
+                node_ids.add(node_id)
+
+        inputs.destroy_if_owned()
+        del self._running_tasks[task_index]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/physical_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/physical_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d2b4f6a86fe6a2c0dfbf3b6d5b1ce3779df7dac
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/physical_operator.py
@@ -0,0 +1,881 @@
+import abc
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import ray
+from .ref_bundle import RefBundle
+from ray._raylet import ObjectRefGenerator
+from ray.data._internal.actor_autoscaler.autoscaling_actor_pool import (
+    AutoscalingActorPool,
+)
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.data._internal.execution.interfaces.op_runtime_metrics import OpRuntimeMetrics
+from ray.data._internal.logical.interfaces import LogicalOperator, Operator
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.stats import StatsDict, Timer
+from ray.data.block import Block, BlockMetadata
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+
+    from ray.data.block import BlockMetadataWithSchema
+
+logger = logging.getLogger(__name__)
+
+# Timeout for getting metadata from Ray object references (in seconds)
+METADATA_GET_TIMEOUT_S = 1.0
+
+# Timeout for waiting for metadata object to become available (in seconds)
+METADATA_WAIT_TIMEOUT_S = 0.1
+
+# TODO(hchen): Ray Core should have a common interface for these two types.
+Waitable = Union[ray.ObjectRef, ObjectRefGenerator]
+
+
+class OpTask(ABC):
+    """Abstract class that represents a task that is created by an PhysicalOperator.
+
+    The task can be either a regular task or an actor task.
+    """
+
+    def __init__(
+        self,
+        task_index: int,
+        task_resource_bundle: Optional[ExecutionResources] = None,
+    ):
+        self._task_index: int = task_index
+        self._task_resource_bundle: Optional[ExecutionResources] = task_resource_bundle
+
+    def task_index(self) -> int:
+        """Return the index of the task."""
+        return self._task_index
+
+    def get_requested_resource_bundle(self) -> Optional[ExecutionResources]:
+        return self._task_resource_bundle
+
+    @abstractmethod
+    def get_waitable(self) -> Waitable:
+        """Return the ObjectRef or ObjectRefGenerator to wait on."""
+        ...
+
+    def _cancel(self, force: bool):
+        object_ref = self.get_waitable()
+
+        # Get generator's `ObjectRef`
+        if isinstance(object_ref, ObjectRefGenerator):
+            object_ref = object_ref._generator_ref
+
+        is_actor_task = not object_ref.task_id().actor_id().is_nil()
+
+        ray.cancel(
+            object_ref,
+            recursive=True,
+            # NOTE: Actor tasks can't be force-cancelled
+            force=force and not is_actor_task,
+        )
+
+
+class DataOpTask(OpTask):
+    """Represents an OpTask that handles Block data."""
+
+    def __init__(
+        self,
+        task_index: int,
+        streaming_gen: ObjectRefGenerator,
+        output_ready_callback: Callable[[RefBundle], None] = lambda bundle: None,
+        task_done_callback: Callable[[Optional[Exception]], None] = lambda exc: None,
+        block_ready_callback: Callable[
+            [ray.ObjectRef[Block]], None
+        ] = lambda block_ref: None,
+        metadata_ready_callback: Callable[
+            [ray.ObjectRef[BlockMetadata]], None
+        ] = lambda metadata_ref: None,
+        task_resource_bundle: Optional[ExecutionResources] = None,
+    ):
+        """Create a DataOpTask
+        Args:
+            task_index: Index of the task. Used for callbacks.
+            streaming_gen: The streaming generator of this task. It should yield blocks.
+            output_ready_callback: The callback to call when a new RefBundle is output
+                from the generator.
+            task_done_callback: The callback to call when the task is done.
+            block_ready_callback: A callback that's invoked when a new block reference
+                is ready. This is exposed as a seam for testing.
+            metadata_ready_callback: A callback that's invoked when a new block metadata
+                reference is ready. This is exposed as a seam for testing.
+            task_resource_bundle: The execution resources of this task.
+        """
+        super().__init__(task_index, task_resource_bundle)
+        # TODO(hchen): Right now, the streaming generator is required to yield a Block
+        # and a BlockMetadata each time. We should unify task submission with an unified
+        # interface. So each individual operator don't need to take care of the
+        # BlockMetadata.
+        self._streaming_gen = streaming_gen
+        self._output_ready_callback = output_ready_callback
+        self._task_done_callback = task_done_callback
+        self._block_ready_callback = block_ready_callback
+        self._metadata_ready_callback = metadata_ready_callback
+
+        # If the generator hasn't produced block metadata yet, or if the block metadata
+        # object isn't available after we get a reference, we need store the pending
+        # references and wait until Ray (re)constructs the block metadata. Either case
+        # can happen if a node dies after producing a block.
+        self._pending_block_ref: ray.ObjectRef[Block] = ray.ObjectRef.nil()
+        self._pending_meta_ref: ray.ObjectRef[BlockMetadata] = ray.ObjectRef.nil()
+
+        self._has_finished = False
+
+    def get_waitable(self) -> ObjectRefGenerator:
+        return self._streaming_gen
+
+    def on_data_ready(self, max_bytes_to_read: Optional[int]) -> int:
+        """Callback when data is ready to be read from the streaming generator.
+
+        Args:
+            max_bytes_to_read: Max bytes of blocks to read. If None, all available
+                will be read.
+        Returns: The number of blocks read.
+        """
+        bytes_read = 0
+        while max_bytes_to_read is None or bytes_read < max_bytes_to_read:
+            if self._pending_block_ref.is_nil():
+                assert self._pending_meta_ref.is_nil(), (
+                    "This method expects streaming generators to yield blocks then "
+                    "metadata. So, if we have a reference to metadata but not the "
+                    "block, it means there's an error in the implementation."
+                )
+
+                try:
+                    self._pending_block_ref = self._streaming_gen._next_sync(
+                        timeout_s=0
+                    )
+                except StopIteration:
+                    self._task_done_callback(None)
+                    self._has_finished = True
+                    break
+
+                if self._pending_block_ref.is_nil():
+                    # The generator currently doesn't have new output.
+                    # And it's not stopped yet.
+                    break
+
+                self._block_ready_callback(self._pending_block_ref)
+
+            if self._pending_meta_ref.is_nil():
+                try:
+                    self._pending_meta_ref = self._streaming_gen._next_sync(
+                        timeout_s=METADATA_WAIT_TIMEOUT_S
+                    )
+                except StopIteration:
+                    # The generator should always yield 2 values (block and metadata)
+                    # each time. If we get a StopIteration here, it means an error
+                    # happened in the task.
+                    # And in this case, the block_ref is the exception object.
+                    # TODO(hchen): Ray Core should have a better interface for
+                    # detecting and obtaining the exception.
+                    try:
+                        ray.get(self._pending_block_ref)
+                        assert False, "Above ray.get should raise an exception."
+                    except Exception as ex:
+                        self._task_done_callback(ex)
+                        self._has_finished = True
+                        raise ex from None
+
+                if self._pending_meta_ref.is_nil():
+                    # We have a reference to the block but the metadata isn't ready
+                    # yet.
+                    break
+
+                self._metadata_ready_callback(self._pending_meta_ref)
+
+            try:
+                # The timeout for `ray.get` includes the time required to ship the
+                # block metadata to this node. So, if we set the timeout to 0, `ray.get`
+                # will timeout and possible cancel the download. To avoid this issue,
+                # we set the timeout to a small non-zero value.
+                meta_with_schema: "BlockMetadataWithSchema" = ray.get(
+                    self._pending_meta_ref, timeout=METADATA_GET_TIMEOUT_S
+                )
+            except ray.exceptions.GetTimeoutError:
+                # We have a reference to the block and its metadata, but the metadata
+                # object isn't available. This can happen if the node dies.
+                logger.warning(
+                    f"Metadata object not ready for "
+                    f"ref={self._pending_meta_ref.hex()} "
+                    f"(operator={self.__class__.__name__}). "
+                    f"Metadata may still be computing or worker may have failed and "
+                    f"object is being reconstructed. Will retry in next iteration."
+                )
+                break
+
+            meta = meta_with_schema.metadata
+            self._output_ready_callback(
+                RefBundle(
+                    [(self._pending_block_ref, meta)],
+                    owns_blocks=True,
+                    schema=meta_with_schema.schema,
+                ),
+            )
+            self._pending_block_ref = ray.ObjectRef.nil()
+            self._pending_meta_ref = ray.ObjectRef.nil()
+
+            bytes_read += meta.size_bytes
+
+        return bytes_read
+
+    @property
+    def has_finished(self) -> bool:
+        return self._has_finished
+
+
+class MetadataOpTask(OpTask):
+    """Represents an OpTask that only handles metadata, instead of Block data."""
+
+    def __init__(
+        self,
+        task_index: int,
+        object_ref: ray.ObjectRef,
+        task_done_callback: Callable[[], None],
+        task_resource_bundle: Optional[ExecutionResources] = None,
+    ):
+        """
+        Args:
+            object_ref: The ObjectRef of the task.
+            task_done_callback: The callback to call when the task is done.
+        """
+        super().__init__(task_index, task_resource_bundle)
+        self._object_ref = object_ref
+        self._task_done_callback = task_done_callback
+
+    def get_waitable(self) -> ray.ObjectRef:
+        return self._object_ref
+
+    def on_task_finished(self):
+        """Callback when the task is finished."""
+        self._task_done_callback()
+
+
+@dataclass
+class _ActorPoolInfo:
+    """Breakdown of the state of the actors used by the ``PhysicalOperator``"""
+
+    running: int
+    pending: int
+    restarting: int
+
+    def __str__(self):
+        return (
+            f"running={self.running}, restarting={self.restarting}, "
+            f"pending={self.pending}"
+        )
+
+
+class PhysicalOperator(Operator):
+    """Abstract class for physical operators.
+
+    An operator transforms one or more input streams of RefBundles into a single
+    output stream of RefBundles.
+
+    Physical operators are stateful and non-serializable; they live on the driver side
+    of the Dataset only.
+
+    Here's a simple example of implementing a basic "Map" operator:
+
+        class MapOperator(PhysicalOperator):
+            def __init__(self):
+                self.active_tasks = []
+
+            def add_input(self, refs, _):
+                self.active_tasks.append(map_task.remote(refs))
+
+            def has_next(self):
+                ready, _ = ray.wait(self.active_tasks, timeout=0)
+                return len(ready) > 0
+
+            def get_next(self):
+                ready, remaining = ray.wait(self.active_tasks, num_returns=1)
+                self.active_tasks = remaining
+                return ready[0]
+
+    Note that the above operator fully supports both bulk and streaming execution,
+    since `add_input` and `get_next` can be called in any order. In bulk execution
+    (now deprecated), all inputs would be added up-front, but in streaming
+    execution (now the default execution mode) the calls could be interleaved.
+    """
+
+    _OPERATOR_ID_LABEL_KEY = "__data_operator_id"
+
+    def __init__(
+        self,
+        name: str,
+        input_dependencies: List["PhysicalOperator"],
+        data_context: DataContext,
+        target_max_block_size_override: Optional[int] = None,
+    ):
+        super().__init__(name, input_dependencies)
+
+        for x in input_dependencies:
+            assert isinstance(x, PhysicalOperator), x
+        self._inputs_complete = not input_dependencies
+        self._output_block_size_option_override = OutputBlockSizeOption.of(
+            target_max_block_size=target_max_block_size_override
+        )
+        self._started = False
+        self._shutdown = False
+        self._in_task_submission_backpressure = False
+        self._in_task_output_backpressure = False
+        self._estimated_num_output_bundles = None
+        self._estimated_output_num_rows = None
+        self._execution_finished = False
+        # The LogicalOperator(s) which were translated to create this PhysicalOperator.
+        # Set via `PhysicalOperator.set_logical_operators()`.
+        self._logical_operators: List[LogicalOperator] = []
+        self._data_context = data_context
+        self._id = str(uuid.uuid4())
+        # Initialize metrics after data_context is set
+        self._metrics = OpRuntimeMetrics(self)
+
+    def __reduce__(self):
+        raise ValueError("Operator is not serializable.")
+
+    @property
+    def id(self) -> str:
+        """Return a unique identifier for this operator."""
+        return self._id
+
+    @property
+    def data_context(self) -> DataContext:
+        return self._data_context
+
+    # Override the following 3 methods to correct type hints.
+
+    @property
+    def input_dependencies(self) -> List["PhysicalOperator"]:
+        return super().input_dependencies  # type: ignore
+
+    @property
+    def output_dependencies(self) -> List["PhysicalOperator"]:
+        return super().output_dependencies  # type: ignore
+
+    def post_order_iter(self) -> Iterator["PhysicalOperator"]:
+        return super().post_order_iter()  # type: ignore
+
+    def set_logical_operators(
+        self,
+        *logical_ops: LogicalOperator,
+    ):
+        self._logical_operators = list(logical_ops)
+
+    @property
+    def target_max_block_size_override(self) -> Optional[int]:
+        """
+        Target max block size output by this operator. If this returns None,
+        then the default from DataContext should be used.
+        """
+        if self._output_block_size_option_override is None:
+            return None
+        else:
+            return self._output_block_size_option_override.target_max_block_size
+
+    def override_target_max_block_size(self, target_max_block_size: Optional[int]):
+        self._output_block_size_option_override = OutputBlockSizeOption.of(
+            target_max_block_size=target_max_block_size
+        )
+
+    def mark_execution_finished(self):
+        """Manually mark that this operator has finished execution."""
+        self._execution_finished = True
+
+    def execution_finished(self) -> bool:
+        """Return True when this operator has finished execution.
+
+        The outputs may or may not have been taken.
+        """
+        return self._execution_finished
+
+    def completed(self) -> bool:
+        """Returns whether this operator has been fully completed.
+
+        An operator is completed iff:
+            * The operator has finished execution (i.e., `execution_finished()` is True).
+            * All outputs have been taken (i.e., `has_next()` is False) from it.
+        """
+        from ..operators.base_physical_operator import InternalQueueOperatorMixin
+
+        internal_input_queue_num_blocks = 0
+        internal_output_queue_num_blocks = 0
+        if isinstance(self, InternalQueueOperatorMixin):
+            internal_input_queue_num_blocks = self.internal_input_queue_num_blocks()
+            internal_output_queue_num_blocks = self.internal_output_queue_num_blocks()
+
+        if not self._execution_finished:
+            if (
+                self._inputs_complete
+                and internal_input_queue_num_blocks == 0
+                and self.num_active_tasks() == 0
+            ):
+                # NOTE: Operator is considered completed iff
+                #   - All input blocks have been ingested
+                #   - Internal queue is empty
+                #   - There are no active or pending tasks
+                self._execution_finished = True
+
+        # NOTE: We check for (internal_output_queue_size == 0) and
+        # (not self.has_next()) because _OrderedOutputQueue can
+        # return False for self.has_next(), but have a non-empty queue size.
+        # Draining the internal output queue is important to free object refs.
+        return (
+            self._execution_finished
+            and not self.has_next()
+            and internal_output_queue_num_blocks == 0
+        )
+
+    def get_stats(self) -> StatsDict:
+        """Return recorded execution stats for use with DatasetStats."""
+        raise NotImplementedError
+
+    @property
+    def metrics(self) -> OpRuntimeMetrics:
+        """Returns the runtime metrics of this operator."""
+        self._metrics._extra_metrics = self._extra_metrics()
+        return self._metrics
+
+    def _extra_metrics(self) -> Dict[str, Any]:
+        """Subclasses should override this method to report extra metrics
+        that are specific to them."""
+        return {}
+
+    def _get_logical_args(self) -> Dict[str, Dict[str, Any]]:
+        """Return the logical arguments that were translated to create this
+        PhysicalOperator."""
+        res = {}
+        for i, logical_op in enumerate(self._logical_operators):
+            logical_op_id = f"{logical_op}_{i}"
+            res[logical_op_id] = logical_op._get_args()
+        return res
+
+    # TODO(@balaji): Disambiguate this with `incremental_resource_usage`.
+    def per_task_resource_allocation(
+        self: "PhysicalOperator",
+    ) -> ExecutionResources:
+        """The amount of logical resources used by each task.
+
+        For regular tasks, these are the resources required to schedule a task. For
+        actor tasks, these are the resources required to schedule an actor divided by
+        the number of actor threads (i.e., `max_concurrency`).
+
+        Returns:
+            The resource requirement per task.
+        """
+        return ExecutionResources.zero()
+
+    def max_task_concurrency(self: "PhysicalOperator") -> Optional[int]:
+        """The maximum number of tasks that can be run concurrently.
+
+        Some operators manually configure a maximum concurrency. For example, if you
+        specify `concurrency` in `map_batches`.
+        """
+        return None
+
+    # TODO(@balaji): Disambiguate this with `base_resource_usage`.
+    def min_scheduling_resources(
+        self: "PhysicalOperator",
+    ) -> ExecutionResources:
+        """The minimum resource bundle required to schedule a worker.
+
+        For regular tasks, this is the resources required to schedule a task. For actor
+        tasks, this is the resources required to schedule an actor.
+        """
+        return ExecutionResources.zero()
+
+    def progress_str(self) -> str:
+        """Return any extra status to be displayed in the operator progress bar.
+
+        For example, `<N> actors` to show current number of actors in an actor pool.
+        """
+        return ""
+
+    def num_outputs_total(self) -> Optional[int]:
+        """Returns the total number of output bundles of this operator,
+        or ``None`` if unable to provide a reasonable estimate (for example,
+        if no tasks have finished yet).
+
+        The value returned may be an estimate based off the consumption so far.
+        This is useful for reporting progress.
+
+        Subclasses should either override this method, or update
+        ``self._estimated_num_output_bundles`` appropriately.
+        """
+        return self._estimated_num_output_bundles
+
+    def num_output_rows_total(self) -> Optional[int]:
+        """Returns the total number of output rows of this operator,
+        or ``None`` if unable to provide a reasonable estimate (for example,
+        if no tasks have finished yet).
+
+        The value returned may be an estimate based off the consumption so far.
+        This is useful for reporting progress.
+
+        Subclasses should either override this method, or update
+        ``self._estimated_output_num_rows`` appropriately.
+        """
+        return self._estimated_output_num_rows
+
+    def start(self, options: ExecutionOptions) -> None:
+        """Called by the executor when execution starts for an operator.
+
+        Args:
+            options: The global options used for the overall execution.
+        """
+        self._started = True
+
+    def should_add_input(self) -> bool:
+        """Return whether it is desirable to add input to this operator right now.
+
+        Operators can customize the implementation of this method to apply additional
+        backpressure (e.g., waiting for internal actors to be created).
+        """
+        return True
+
+    def add_input(self, refs: RefBundle, input_index: int) -> None:
+        """Called when an upstream result is available.
+
+        Inputs may be added in any order, and calls to `add_input` may be interleaved
+        with calls to `get_next` / `has_next` to implement streaming execution.
+
+        Subclasses should override `_add_input_inner` instead of this method.
+
+        Args:
+            refs: The ref bundle that should be added as input.
+            input_index: The index identifying the input dependency producing the
+                input. For most operators, this is always `0` since there is only
+                one upstream input operator.
+        """
+        assert 0 <= input_index < len(self._input_dependencies), (
+            f"Input index out of bounds (total inputs {len(self._input_dependencies)}, "
+            f"index is {input_index})"
+        )
+
+        self._metrics.on_input_received(refs)
+        self._add_input_inner(refs, input_index)
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        """Subclasses should override this method to implement `add_input`."""
+        raise NotImplementedError
+
+    def input_done(self, input_index: int) -> None:
+        """Called when the upstream operator at index `input_index` has completed().
+
+        After this is called, the executor guarantees that no more inputs will be added
+        via `add_input` for the given input index.
+        """
+        pass
+
+    def all_inputs_done(self) -> None:
+        """Called when all upstream operators have completed().
+
+        After this is called, the executor guarantees that no more inputs will be added
+        via `add_input` for any input index.
+        """
+        self._inputs_complete = True
+
+    def has_next(self) -> bool:
+        """Returns when a downstream output is available.
+
+        When this returns true, it is safe to call `get_next()`.
+        """
+        raise NotImplementedError
+
+    def get_next(self) -> RefBundle:
+        """Get the next downstream output.
+
+        It is only allowed to call this if `has_next()` has returned True.
+
+        Subclasses should override `_get_next_inner` instead of this method.
+        """
+        output = self._get_next_inner()
+        self._metrics.on_output_taken(output)
+        return output
+
+    def _get_next_inner(self) -> RefBundle:
+        """Subclasses should override this method to implement `get_next`."""
+        raise NotImplementedError
+
+    def get_active_tasks(self) -> List[OpTask]:
+        """Get a list of the active tasks of this operator.
+
+        Subclasses should return *all* running normal/actor tasks. The
+        StreamingExecutor will wait on these tasks and trigger callbacks.
+        """
+        return []
+
+    def num_active_tasks(self) -> int:
+        """Return the number of active tasks.
+
+        This method is used for 2 purposes:
+        * Determine if this operator is completed.
+        * Displaying active task info in the progress bar.
+        Thus, the return value can be less than `len(get_active_tasks())`,
+        if some tasks are not needed for the above purposes. E.g., for the
+        actor pool map operator, readiness checking tasks can be excluded
+        from `num_active_tasks`, but they should be included in
+        `get_active_tasks`.
+
+        Subclasses can override this as a performance optimization.
+        """
+        return len(self.get_active_tasks())
+
+    def throttling_disabled(self) -> bool:
+        """Whether to disable resource throttling for this operator.
+
+        This should return True for operators that only manipulate bundle metadata
+        (e.g., the OutputSplitter operator). This hints to the execution engine that
+        these operators should not be throttled based on resource usage.
+        """
+        return False
+
+    def shutdown(self, timer: Timer, force: bool = False) -> None:
+        """Abort execution and release all resources used by this operator.
+
+        This release any Ray resources acquired by this operator such as active
+        tasks, actors, and objects.
+        """
+        if self._shutdown:
+            return
+        elif not self._started:
+            raise ValueError("Operator must be started before being shutdown.")
+
+        # Mark operator as shut down
+        self._shutdown = True
+        # Time shutdown sequence duration
+        with timer.timer():
+            self._do_shutdown(force)
+
+    def _do_shutdown(self, force: bool):
+        # Default implementation simply cancels any outstanding active task
+        self._cancel_active_tasks(force=force)
+
+    def current_processor_usage(self) -> ExecutionResources:
+        """Returns the current estimated CPU and GPU usage of this operator, excluding
+        object store memory.
+
+        This method is called by the executor to decide how to allocate processors
+        between different operators.
+        """
+        return ExecutionResources(0, 0, 0)
+
+    def running_processor_usage(self) -> ExecutionResources:
+        """Returns the estimated running CPU and GPU usage of this operator, excluding
+        object store memory.
+
+        This method is called by the resource manager and the streaming
+        executor to display the number of currently running CPUs and GPUs in the
+        progress bar.
+
+        Note, this method returns `current_processor_usage() -
+        pending_processor_usage()` by default. Subclasses should only override
+        `pending_processor_usage()` if needed.
+        """
+        usage = self.current_processor_usage()
+        usage = usage.subtract(self.pending_processor_usage())
+        return usage
+
+    def pending_processor_usage(self) -> ExecutionResources:
+        """Returns the estimated pending CPU and GPU usage of this operator, excluding
+        object store memory.
+
+        This method is called by the resource manager and the streaming
+        executor to display the number of currently pending actors in the
+        progress bar.
+        """
+        return ExecutionResources(0, 0, 0)
+
+    def min_max_resource_requirements(
+        self,
+    ) -> Tuple[ExecutionResources, ExecutionResources]:
+        """Returns lower/upper boundary of resource requirements for this operator:
+
+        - Minimal: lower bound (min) of resources required to start this operator
+        (for most operators this is 0, except the ones that utilize actors)
+        - Maximum: upper bound (max) of how many resources this operator could
+        utilize.
+        """
+        return ExecutionResources.zero(), ExecutionResources.inf()
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        """Returns the incremental resources required for processing another input.
+
+        For example, an operator that launches a task per input could return
+        ExecutionResources(cpu=1) as its incremental usage.
+        """
+        return ExecutionResources()
+
+    def notify_in_task_submission_backpressure(self, in_backpressure: bool) -> None:
+        """Called periodically from the executor to update internal in backpressure
+        status for stats collection purposes.
+
+        Args:
+            in_backpressure: Value this operator's in_backpressure should be set to.
+        """
+        # only update on change to in_backpressure
+        if self._in_task_submission_backpressure != in_backpressure:
+            self._metrics.on_toggle_task_submission_backpressure(in_backpressure)
+            self._in_task_submission_backpressure = in_backpressure
+
+    def notify_in_task_output_backpressure(self, in_backpressure: bool) -> None:
+        """Called periodically from the executor to update internal output backpressure
+        status for stats collection purposes.
+
+        Args:
+            in_backpressure: Value this operator's output backpressure should be set to.
+        """
+        # only update on change to in_backpressure
+        if self._in_task_output_backpressure != in_backpressure:
+            self._metrics.on_toggle_task_output_backpressure(in_backpressure)
+            self._in_task_output_backpressure = in_backpressure
+
+    def get_autoscaling_actor_pools(self) -> List[AutoscalingActorPool]:
+        """Return a list of `AutoscalingActorPool`s managed by this operator."""
+        return []
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        """Return whether this operator implements accurate memory accounting.
+
+        An operator that implements accurate memory accounting should properly
+        report its memory usage via the following APIs:
+          - `self._metrics.on_input_queued`.
+          - `self._metrics.on_input_dequeued`.
+          - `self._metrics.on_output_queued`.
+          - `self._metrics.on_output_dequeued`.
+        """
+        # TODO(hchen): Currently we only enable `ReservationOpResourceAllocator` when
+        # all operators in the dataset have implemented accurate memory accounting.
+        # Eventually all operators should implement accurate memory accounting.
+        return False
+
+    def supports_fusion(self) -> bool:
+        """Returns ```True``` if this operator can be fused with other operators."""
+        return False
+
+    def update_resource_usage(self) -> None:
+        """Updates resource usage of this operator at runtime.
+
+        This method will be called at runtime in each StreamingExecutor iteration.
+        Subclasses can override it to account for dynamic resource usage updates due to
+        restarting actors, retrying tasks, lost objects, etc.
+        """
+        pass
+
+    def get_actor_info(self) -> _ActorPoolInfo:
+        """Returns the current status of actors being used by the operator"""
+        return _ActorPoolInfo(running=0, pending=0, restarting=0)
+
+    def _cancel_active_tasks(self, force: bool):
+        tasks: List[OpTask] = self.get_active_tasks()
+
+        # Interrupt all (still) running tasks immediately
+        for task in tasks:
+            task._cancel(force=force)
+
+        # In case of forced cancellation block until task actually return
+        # to guarantee all tasks are done upon return from this method
+        if force:
+            # Wait for all tasks to get cancelled before returning
+            for task in tasks:
+                try:
+                    ray.get(task.get_waitable())
+                except ray.exceptions.RayError:
+                    # Cancellation either succeeded, or the task might have already
+                    # failed with a different error, or cancellation failed.
+                    # In all cases, we swallow the exception.
+                    pass
+
+    def upstream_op_num_outputs(self):
+        upstream_op_num_outputs = sum(
+            op.num_outputs_total() or 0 for op in self.input_dependencies
+        )
+        return upstream_op_num_outputs
+
+    def get_max_concurrency_limit(self) -> Optional[int]:
+        """Max value of how many tasks this operator could run
+        concurrently (if limited)"""
+        return None
+
+
+class ReportsExtraResourceUsage(abc.ABC):
+    @abc.abstractmethod
+    def extra_resource_usage(self: PhysicalOperator) -> ExecutionResources:
+        """Returns resources used by this operator beyond standard accounting."""
+        ...
+
+
+def estimate_total_num_of_blocks(
+    num_tasks_submitted: int,
+    upstream_op_num_outputs: int,
+    metrics: OpRuntimeMetrics,
+    total_num_tasks: Optional[int] = None,
+) -> Tuple[int, int, int]:
+    """This method is trying to estimate total number of blocks/rows based on
+    - How many outputs produced by the input deps
+    - How many blocks/rows produced by tasks of this operator
+    """
+
+    if (
+        upstream_op_num_outputs > 0
+        and metrics.average_num_inputs_per_task
+        and metrics.average_num_outputs_per_task
+        and metrics.average_rows_outputs_per_task
+    ):
+        estimated_num_tasks = total_num_tasks
+        if estimated_num_tasks is None:
+            estimated_num_tasks = (
+                upstream_op_num_outputs / metrics.average_num_inputs_per_task
+            )
+
+        estimated_num_output_bundles = round(
+            estimated_num_tasks * metrics.average_num_outputs_per_task
+        )
+        estimated_output_num_rows = round(
+            estimated_num_tasks * metrics.average_rows_outputs_per_task
+        )
+
+        return (
+            estimated_num_tasks,
+            estimated_num_output_bundles,
+            estimated_output_num_rows,
+        )
+
+    return (0, 0, 0)
+
+
+def _create_sub_pb(
+    name: str, total_output_rows: Optional[int], position: int
+) -> Tuple[ProgressBar, int]:
+    progress_bar = ProgressBar(
+        name,
+        total_output_rows or 1,
+        unit="row",
+        position=position,
+    )
+    # NOTE: call `set_description` to trigger the initial print of progress
+    # bar on console.
+    progress_bar.set_description(f"  *- {name}")
+    position += 1
+    return progress_bar, position
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/ref_bundle.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/ref_bundle.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b011eac4b896083eaf857d0e18658c17776e617
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/ref_bundle.py
@@ -0,0 +1,399 @@
+import itertools
+import math
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple
+
+import ray
+from .common import NodeIdStr
+from ray.data._internal.memory_tracing import trace_deallocation
+from ray.data.block import Block, BlockAccessor, BlockMetadata, Schema
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+
+@dataclass
+class BlockSlice:
+    """A slice of a block."""
+
+    # Starting row offset (inclusive) within the block.
+    start_offset: int
+    # Ending row offset (exclusive) within the block.
+    end_offset: int
+
+    @property
+    def num_rows(self) -> int:
+        return self.end_offset - self.start_offset
+
+
+@dataclass
+class RefBundle:
+    """A group of data block references and their metadata.
+
+    Operators take in and produce streams of RefBundles.
+
+    Most commonly a RefBundle consists of a single block object reference.
+    In some cases, e.g., due to block splitting, or for a reduce task, there may
+    be more than one block.
+
+    Block bundles have ownership semantics, i.e., shared ownership (similar to C++
+    shared_ptr, multiple operators share the same block bundle), or unique ownership
+    (similar to C++ unique_ptr, only one operator owns the block bundle). This
+    allows operators to know whether they can destroy blocks when they don't need
+    them. Destroying blocks eagerly is more efficient than waiting for Python GC /
+    Ray reference counting to kick in.
+    """
+
+    # The size_bytes must be known in the metadata, num_rows is optional.
+    blocks: Tuple[Tuple[ObjectRef[Block], BlockMetadata], ...]
+
+    # The schema of the blocks in this bundle. This is optional, and may be None
+    # if blocks are empty.
+    schema: Optional["Schema"]
+
+    # Whether we own the blocks (can safely destroy them).
+    owns_blocks: bool
+
+    # The slices of the blocks in this bundle. After __post_init__, this is always
+    # a list with length equal to len(blocks). Individual entries can be None to
+    # represent a full block (equivalent to BlockSlice(0, num_rows)).
+    # Pass None during construction to initialize all slices as None (full blocks).
+    slices: Optional[List[Optional[BlockSlice]]] = None
+
+    # This attribute is used by the split() operator to assign bundles to logical
+    # output splits. It is otherwise None.
+    output_split_idx: Optional[int] = None
+
+    # Object metadata (size, locations, spilling status)
+    _cached_object_meta: Optional[Dict[ObjectRef, "_ObjectMetadata"]] = None
+
+    # Preferred locations for this bundle determined based on the locations
+    # of individual objects and their corresponding size, ie location with the
+    # largest total number of bytes present there has the highest preference.
+    _cached_preferred_locations: Optional[Dict[NodeIdStr, int]] = None
+
+    def __post_init__(self):
+        if not isinstance(self.blocks, tuple):
+            object.__setattr__(self, "blocks", tuple(self.blocks))
+
+        if self.slices is None:
+            self.slices = [None] * len(self.blocks)
+        else:
+            assert len(self.blocks) == len(
+                self.slices
+            ), "Number of blocks and slices must match"
+            # Validate slice ranges
+            for (_, metadata), block_slice in zip(self.blocks, self.slices):
+                if block_slice is not None:
+                    assert (
+                        block_slice.start_offset >= 0
+                    ), f"Slice start_offset must be non-negative: {block_slice.start_offset}"
+                    assert (
+                        block_slice.end_offset >= block_slice.start_offset
+                    ), f"Slice end_offset must be >= start_offset: [{block_slice.start_offset}, {block_slice.end_offset})"
+                    if metadata.num_rows is not None:
+                        assert (
+                            block_slice.end_offset <= metadata.num_rows
+                        ), f"Slice range [{block_slice.start_offset}, {block_slice.end_offset}) exceeds block num_rows: {metadata.num_rows}"
+
+        for b in self.blocks:
+            assert isinstance(b, tuple), b
+            assert len(b) == 2, b
+            assert isinstance(b[0], ray.ObjectRef), b[0]
+            assert isinstance(b[1], BlockMetadata), b[1]
+            if b[1].size_bytes is None:
+                raise ValueError(
+                    "The size in bytes of the block must be known: {}".format(b)
+                )
+
+    def __setattr__(self, key, value):
+        if hasattr(self, key) and key in ["blocks", "owns_blocks"]:
+            raise ValueError(f"The `{key}` field of RefBundle cannot be updated.")
+        object.__setattr__(self, key, value)
+
+    @property
+    def block_refs(self) -> List[ObjectRef[Block]]:
+        """List of block references in this bundle."""
+        return [block_ref for block_ref, _ in self.blocks]
+
+    @property
+    def metadata(self) -> List[BlockMetadata]:
+        """List of block metadata in this bundle."""
+        return [metadata for _, metadata in self.blocks]
+
+    def num_rows(self) -> Optional[int]:
+        """Number of rows present in this bundle, if known.
+
+        Iterates through blocks and their corresponding slices to calculate the total.
+        Note: Block metadata always refers to the full block, not the slice.
+
+        - If block_slice is None, uses the full block's metadata.num_rows
+        - If block_slice is present, uses the slice's num_rows (partial block portion)
+        - Returns None if any full block has unknown row count (metadata.num_rows is None)
+        """
+        total = 0
+        for metadata, block_slice in zip(self.metadata, self.slices):
+            if block_slice is None:
+                if metadata.num_rows is None:
+                    return None
+                total += metadata.num_rows
+            else:
+                total += block_slice.num_rows
+        return total
+
+    def size_bytes(self) -> int:
+        """Size of the blocks of this bundle in bytes.
+
+        Iterates through blocks and their corresponding slices to calculate the total size.
+        Note: Block metadata always refers to the full block, not the slice.
+
+        - If block_slice is None, uses the full block's metadata.size_bytes
+        - If block_slice is present but num_rows is unknown or zero, uses full metadata.size_bytes
+        - If block_slice represents a partial block, estimates size proportionally based on
+          (metadata.size_bytes / metadata.num_rows) * block_slice.num_rows
+        - Otherwise, uses the full metadata.size_bytes
+        """
+        total = 0
+        for (_, metadata), block_slice in zip(self.blocks, self.slices):
+            if block_slice is None:
+                # Full block
+                total += metadata.size_bytes
+            elif metadata.num_rows is None or metadata.num_rows == 0:
+                # Unknown num_rows or empty block - use full metadata size
+                total += metadata.size_bytes
+            elif metadata.num_rows != block_slice.num_rows:
+                # Partial block - estimate size based on rows
+                per_row = metadata.size_bytes / metadata.num_rows
+                total += max(1, int(math.ceil(per_row * block_slice.num_rows)))
+            else:
+                total += metadata.size_bytes
+        return total
+
+    def destroy_if_owned(self) -> int:
+        """Clears the object store memory for these blocks if owned.
+
+        Returns:
+            The number of bytes freed.
+        """
+        should_free = self.owns_blocks and DataContext.get_current().eager_free
+        for block_ref in self.block_refs:
+            trace_deallocation(
+                block_ref, "RefBundle.destroy_if_owned", free=should_free
+            )
+        return self.size_bytes() if should_free else 0
+
+    def get_preferred_object_locations(self) -> Dict[NodeIdStr, int]:
+        """Returns a mapping of node IDs to total bytes stored on each node.
+
+        Returns:
+            Dict mapping node ID to total bytes stored on that node
+        """
+        meta = self._get_cached_metadata()
+
+        if self._cached_preferred_locations is None:
+            preferred_locs: Dict[NodeIdStr, int] = defaultdict(int)
+
+            for ref, obj_meta in meta.items():
+                for loc in obj_meta.locs:
+                    preferred_locs[loc] += obj_meta.size
+
+            self._cached_preferred_locations = preferred_locs
+
+        return self._cached_preferred_locations
+
+    def _get_cached_metadata(self) -> Dict[ObjectRef, "_ObjectMetadata"]:
+        if self._cached_object_meta is None:
+            # This call is pretty fast for owned objects (~5k/s), so we don't need to
+            # batch it for now.
+            meta = ray.experimental.get_local_object_locations(self.block_refs)
+            # Extract locations
+            object_metas: Dict[ObjectRef, _ObjectMetadata] = {
+                ref: _ObjectMetadata(
+                    size=meta[ref]["object_size"],
+                    spilled=meta[ref]["did_spill"],
+                    locs=meta[ref]["node_ids"],
+                )
+                for ref in self.block_refs
+            }
+
+            self._cached_object_meta = object_metas
+
+        return self._cached_object_meta
+
+    def slice(self, needed_rows: int) -> Tuple["RefBundle", "RefBundle"]:
+        """Slice a Ref Bundle into the first bundle containing the first `needed_rows` rows and the remaining bundle containing the remaining rows.
+
+        Args:
+            needed_rows: Number of rows to take from the head of the bundle.
+
+        Returns:
+            A tuple of (sliced_bundle, remaining_bundle). The needed rows must be less than the number of rows in the bundle.
+        """
+        assert needed_rows > 0, "needed_rows must be positive."
+        assert (
+            self.num_rows() is not None
+        ), "Cannot slice a RefBundle with unknown number of rows."
+        assert (
+            needed_rows < self.num_rows()
+        ), f"To slice a RefBundle, the number of requested rows must be less than the number of rows in the bundle. Requested {needed_rows} rows but bundle only has {self.num_rows()} rows."
+
+        block_slices = []
+        for metadata, block_slice in zip(self.metadata, self.slices):
+            if block_slice is None:
+                # None represents a full block, convert to explicit BlockSlice
+                assert (
+                    metadata.num_rows is not None
+                ), "Cannot derive block slice for a RefBundle with unknown block row counts."
+                block_slices.append(
+                    BlockSlice(start_offset=0, end_offset=metadata.num_rows)
+                )
+            else:
+                block_slices.append(block_slice)
+
+        consumed_blocks: List[Tuple[ObjectRef[Block], BlockMetadata]] = []
+        consumed_slices: List[BlockSlice] = []
+        remaining_blocks: List[Tuple[ObjectRef[Block], BlockMetadata]] = []
+        remaining_slices: List[BlockSlice] = []
+
+        rows_to_take = needed_rows
+
+        for (block_ref, metadata), block_slice in zip(self.blocks, block_slices):
+            block_rows = block_slice.num_rows
+            if rows_to_take >= block_rows:
+                consumed_blocks.append((block_ref, metadata))
+                consumed_slices.append(block_slice)
+                rows_to_take -= block_rows
+            else:
+                if rows_to_take == 0:
+                    remaining_blocks.append((block_ref, metadata))
+                    remaining_slices.append(block_slice)
+                    continue
+                consume_slice = BlockSlice(
+                    start_offset=block_slice.start_offset,
+                    end_offset=block_slice.start_offset + rows_to_take,
+                )
+                consumed_blocks.append((block_ref, metadata))
+                consumed_slices.append(consume_slice)
+
+                leftover_rows = block_rows - rows_to_take
+                if leftover_rows > 0:
+                    remainder_slice = BlockSlice(
+                        start_offset=consume_slice.end_offset,
+                        end_offset=block_slice.end_offset,
+                    )
+                    remaining_blocks.append((block_ref, metadata))
+                    remaining_slices.append(remainder_slice)
+
+                rows_to_take = 0
+
+        sliced_bundle = RefBundle(
+            blocks=tuple(consumed_blocks),
+            schema=self.schema,
+            owns_blocks=False,
+            slices=consumed_slices if consumed_slices else None,
+        )
+
+        remaining_bundle = RefBundle(
+            blocks=tuple(remaining_blocks),
+            schema=self.schema,
+            owns_blocks=False,
+            slices=remaining_slices if remaining_slices else None,
+        )
+
+        return sliced_bundle, remaining_bundle
+
+    @classmethod
+    def merge_ref_bundles(cls, bundles: List["RefBundle"]) -> "RefBundle":
+        assert bundles, "Cannot merge an empty list of RefBundles."
+        merged_blocks = list(itertools.chain(*[bundle.blocks for bundle in bundles]))
+        merged_slices = list(itertools.chain(*[bundle.slices for bundle in bundles]))
+        return cls(
+            blocks=tuple(merged_blocks),
+            schema=bundles[0].schema,  # Assume all bundles have the same schema
+            owns_blocks=bundles[
+                0
+            ].owns_blocks,  # Assume all bundles have the same ownership
+            slices=merged_slices,
+        )
+
+    def __eq__(self, other) -> bool:
+        return self is other
+
+    def __hash__(self) -> int:
+        return id(self)
+
+    def __len__(self) -> int:
+        return len(self.blocks)
+
+    def __str__(self) -> str:
+        lines = [
+            f"RefBundle({len(self.blocks)} blocks,",
+            f"  {self.num_rows()} rows,",
+            f"  schema={self.schema},",
+            f"  owns_blocks={self.owns_blocks},",
+            "  blocks=(",
+        ]
+
+        # Loop through each block and show details
+        for i, ((block_ref, metadata), block_slice) in enumerate(
+            zip(self.blocks, self.slices)
+        ):
+            row_str = (
+                f"{metadata.num_rows} rows"
+                if metadata.num_rows is not None
+                else "unknown rows"
+            )
+            bytes_str = f"{metadata.size_bytes} bytes"
+            slice_str = (
+                f"slice={block_slice}"
+                if block_slice is not None
+                else "slice=None (full block)"
+            )
+
+            lines.append(f"    {i}: {row_str}, {bytes_str}, {slice_str}")
+
+        lines.append("  )")
+        lines.append(")")
+
+        return "\n".join(lines)
+
+
+@dataclass
+class _ObjectMetadata:
+    # Object size in bytes
+    size: int
+    # Flag whether object has been spilled
+    spilled: bool
+    # List of nodes object exists on
+    locs: List[NodeIdStr] = None
+
+
+def _ref_bundles_iterator_to_block_refs_list(
+    ref_bundles: Iterator[RefBundle],
+) -> List[ObjectRef[Block]]:
+    """Convert an iterator of RefBundles to a list of Block object references."""
+    return [
+        block_ref for ref_bundle in ref_bundles for block_ref in ref_bundle.block_refs
+    ]
+
+
+def _iter_sliced_blocks(
+    blocks: Iterable[Block],
+    slices: List[Optional[BlockSlice]],
+) -> Iterator[Block]:
+    blocks_list = list(blocks)
+    for block, block_slice in zip(blocks_list, slices):
+        if block_slice is None:
+            # None represents a full block - yield it as is
+            yield block
+        else:
+            accessor = BlockAccessor.for_block(block)
+            start = block_slice.start_offset
+            end = block_slice.end_offset
+            assert start <= end, "start must be less than end"
+            assert start >= 0, "start must be non-negative"
+            assert (
+                end <= accessor.num_rows()
+            ), "end must be less than or equal to the number of rows in the block"
+
+            yield accessor.slice(start, end, copy=False)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/task_context.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/task_context.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ff0f60f96708482dcb2d9eade926b5e5fb33947
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/task_context.py
@@ -0,0 +1,76 @@
+import threading
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+from ray.data._internal.progress_bar import ProgressBar
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.operators.map_transformer import MapTransformer
+
+
+_thread_local = threading.local()
+
+
+@dataclass
+class TaskContext:
+    """This describes the information of a task running block transform."""
+
+    # The index of task. Each task has a unique task index within the same
+    # operator.
+    task_idx: int
+
+    # Name of the operator that this task belongs to.
+    op_name: str
+
+    # The dictionary of sub progress bar to update. The key is name of sub progress
+    # bar. Note this is only used on driver side.
+    # TODO(chengsu): clean it up from TaskContext with new optimizer framework.
+    sub_progress_bar_dict: Optional[Dict[str, ProgressBar]] = None
+
+    # NOTE(hchen): `upstream_map_transformer` and `upstream_map_ray_remote_args`
+    # are only used for `RandomShuffle`. DO NOT use them for other operators.
+    # Ideally, they should be handled by the optimizer, and should be transparent
+    # to the specific operators.
+    # But for `RandomShuffle`, the AllToAllOperator doesn't do the shuffle itself.
+    # It uses `ExchangeTaskScheduler` to launch new tasks to do the shuffle.
+    # That's why we need to pass them to `ExchangeTaskScheduler`.
+    # TODO(hchen): Use a physical operator to do the shuffle directly.
+
+    # The underlying function called in a MapOperator; this is used when fusing
+    # an AllToAllOperator with an upstream MapOperator.
+    upstream_map_transformer: Optional["MapTransformer"] = None
+
+    # The Ray remote arguments of the fused upstream MapOperator.
+    # This should be set if upstream_map_transformer is set.
+    upstream_map_ray_remote_args: Optional[Dict[str, Any]] = None
+
+    # Override of the target max-block-size for the task
+    target_max_block_size_override: Optional[int] = None
+
+    # Additional keyword arguments passed to the task.
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def get_current(cls) -> Optional["TaskContext"]:
+        """Get the TaskContext for the current thread.
+        Returns None if no TaskContext has been set.
+        """
+
+        return getattr(_thread_local, "task_context", None)
+
+    @classmethod
+    def set_current(cls, context):
+        """Set the TaskContext for the current thread.
+
+        Args:
+            context: The TaskContext instance to set for this thread
+        """
+
+        _thread_local.task_context = context
+
+    @classmethod
+    def reset_current(cls):
+        """Clear the current thread's TaskContext."""
+
+        if hasattr(_thread_local, "task_context"):
+            delattr(_thread_local, "task_context")
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/transform_fn.py b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/transform_fn.py
new file mode 100644
index 0000000000000000000000000000000000000000..b867c28a291db9f2afc52f343ae79e2aca7d2281
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/interfaces/transform_fn.py
@@ -0,0 +1,14 @@
+from typing import Callable, List, Tuple
+
+from .ref_bundle import RefBundle
+from .task_context import TaskContext
+from ray.data._internal.stats import StatsDict
+
+# Result type of AllToAllTransformFn.
+AllToAllTransformFnResult = Tuple[List[RefBundle], StatsDict]
+
+# Block transform function applied in AllToAllOperator.
+AllToAllTransformFn = Callable[
+    [List[RefBundle], TaskContext],
+    AllToAllTransformFnResult,
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/legacy_compat.py b/lib/python3.12/site-packages/ray/data/_internal/execution/legacy_compat.py
new file mode 100644
index 0000000000000000000000000000000000000000..fefb1592c7366bfd9fae8c12f24c6096930a63cd
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/legacy_compat.py
@@ -0,0 +1,193 @@
+"""This file contains temporary helper functions for legacy plan/executor interaction.
+
+It should be deleted once we fully move to the new executor backend.
+"""
+import logging
+from typing import Iterator, Optional, Tuple
+
+from ray.data._internal.block_list import BlockList
+from ray.data._internal.execution.interfaces import (
+    Executor,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.interfaces.executor import OutputIterator
+from ray.data._internal.execution.streaming_executor_state import Topology
+from ray.data._internal.logical.util import record_operators_usage
+from ray.data._internal.plan import ExecutionPlan
+from ray.data._internal.stats import DatasetStats
+from ray.data.block import (
+    BlockMetadata,
+    BlockMetadataWithSchema,
+    _take_first_non_empty_schema,
+)
+
+# Warn about tasks larger than this.
+TASK_SIZE_WARN_THRESHOLD_BYTES = 100000
+
+logger = logging.getLogger(__name__)
+
+
+def execute_to_legacy_bundle_iterator(
+    executor: Executor,
+    plan: ExecutionPlan,
+) -> Iterator[RefBundle]:
+    """Execute a plan with the new executor and return a bundle iterator.
+
+    Args:
+        executor: The executor to use.
+        plan: The legacy plan to execute.
+
+    Returns:
+        The output as a bundle iterator.
+    """
+    dag, stats = _get_execution_dag(
+        executor,
+        plan,
+        preserve_order=False,
+    )
+
+    bundle_iter = executor.execute(dag, initial_stats=stats)
+
+    topology: "Topology" = executor._topology
+
+    class CacheMetadataIterator(OutputIterator):
+        """Wrapper for `bundle_iterator` above.
+
+        For a given iterator which yields output RefBundles,
+        collect the metadata from each output bundle, and yield the
+        original RefBundle. Only after the entire iterator is exhausted,
+        we cache the resulting metadata to the execution plan."""
+
+        def __init__(self, base_iterator: OutputIterator):
+            # Note: the base_iterator should be of type StreamIterator,
+            # defined within `StreamingExecutor.execute()`. It must
+            # support the `get_next()` method.
+            self._base_iterator = base_iterator
+            self._collected_metadata = BlockMetadata(
+                num_rows=0,
+                size_bytes=0,
+                input_files=None,
+                exec_stats=None,
+            )
+
+        def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
+            try:
+                bundle = self._base_iterator.get_next(output_split_idx)
+                self._collect_metadata(bundle)
+                return bundle
+            except StopIteration:
+                # Once the iterator is completely exhausted, we are done
+                # collecting metadata. We can add this cached metadata to the plan.
+
+                # Traverse the topology backwards and find the first available schema
+                schema = next(reversed(topology.values()))._schema
+
+                meta_with_schema = BlockMetadataWithSchema(
+                    metadata=self._collected_metadata,
+                    schema=schema,
+                )
+                plan._snapshot_metadata_schema = meta_with_schema
+                raise
+
+        def _collect_metadata(self, bundle: RefBundle) -> RefBundle:
+            """Collect the metadata from each output bundle and accumulate
+            results, so we can access important information, such as
+            row count, schema, etc., after iteration completes."""
+            self._collected_metadata.num_rows += bundle.num_rows()
+            self._collected_metadata.size_bytes += bundle.size_bytes()
+            return bundle
+
+    return CacheMetadataIterator(bundle_iter)
+
+
+def execute_to_legacy_block_list(
+    executor: Executor,
+    plan: ExecutionPlan,
+    dataset_uuid: str,
+    preserve_order: bool,
+) -> BlockList:
+    """Execute a plan with the new executor and translate it into a legacy block list.
+
+    Args:
+        executor: The executor to use.
+        plan: The legacy plan to execute.
+        dataset_uuid: UUID of the dataset for this execution.
+        preserve_order: Whether to preserve order in execution.
+
+    Returns:
+        The output as a legacy block list.
+    """
+    dag, stats = _get_execution_dag(
+        executor,
+        plan,
+        preserve_order,
+    )
+    bundles = executor.execute(dag, initial_stats=stats)
+    block_list = _bundles_to_block_list(bundles)
+    # Set the stats UUID after execution finishes.
+    _set_stats_uuid_recursive(executor.get_stats(), dataset_uuid)
+    return block_list
+
+
+def _get_execution_dag(
+    executor: Executor,
+    plan: ExecutionPlan,
+    preserve_order: bool,
+) -> Tuple[PhysicalOperator, DatasetStats]:
+    """Get the physical operators DAG from a plan."""
+    from ray.data._internal.logical.optimizers import get_execution_plan
+
+    # Record usage of logical operators if available.
+    if hasattr(plan, "_logical_plan") and plan._logical_plan is not None:
+        record_operators_usage(plan._logical_plan.dag)
+
+    # Get DAG of physical operators and input statistics.
+    dag = get_execution_plan(plan._logical_plan).dag
+    stats = _get_initial_stats_from_plan(plan)
+
+    # Enforce to preserve ordering if the plan has operators
+    # required to do so, such as Zip and Sort.
+    if preserve_order or plan.require_preserve_order():
+        executor._options.preserve_order = True
+
+    return dag, stats
+
+
+def _get_initial_stats_from_plan(plan: ExecutionPlan) -> DatasetStats:
+    if plan._snapshot_bundle is not None:
+        return plan._snapshot_stats
+    # For Datasets created from "read_xxx", `plan._in_stats` contains useless data.
+    # For Datasets created from "from_xxx", we need to use `plan._in_stats` as
+    # the initial stats. Because the `FromXxx` logical operators will be translated to
+    # "InputDataBuffer" physical operators, which will be ignored when generating
+    # stats, see `StreamingExecutor._generate_stats`.
+    # TODO(hchen): Unify the logic by saving the initial stats in `InputDataBuffer
+    if plan.has_lazy_input():
+        return DatasetStats(metadata={}, parent=None)
+    else:
+        return plan._in_stats
+
+
+def _bundles_to_block_list(bundles: Iterator[RefBundle]) -> BlockList:
+    blocks, metadata = [], []
+    owns_blocks = True
+    bundle_list = list(bundles)
+    schema = _take_first_non_empty_schema(
+        ref_bundle.schema for ref_bundle in bundle_list
+    )
+
+    for ref_bundle in bundle_list:
+        if not ref_bundle.owns_blocks:
+            owns_blocks = False
+        blocks.extend(ref_bundle.block_refs)
+        metadata.extend(ref_bundle.metadata)
+
+    return BlockList(blocks, metadata, owned_by_consumer=owns_blocks, schema=schema)
+
+
+def _set_stats_uuid_recursive(stats: DatasetStats, dataset_uuid: str) -> None:
+    if not stats.dataset_uuid:
+        stats.dataset_uuid = dataset_uuid
+    for parent in stats.parents or []:
+        _set_stats_uuid_recursive(parent, dataset_uuid)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d336807561288d80a279d9c31713a69486d0f8b7
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__init__.py
@@ -0,0 +1,3 @@
+from .actor_location import ActorLocationTracker, get_or_create_actor_location_tracker
+
+__all__ = ["get_or_create_actor_location_tracker", "ActorLocationTracker"]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7603ffbbc87ff781907f7b91989df5ec2124f189
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__pycache__/actor_location.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__pycache__/actor_location.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2f979669dfc13fa27b15a42ab6c2bb92964f429a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/__pycache__/actor_location.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/actor_location.py b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/actor_location.py
new file mode 100644
index 0000000000000000000000000000000000000000..6acab96fc58e5403665fbcba8e52bd020f111e91
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/node_trackers/actor_location.py
@@ -0,0 +1,41 @@
+import threading
+from typing import List
+
+import ray
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+
+@ray.remote(num_cpus=0, max_restarts=-1, max_task_retries=-1)
+class ActorLocationTracker:
+    def __init__(self):
+        self._actor_locations = {}
+        self._actor_locations_lock = threading.Lock()
+
+    def update_actor_location(self, logical_actor_id: str, node_id: str):
+        with self._actor_locations_lock:
+            self._actor_locations[logical_actor_id] = node_id
+
+    def get_actor_locations(self, logical_actor_ids: List[str]):
+        return {
+            logical_actor_id: self._actor_locations.get(logical_actor_id, None)
+            for logical_actor_id in logical_actor_ids
+        }
+
+
+def get_or_create_actor_location_tracker():
+
+    # Pin the actor location tracker to the local node so it fate-shares with the driver.
+    # NOTE: for Ray Client, the ray.get_runtime_context().get_node_id() should
+    # point to the head node.
+    scheduling_strategy = NodeAffinitySchedulingStrategy(
+        ray.get_runtime_context().get_node_id(),
+        soft=False,
+    )
+    return ActorLocationTracker.options(
+        name="ActorLocationTracker",
+        namespace="ActorLocationTracker",
+        get_if_exists=True,
+        lifetime="detached",
+        scheduling_strategy=scheduling_strategy,
+        max_concurrency=8,
+    ).remote()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5b1098a9fe5e996c3771bd15c0aacd5d414d62fb
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/actor_pool_map_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/actor_pool_map_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1dbd51dabdd39c56ae15ff11fc854413d5a788e7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/actor_pool_map_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/aggregate_num_rows.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/aggregate_num_rows.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e8b4323d903792f5df3e45dfe78aa3e87d7ab08d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/aggregate_num_rows.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/base_physical_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/base_physical_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..de6d97c212e06c933d59cba65a73beaae1c1ea57
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/base_physical_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/hash_aggregate.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/hash_aggregate.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..81768594efd5a792a0499154e41a37a1807d4116
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/hash_aggregate.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/hash_shuffle.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/hash_shuffle.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f1d0e1714c0e7f246efd5f11bee94c2ab53979be
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/hash_shuffle.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/input_data_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/input_data_buffer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b6809f94aefc1ce3d6dbde335131f28d268bc0ad
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/input_data_buffer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/join.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/join.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ee7e253d58ee056ec2e486b4fa517239befdd7b4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/join.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/limit_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/limit_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..20f44db775f153efbe3d99312fbacffcbc537070
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/limit_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/map_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/map_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f24632ee1c71fc1bf9d55bc567c78e81fe092c15
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/map_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/map_transformer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/map_transformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d71f00e9ee487f538696733bded6edec785f255c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/map_transformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/output_splitter.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/output_splitter.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..61ccd8070d8851dd32477b1c24ada67d75b6132e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/output_splitter.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/sub_progress.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/sub_progress.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..223885d525d2abbfd17dae42c9161840aace2bac
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/sub_progress.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/task_pool_map_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/task_pool_map_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c4e04f793d3fa3240e40097277b80ce15a5936d4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/task_pool_map_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/union_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/union_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..654c5fc761c562e5ce338c6e2e12863b2600a472
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/union_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/zip_operator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/zip_operator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..deeaf46e4c39df888af3ef51fd5205b74f97808b
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/__pycache__/zip_operator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/actor_pool_map_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/actor_pool_map_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..017a4e9a3192b8d7e805f364fb8c281a47c4efe9
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/actor_pool_map_operator.py
@@ -0,0 +1,1184 @@
+import abc
+import logging
+import time
+import uuid
+import warnings
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
+
+import ray
+from ray.actor import ActorHandle
+from ray.core.generated import gcs_pb2
+from ray.data._internal.actor_autoscaler import (
+    AutoscalingActorPool,
+)
+from ray.data._internal.actor_autoscaler.autoscaling_actor_pool import (
+    ActorPoolScalingRequest,
+)
+from ray.data._internal.compute import ActorPoolStrategy
+from ray.data._internal.execution.bundle_queue import create_bundle_queue
+from ray.data._internal.execution.bundle_queue.bundle_queue import BundleQueue
+from ray.data._internal.execution.interfaces import (
+    BlockSlice,
+    ExecutionOptions,
+    ExecutionResources,
+    NodeIdStr,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.physical_operator import _ActorPoolInfo
+from ray.data._internal.execution.node_trackers.actor_location import (
+    ActorLocationTracker,
+    get_or_create_actor_location_tracker,
+)
+from ray.data._internal.execution.operators.map_operator import (
+    BaseRefBundler,
+    MapOperator,
+    _map_task,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.execution.util import locality_string
+from ray.data._internal.remote_fn import _add_system_error_to_retry_exceptions
+from ray.data.block import Block, BlockMetadata
+from ray.data.context import (
+    DEFAULT_ACTOR_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR,
+    DataContext,
+)
+from ray.types import ObjectRef
+from ray.util.common import INT32_MAX
+
+logger = logging.getLogger(__name__)
+
+
+class ActorPoolMapOperator(MapOperator):
+    """A MapOperator implementation that executes tasks on an actor pool.
+
+    NOTE: This class is NOT thread-safe
+
+    This class manages the state of a pool of actors used for task execution, as well
+    as dispatch of tasks to those actors.
+
+    It operates in two modes. In bulk mode, tasks are queued internally and executed
+    when the operator has free actor slots. In streaming mode, the streaming executor
+    only adds input when `should_add_input() = True` (i.e., there are free slots).
+    This allows for better control of backpressure (e.g., suppose we go over memory
+    limits after adding put, then there isn't any way to "take back" the inputs prior
+    to actual execution).
+    """
+
+    def __init__(
+        self,
+        map_transformer: MapTransformer,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        compute_strategy: ActorPoolStrategy,
+        name: str = "ActorPoolMap",
+        min_rows_per_bundle: Optional[int] = None,
+        ref_bundler: Optional[BaseRefBundler] = None,
+        supports_fusion: bool = True,
+        map_task_kwargs: Optional[Dict[str, Any]] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+        ray_actor_task_remote_args: Optional[Dict[str, Any]] = None,
+        target_max_block_size_override: Optional[int] = None,
+    ):
+        """Create an ActorPoolMapOperator instance.
+
+        Args:
+            map_transformer: Instance of `MapTransformer` that will be applied
+                to each ref bundle input.
+            input_op: Operator generating input data for this op.
+            data_context: The DataContext instance containing configuration settings.
+            compute_strategy: `ComputeStrategy` used for this operator.
+            name: The name of this operator.
+            min_rows_per_bundle: The number of rows to gather per batch passed to the
+                transform_fn, or None to use the block size. Setting the batch size is
+                important for the performance of GPU-accelerated transform functions.
+                The actual rows passed may be less if the dataset is small.
+            ref_bundler: The ref bundler to use for this operator.
+            supports_fusion: Whether this operator supports fusion with other operators.
+            map_task_kwargs: A dictionary of kwargs to pass to the map task. You can
+                access these kwargs through the `TaskContext.kwargs` dictionary.
+            ray_remote_args_fn: A function that returns a dictionary of remote args
+                passed to each map worker. The purpose of this argument is to generate
+                dynamic arguments for each actor/task, and will be called each time
+                prior to initializing the worker. Args returned from this dict will
+                always override the args in ``ray_remote_args``. Note: this is an
+                advanced, experimental feature.
+            ray_remote_args: Customize the ray remote args for this op's tasks.
+                See :func:`ray.remote` for details.
+            ray_actor_task_remote_args: Ray Core options passed to map actor tasks.
+            target_max_block_size_override: The target maximum number of bytes to
+                include in an output block.
+        """
+        super().__init__(
+            map_transformer,
+            input_op,
+            data_context,
+            name,
+            target_max_block_size_override,
+            min_rows_per_bundle,
+            ref_bundler,
+            supports_fusion,
+            map_task_kwargs,
+            ray_remote_args_fn,
+            ray_remote_args,
+        )
+
+        self._min_rows_per_bundle = min_rows_per_bundle
+        self._ray_remote_args_fn = ray_remote_args_fn
+        self._ray_remote_args = self._apply_default_remote_args(
+            self._ray_remote_args, self.data_context
+        )
+        self._ray_actor_task_remote_args = self._apply_default_actor_task_remote_args(
+            ray_actor_task_remote_args, self.data_context
+        )
+
+        per_actor_resource_usage = ExecutionResources(
+            cpu=self._ray_remote_args.get("num_cpus"),
+            gpu=self._ray_remote_args.get("num_gpus"),
+            memory=self._ray_remote_args.get("memory"),
+        )
+
+        max_actor_concurrency = self._ray_remote_args.get("max_concurrency", 1)
+
+        self._actor_pool = _ActorPool(
+            self._start_actor,
+            per_actor_resource_usage,
+            min_size=compute_strategy.min_size,
+            max_size=compute_strategy.max_size,
+            initial_size=compute_strategy.initial_size,
+            max_actor_concurrency=max_actor_concurrency,
+            max_tasks_in_flight_per_actor=(
+                # Unless explicitly overridden by the user, max tasks-in-flight config
+                # will fall back to be:
+                #
+                #   DEFAULT_ACTOR_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR * max_concurrency,
+                compute_strategy.max_tasks_in_flight_per_actor
+                or data_context.max_tasks_in_flight_per_actor
+                or max_actor_concurrency
+                * DEFAULT_ACTOR_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR
+            ),
+            _enable_actor_pool_on_exit_hook=self.data_context._enable_actor_pool_on_exit_hook,
+        )
+        self._actor_task_selector = self._create_task_selector(self._actor_pool)
+        # A queue of bundles awaiting dispatch to actors.
+        self._bundle_queue = create_bundle_queue()
+        # HACK: Without this, all actors show up as `_MapWorker` in Grafana, so we can’t
+        # tell which operator they belong to. To fix that, we dynamically create a new
+        # class per operator with a unique name.
+        self._map_worker_cls = type(f"MapWorker({self.name})", (_MapWorker,), {})
+        # Cached actor class.
+        self._actor_cls = None
+        # Whether no more submittable bundles will be added.
+        self._inputs_done = False
+        self._actor_locality_enabled: Optional[bool] = None
+
+        # Locality metrics
+        self._locality_hits = 0
+        self._locality_misses = 0
+
+    @staticmethod
+    def _create_task_selector(actor_pool: "_ActorPool") -> "_ActorTaskSelector":
+        return _ActorTaskSelectorImpl(actor_pool)
+
+    @staticmethod
+    def _apply_default_actor_task_remote_args(
+        ray_actor_task_remote_args: Optional[Dict[str, Any]], data_context: DataContext
+    ) -> Dict[str, Any]:
+        """Apply defaults to the actor task remote args."""
+        if ray_actor_task_remote_args is None:
+            ray_actor_task_remote_args = {}
+
+        ray_actor_task_remote_args = ray_actor_task_remote_args.copy()
+
+        actor_task_errors = data_context.actor_task_retry_on_errors
+        if actor_task_errors:
+            ray_actor_task_remote_args["retry_exceptions"] = actor_task_errors
+        _add_system_error_to_retry_exceptions(ray_actor_task_remote_args)
+
+        if (
+            "_generator_backpressure_num_objects" not in ray_actor_task_remote_args
+            and data_context._max_num_blocks_in_streaming_gen_buffer is not None
+        ):
+            # The `_generator_backpressure_num_objects` parameter should be
+            # `2 * _max_num_blocks_in_streaming_gen_buffer` because we yield
+            # 2 objects for each block: the block and the block metadata.
+            ray_actor_task_remote_args["_generator_backpressure_num_objects"] = (
+                2 * data_context._max_num_blocks_in_streaming_gen_buffer
+            )
+
+        return ray_actor_task_remote_args
+
+    def internal_input_queue_num_blocks(self) -> int:
+        # NOTE: Internal queue size for ``ActorPoolMapOperator`` includes both
+        #   - Input blocks bundler, alas
+        #   - Own bundle's queue
+        return self._block_ref_bundler.num_blocks() + self._bundle_queue.num_blocks()
+
+    def internal_input_queue_num_bytes(self) -> int:
+        return (
+            self._bundle_queue.estimate_size_bytes()
+            + self._block_ref_bundler.size_bytes()
+        )
+
+    def start(self, options: ExecutionOptions):
+        self._actor_locality_enabled = options.actor_locality_enabled
+        super().start(options)
+
+        # Create the actor workers and add them to the pool.
+        self._actor_cls = ray.remote(**self._ray_remote_args)(self._map_worker_cls)
+        self._actor_pool.scale(
+            ActorPoolScalingRequest(
+                delta=self._actor_pool.initial_size(), reason="scaling to initial size"
+            )
+        )
+
+        # If `wait_for_min_actors_s` is specified and is positive, then
+        # Actor Pool will block until min number of actors is provisioned.
+        #
+        # Otherwise, all actors will be provisioned asynchronously.
+        if self.data_context.wait_for_min_actors_s > 0:
+            refs = self._actor_pool.get_pending_actor_refs()
+
+            logger.debug(
+                f"{self._name}: Waiting for {len(refs)} pool actors to start "
+                f"(for {self.data_context.wait_for_min_actors_s}s)..."
+            )
+
+            try:
+                timeout = self.data_context.wait_for_min_actors_s
+                ray.get(refs, timeout=timeout)
+            except ray.exceptions.GetTimeoutError:
+                raise ray.exceptions.GetTimeoutError(
+                    "Timed out while starting actors. "
+                    "This may mean that the cluster does not have "
+                    "enough resources for the requested actor pool."
+                )
+
+    def should_add_input(self) -> bool:
+        return self._actor_pool.num_free_task_slots() > 0
+
+    def _start_actor(
+        self, labels: Dict[str, str], logical_actor_id: str
+    ) -> Tuple[ActorHandle, ObjectRef]:
+        """Start a new actor and add it to the actor pool as a pending actor.
+
+        Args:
+            labels: The key-value labels to launch the actor with.
+            logical_actor_id: The logical id of the actor.
+
+        Returns:
+            A tuple of the actor handle and the object ref to the actor's location.
+        """
+        assert self._actor_cls is not None
+        ctx = self.data_context
+        if self._ray_remote_args_fn:
+            self._refresh_actor_cls()
+        actor = self._actor_cls.options(
+            _labels={self._OPERATOR_ID_LABEL_KEY: self.id, **labels}
+        ).remote(
+            ctx=ctx,
+            logical_actor_id=logical_actor_id,
+            src_fn_name=self.name,
+            map_transformer=self._map_transformer,
+            actor_location_tracker=get_or_create_actor_location_tracker(),
+        )
+        res_ref = actor.get_location.remote()
+
+        def _task_done_callback(res_ref):
+            # res_ref is a future for a now-ready actor; move actor from pending to the
+            # active actor pool.
+            has_actor = self._actor_pool.pending_to_running(res_ref)
+            if not has_actor:
+                # Actor has already been killed.
+                return
+            # A new actor has started, we try to dispatch queued tasks.
+            self._dispatch_tasks()
+
+        self._submit_metadata_task(
+            res_ref,
+            lambda: _task_done_callback(res_ref),
+        )
+        return actor, res_ref
+
+    def _add_bundled_input(self, bundle: RefBundle):
+        self._bundle_queue.add(bundle)
+        self._metrics.on_input_queued(bundle)
+        # Try to dispatch all bundles in the queue, including this new bundle.
+        self._dispatch_tasks()
+
+    def _dispatch_tasks(self):
+        """Try to dispatch tasks from the bundle buffer to the actor pool.
+
+        This is called when:
+            * a new input bundle is added,
+            * a task finishes,
+            * a new worker has been created.
+        """
+        for bundle, actor in self._actor_task_selector.select_actors(
+            self._bundle_queue, self._actor_locality_enabled
+        ):
+            # Submit the map task.
+            self._metrics.on_input_dequeued(bundle)
+            input_blocks = [block for block, _ in bundle.blocks]
+            self._actor_pool.on_task_submitted(actor)
+
+            ctx = TaskContext(
+                task_idx=self._next_data_task_idx,
+                op_name=self.name,
+                target_max_block_size_override=self.target_max_block_size_override,
+            )
+            gen = actor.submit.options(
+                num_returns="streaming",
+                **self._ray_actor_task_remote_args,
+            ).remote(
+                self.data_context,
+                ctx,
+                *input_blocks,
+                slices=bundle.slices,
+                **self.get_map_task_kwargs(),
+            )
+
+            def _task_done_callback(actor_to_return):
+                # Return the actor that was running the task to the pool.
+                self._actor_pool.on_task_completed(actor_to_return)
+                # Dipsatch more tasks.
+                self._dispatch_tasks()
+
+            from functools import partial
+
+            self._submit_data_task(
+                gen, bundle, partial(_task_done_callback, actor_to_return=actor)
+            )
+
+            # Update locality metrics
+            if (
+                self._actor_pool.running_actors()[actor].actor_location
+                in bundle.get_preferred_object_locations()
+            ):
+                self._locality_hits += 1
+            else:
+                self._locality_misses += 1
+
+    def _refresh_actor_cls(self):
+        """When `self._ray_remote_args_fn` is specified, this method should
+        be called prior to initializing the new worker in order to get new
+        remote args passed to the worker. It updates `self.cls` with the same
+        `_MapWorker` class, but with the new remote args from
+        `self._ray_remote_args_fn`."""
+        assert self._ray_remote_args_fn, "_ray_remote_args_fn must be provided"
+        remote_args = self._ray_remote_args.copy()
+        new_remote_args = self._ray_remote_args_fn()
+
+        # Override args from user-defined remote args function.
+        new_and_overriden_remote_args = {}
+        for k, v in new_remote_args.items():
+            remote_args[k] = v
+            new_and_overriden_remote_args[k] = v
+        self._actor_cls = ray.remote(**remote_args)(self._map_worker_cls)
+        return new_and_overriden_remote_args
+
+    def all_inputs_done(self):
+        # Call base implementation to handle any leftover bundles. This may or may not
+        # trigger task dispatch.
+        super().all_inputs_done()
+
+        # Mark inputs as done so future task dispatch will kill all inactive workers
+        # once the bundle queue is exhausted.
+        self._inputs_done = True
+
+        if self._metrics.num_inputs_received < self._actor_pool.min_size():
+            warnings.warn(
+                f"The minimum number of concurrent actors for '{self.name}' is set to "
+                f"{self._actor_pool.min_size()}, but the operator only received "
+                f"{self._metrics.num_inputs_received} input(s). This means that the "
+                f"operator can launch at most {self._metrics.num_inputs_received} "
+                f"task(s), and won't fully utilize the available concurrency. "
+                "You might be able to increase the number of concurrent tasks by "
+                "configuring `override_num_blocks` earlier in the pipeline."
+            )
+
+    def _do_shutdown(self, force: bool = False):
+        self._actor_pool.shutdown(force=force)
+        # NOTE: It's critical for Actor Pool to release actors before calling into
+        #       the base method that will attempt to cancel and join pending.
+        super()._do_shutdown(force)
+
+    def progress_str(self) -> str:
+        if self._actor_locality_enabled:
+            return locality_string(
+                self._locality_hits,
+                self._locality_misses,
+            )
+        return "[locality off]"
+
+    def min_max_resource_requirements(
+        self,
+    ) -> Tuple[ExecutionResources, ExecutionResources]:
+        min_actors = self._actor_pool.min_size()
+        assert min_actors is not None, min_actors
+
+        num_cpus_per_actor = self._ray_remote_args.get("num_cpus", 0)
+        num_gpus_per_actor = self._ray_remote_args.get("num_gpus", 0)
+        memory_per_actor = self._ray_remote_args.get("memory", 0)
+
+        min_resource_usage = ExecutionResources(
+            cpu=num_cpus_per_actor * min_actors,
+            gpu=num_gpus_per_actor * min_actors,
+            memory=memory_per_actor * min_actors,
+            # To ensure that all actors are utilized, reserve enough resource budget
+            # to launch one task for each worker.
+            object_store_memory=(
+                self._metrics.obj_store_mem_max_pending_output_per_task or 0
+            )
+            * min_actors,
+        )
+
+        return min_resource_usage, ExecutionResources.for_limits()
+
+    def current_processor_usage(self) -> ExecutionResources:
+        # Both pending and running actors count towards our current resource usage.
+        num_active_workers = self._actor_pool.current_size()
+        return ExecutionResources(
+            cpu=self._ray_remote_args.get("num_cpus", 0) * num_active_workers,
+            gpu=self._ray_remote_args.get("num_gpus", 0) * num_active_workers,
+        )
+
+    def pending_processor_usage(self) -> ExecutionResources:
+        # Both pending and restarting actors count towards pending processor usage
+        num_pending_workers = (
+            self._actor_pool.num_pending_actors()
+            + self._actor_pool.num_restarting_actors()
+        )
+        return ExecutionResources(
+            cpu=self._ray_remote_args.get("num_cpus", 0) * num_pending_workers,
+            gpu=self._ray_remote_args.get("num_gpus", 0) * num_pending_workers,
+        )
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        # Submitting tasks to existing actors doesn't require additional
+        # CPU/GPU resources.
+        return ExecutionResources(
+            cpu=0,
+            gpu=0,
+            object_store_memory=self._metrics.obj_store_mem_max_pending_output_per_task
+            or 0,
+        )
+
+    def _extra_metrics(self) -> Dict[str, Any]:
+        res = {}
+        if self._actor_locality_enabled:
+            res["locality_hits"] = self._locality_hits
+            res["locality_misses"] = self._locality_misses
+        res["pending_actors"] = self._actor_pool.num_pending_actors()
+        res["restarting_actors"] = self._actor_pool.num_restarting_actors()
+        return res
+
+    @staticmethod
+    def _apply_default_remote_args(
+        ray_remote_args: Dict[str, Any], data_context: DataContext
+    ) -> Dict[str, Any]:
+        """Apply defaults to the actor creation remote args."""
+        ray_remote_args = ray_remote_args.copy()
+        if "scheduling_strategy" not in ray_remote_args:
+            ray_remote_args["scheduling_strategy"] = data_context.scheduling_strategy
+        # Enable actor fault tolerance by default, with infinite actor recreations and
+        # up to N retries per task. The user can customize this in map_batches via
+        # extra kwargs (e.g., map_batches(..., max_restarts=0) to disable).
+        if "max_restarts" not in ray_remote_args:
+            ray_remote_args["max_restarts"] = -1
+        if (
+            "max_task_retries" not in ray_remote_args
+            and ray_remote_args.get("max_restarts") != 0
+        ):
+            ray_remote_args["max_task_retries"] = -1
+
+        # Allow actor tasks to execute out of order by default. This prevents actors
+        # from idling when the first actor task is blocked.
+        #
+        # `MapOperator` should still respect `preserve_order` in this case.
+        if "allow_out_of_order_execution" not in ray_remote_args:
+            ray_remote_args["allow_out_of_order_execution"] = True
+
+        return ray_remote_args
+
+    def get_autoscaling_actor_pools(self) -> List[AutoscalingActorPool]:
+        return [self._actor_pool]
+
+    def per_task_resource_allocation(
+        self: "PhysicalOperator",
+    ) -> ExecutionResources:
+        # For Actor tasks resource allocation is determined as:
+        #   - Per actor resource allocation divided by
+        #   - Actor's max task concurrency
+        max_concurrency = self._actor_pool.max_actor_concurrency()
+        per_actor_resource_usage = self._actor_pool.per_actor_resource_usage()
+        return per_actor_resource_usage.scale(1 / max_concurrency)
+
+    def min_scheduling_resources(
+        self: "PhysicalOperator",
+    ) -> ExecutionResources:
+        return self._actor_pool.per_actor_resource_usage()
+
+    def update_resource_usage(self) -> None:
+        """Updates resources usage."""
+        for actor in self._actor_pool.get_running_actor_refs():
+            actor_state = actor._get_local_state()
+            if actor_state in (None, gcs_pb2.ActorTableData.ActorState.DEAD):
+                # actor._get_local_state can return None if the state is Unknown
+                # If actor_state is None or dead, there is nothing to do.
+                continue
+            elif actor_state != gcs_pb2.ActorTableData.ActorState.ALIVE:
+                # The actors can be either ALIVE or RESTARTING here because they will
+                # be restarted indefinitely until execution finishes.
+                assert (
+                    actor_state == gcs_pb2.ActorTableData.ActorState.RESTARTING
+                ), actor_state
+                self._actor_pool.update_running_actor_state(actor, True)
+            else:
+                self._actor_pool.update_running_actor_state(actor, False)
+
+    def get_actor_info(self) -> _ActorPoolInfo:
+        """Returns Actor counts for Alive, Restarting and Pending Actors."""
+        return self._actor_pool.get_actor_info()
+
+    def get_max_concurrency_limit(self) -> Optional[int]:
+        return self._actor_pool.max_size() * self._actor_pool.max_actor_concurrency()
+
+
+class _MapWorker:
+    """An actor worker for MapOperator."""
+
+    def __init__(
+        self,
+        ctx: DataContext,
+        src_fn_name: str,
+        map_transformer: MapTransformer,
+        logical_actor_id: str,
+        actor_location_tracker: ray.actor.ActorHandle[ActorLocationTracker],
+    ):
+        self.src_fn_name: str = src_fn_name
+        self._map_transformer = map_transformer
+        # Initialize the data context for this actor after setting the src_fn_name in order to not
+        # break __repr__. It's possible that logging setup fails.
+        DataContext._set_current(ctx)
+        # Initialize state for this actor.
+        self._map_transformer.init()
+        self._logical_actor_id = logical_actor_id
+        actor_location_tracker.update_actor_location.remote(
+            self._logical_actor_id, ray.get_runtime_context().get_node_id()
+        )
+
+    def get_location(self) -> NodeIdStr:
+        return ray.get_runtime_context().get_node_id()
+
+    def submit(
+        self,
+        data_context: DataContext,
+        ctx: TaskContext,
+        *blocks: Block,
+        slices: Optional[List[BlockSlice]] = None,
+        **kwargs: Dict[str, Any],
+    ) -> Iterator[Union[Block, List[BlockMetadata]]]:
+        yield from _map_task(
+            self._map_transformer,
+            data_context,
+            ctx,
+            *blocks,
+            slices=slices,
+            **kwargs,
+        )
+
+    def __repr__(self):
+        return f"MapWorker({self.src_fn_name})"
+
+    def on_exit(self):
+        """Called when the actor is about to exist.
+        This enables performing cleanup operations via `UDF.__del__`.
+
+        Note, this only ensures cleanup is performed when the job exists gracefully.
+        If the driver or the actor is forcefully killed, `__del__` will not be called.
+        """
+        # `_map_actor_context` is a global variable that references the UDF object.
+        # Delete it to trigger `UDF.__del__`.
+        del ray.data._map_actor_context
+        ray.data._map_actor_context = None
+
+
+@dataclass
+class _ActorState:
+    """Actor state"""
+
+    # Number of tasks in flight per actor
+    num_tasks_in_flight: int
+
+    # Node id of each ready actor
+    actor_location: str
+
+    # Is Actor state restarting or alive
+    is_restarting: bool
+
+
+class _ActorTaskSelector(abc.ABC):
+    def __init__(self, actor_pool: "_ActorPool"):
+        """Initialize the actor task selector.
+
+        Args:
+            actor_pool: The actor pool to select tasks from.
+        """
+        self._actor_pool = actor_pool
+
+    @abstractmethod
+    def select_actors(
+        self, input_queue: BundleQueue, actor_locality_enabled: bool
+    ) -> Iterator[Tuple[RefBundle, ActorHandle]]:
+        """Select actors for bundles in the input queue.
+
+        Args:
+            input_queue: The input queue to select actors for.
+            actor_locality_enabled: Whether actor locality is enabled.
+
+        Returns:
+            Iterator of tuples of the bundle and the selected actor for that bundle.
+            Iteration stops when there are no more bundles to be selected in the input queue
+        """
+        pass
+
+
+class _ActorTaskSelectorImpl(_ActorTaskSelector):
+    def __init__(self, actor_pool: "_ActorPool"):
+        super().__init__(actor_pool)
+
+    def select_actors(
+        self, input_queue: BundleQueue, actor_locality_enabled: bool
+    ) -> Iterator[Tuple[RefBundle, ActorHandle]]:
+        """Picks actors for task submission based on busyness and locality."""
+        if not self._actor_pool.running_actors():
+            # Actor pool is empty or all actors are still pending.
+            return
+
+        while input_queue:
+            # Filter out actors that are invalid, i.e. actors with number of tasks in
+            # flight >= _max_tasks_in_flight or actor_state is not ALIVE.
+            bundle = input_queue.peek_next()
+            valid_actors = [
+                actor
+                for actor in self._actor_pool.running_actors()
+                if self._actor_pool.running_actors()[actor].num_tasks_in_flight
+                < self._actor_pool.max_tasks_in_flight_per_actor()
+                and not self._actor_pool.running_actors()[actor].is_restarting
+            ]
+
+            if not valid_actors:
+                # All actors are at capacity or actor state is not ALIVE.
+                return
+
+            # Rank all valid actors
+            ranks = self._rank_actors(
+                valid_actors, bundle if actor_locality_enabled else None
+            )
+
+            assert len(ranks) == len(
+                valid_actors
+            ), f"{len(ranks)} != {len(valid_actors)}"
+
+            # Pick the actor with the highest rank (lower value, higher rank)
+            target_actor_idx = min(range(len(valid_actors)), key=lambda idx: ranks[idx])
+
+            target_actor = valid_actors[target_actor_idx]
+
+            # We remove the bundle and yield the actor to the operator. We do not use pop()
+            # in case the queue has changed the order of the bundles.
+            input_queue.remove(bundle)
+            yield bundle, target_actor
+
+    def _rank_actors(
+        self,
+        actors: List[ActorHandle],
+        bundle: Optional[RefBundle],
+    ) -> List[Tuple[int, int]]:
+        """Return ranks for each actor based on node affinity with the blocks in the provided
+        bundle and current Actor's load.
+
+        The rank for each actor is a tuple of
+
+            1. Locality rank: a rank of a node Actor is scheduled on determined based on
+            the ranking of preferred locations for provided ``RefBundle`` (defined by
+            ``RefBundle.get_preferred_locations``). Lower is better.
+            2. Number of tasks currently executed by Actor. Lower is better.
+
+        Args:
+            actors: List of actors to rank
+            bundle: Optional bundle whose locality preferences should be considered
+
+        Returns:
+            List of (locality_rank, num_tasks) tuples, one per input actor
+        """
+        locs_priorities = (
+            {
+                # NOTE: We're negating total bytes to maintain an invariant
+                #       of the rank used -- lower value corresponding to a higher rank
+                node_id: -total_bytes
+                for node_id, total_bytes in bundle.get_preferred_object_locations().items()
+            }
+            if bundle is not None
+            else {}
+        )
+
+        # NOTE: Ranks are ordered in descending order (ie rank[0] is the highest
+        #       and rank[-1] is the lowest)
+        ranks = [
+            (
+                # Priority/rank of the location (based on the object size).
+                # Defaults to int32 max value (ie no rank)
+                locs_priorities.get(
+                    self._actor_pool.running_actors()[actor].actor_location, INT32_MAX
+                ),
+                # Number of tasks currently in flight at the given actor
+                self._actor_pool.running_actors()[actor].num_tasks_in_flight,
+            )
+            for actor in actors
+        ]
+
+        return ranks
+
+
+class _ActorPool(AutoscalingActorPool):
+    """A pool of actors for map task execution.
+
+    This class is in charge of tracking the number of in-flight tasks per actor,
+    providing the least heavily loaded actor to the operator, and killing idle
+    actors when the operator is done submitting work to the pool.
+    """
+
+    _ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S = 10
+    _ACTOR_POOL_GRACEFUL_SHUTDOWN_TIMEOUT_S = 30
+    _LOGICAL_ACTOR_ID_LABEL_KEY = "__ray_data_logical_actor_id"
+
+    def __init__(
+        self,
+        create_actor_fn: "Callable[[Dict[str, str]], Tuple[ActorHandle, ObjectRef[Any]]]",
+        per_actor_resource_usage: ExecutionResources,
+        *,
+        min_size: int,
+        max_size: int,
+        initial_size: int,
+        max_actor_concurrency: int,
+        max_tasks_in_flight_per_actor: int,
+        _enable_actor_pool_on_exit_hook: bool = False,
+    ):
+        """Initialize the actor pool.
+
+        Args:
+            create_actor_fn: This function should take key-value labels as input, and
+                create an actor with those labels. The function should return the actor
+                handle and a reference to the actor's node ID.
+            per_actor_resource_usage: The resource usage per actor.
+            min_size: The minimum number of running actors to be maintained
+                in the pool. Note, that this constraint could be violated when
+                no new work is available for scheduling in the actor pool (ie
+                when operator completes execution).
+            max_size: The maximum number of running actors to be maintained
+                in the pool.
+            initial_size: The initial number of actors to start with.
+            max_actor_concurrency: The maximum number of concurrent tasks a
+                single actor can execute (derived from `ray_remote_args`
+                passed to the operator).
+            max_tasks_in_flight_per_actor: The maximum number of tasks that can
+                be submitted to a single actor at any given time.
+            _enable_actor_pool_on_exit_hook: Whether to enable the actor pool
+                on exit hook.
+        """
+
+        self._min_size: int = min_size
+        self._max_size: int = max_size
+        self._initial_size: int = initial_size
+        self._max_actor_concurrency: int = max_actor_concurrency
+        self._max_tasks_in_flight: int = max_tasks_in_flight_per_actor
+        self._create_actor_fn = create_actor_fn
+        self._per_actor_resource_usage = per_actor_resource_usage
+
+        assert self._min_size >= 1
+        assert self._max_size >= self._min_size
+        assert self._initial_size <= self._max_size
+        assert self._initial_size >= self._min_size
+        assert self._max_tasks_in_flight >= 1
+        assert self._create_actor_fn is not None
+
+        # Timestamp of the last scale up action
+        self._last_upscaled_at: Optional[float] = None
+        self._last_downscaling_debounce_warning_ts: Optional[float] = None
+        # Actors that have started running, including alive and restarting actors.
+        self._running_actors: Dict[ray.actor.ActorHandle, _ActorState] = {}
+        # Actors that are not yet ready (still pending creation).
+        self._pending_actors: Dict[ObjectRef, ray.actor.ActorHandle] = {}
+        # Map from actor handle to its logical ID.
+        self._actor_to_logical_id: Dict[ray.actor.ActorHandle, str] = {}
+        self._enable_actor_pool_on_exit_hook = _enable_actor_pool_on_exit_hook
+        # Cached values for actor / task counts
+        self._num_restarting_actors: int = 0
+        self._num_active_actors: int = 0
+        self._total_num_tasks_in_flight: int = 0
+
+    # === Overriding methods of AutoscalingActorPool ===
+
+    def min_size(self) -> int:
+        return self._min_size
+
+    def max_size(self) -> int:
+        return self._max_size
+
+    def current_size(self) -> int:
+        return self.num_pending_actors() + self.num_running_actors()
+
+    def num_running_actors(self) -> int:
+        return len(self._running_actors)
+
+    def num_restarting_actors(self) -> int:
+        """Restarting actors are all the running actors not in ALIVE state."""
+        return self._num_restarting_actors
+
+    def num_active_actors(self) -> int:
+        """Active actors are all the running actors with inflight tasks."""
+        return self._num_active_actors
+
+    def num_alive_actors(self) -> int:
+        """Alive actors are all the running actors in ALIVE state."""
+        return len(self._running_actors) - self._num_restarting_actors
+
+    def num_pending_actors(self) -> int:
+        return len(self._pending_actors)
+
+    def max_tasks_in_flight_per_actor(self) -> int:
+        return self._max_tasks_in_flight
+
+    def max_actor_concurrency(self) -> int:
+        return self._max_actor_concurrency
+
+    def num_tasks_in_flight(self) -> int:
+        return self._total_num_tasks_in_flight
+
+    def initial_size(self) -> int:
+        return self._initial_size
+
+    def _can_apply(self, config: ActorPoolScalingRequest) -> bool:
+        """Returns whether Actor Pool is able to execute scaling request"""
+
+        if config.delta < 0:
+            # To prevent bouncing back and forth, we disallow scale down for
+            # a "cool-off" period after the most recent scaling up, with an intention
+            # to allow application to actually utilize newly provisioned resources
+            # before making decisions on subsequent actions.
+            #
+            # Note that this action is unidirectional and doesn't apply to
+            # scaling up, ie if actor pool just scaled down, it'd still be able
+            # to scale back up immediately.
+            if (
+                not config.force
+                and self._last_upscaled_at is not None
+                and (
+                    time.time()
+                    <= self._last_upscaled_at
+                    + self._ACTOR_POOL_SCALE_DOWN_DEBOUNCE_PERIOD_S
+                )
+            ):
+                # NOTE: To avoid spamming logs unnecessarily, debounce log is produced once
+                #       per upscaling event
+                if self._last_upscaled_at != self._last_downscaling_debounce_warning_ts:
+                    logger.debug(
+                        f"Ignoring scaling down request (request={config}; reason=debounced from scaling up at {self._last_upscaled_at})"
+                    )
+                    self._last_downscaling_debounce_warning_ts = self._last_upscaled_at
+
+                return False
+
+        return True
+
+    def scale(self, req: ActorPoolScalingRequest) -> Optional[int]:
+        # Verify request could be applied
+        if not self._can_apply(req):
+            return 0
+
+        if req.delta > 0:
+            target_num_actors = req.delta
+
+            logger.debug(
+                f"Scaling up actor pool by {target_num_actors} (reason={req.reason}, "
+                f"{self.get_actor_info()})"
+            )
+
+            for _ in range(target_num_actors):
+                actor, ready_ref = self._create_actor()
+                self.add_pending_actor(actor, ready_ref)
+
+            # Capture last scale up timestamp
+            self._last_upscaled_at = time.time()
+
+            return target_num_actors
+
+        elif req.delta < 0:
+            num_released = 0
+            target_num_actors = abs(req.delta)
+
+            for _ in range(target_num_actors):
+                if self._remove_inactive_actor():
+                    num_released += 1
+
+            if num_released > 0:
+                logger.debug(
+                    f"Scaled down actor pool by {num_released} "
+                    f"(reason={req.reason}; {self.get_actor_info()})"
+                )
+
+            return -num_released
+
+        return None
+
+    def _create_actor(self) -> Tuple[ray.actor.ActorHandle, ObjectRef]:
+        logical_actor_id = str(uuid.uuid4())
+        labels = {self.get_logical_id_label_key(): logical_actor_id}
+        actor, ready_ref = self._create_actor_fn(labels, logical_actor_id)
+        self._actor_to_logical_id[actor] = logical_actor_id
+        return actor, ready_ref
+
+    # === End of overriding methods of AutoscalingActorPool ===
+
+    def running_actors(self) -> Dict[ray.actor.ActorHandle, _ActorState]:
+        return self._running_actors
+
+    def on_task_submitted(self, actor: ray.actor.ActorHandle):
+        self._running_actors[actor].num_tasks_in_flight += 1
+        self._total_num_tasks_in_flight += 1
+
+        if self._running_actors[actor].num_tasks_in_flight == 1:
+            self._num_active_actors += 1
+
+    def update_running_actor_state(
+        self, actor: ray.actor.ActorHandle, is_restarting: bool
+    ) -> None:
+        """Update running actor state.
+
+        Args:
+            actor: The running actor that needs state update.
+            is_restarting: Whether running actor is restarting or alive.
+        """
+        assert actor in self._running_actors
+        if self._running_actors[actor].is_restarting == is_restarting:
+            return
+
+        self._running_actors[actor].is_restarting = is_restarting
+        if is_restarting:
+            self._num_restarting_actors += 1
+        else:
+            self._num_restarting_actors -= 1
+
+    def add_pending_actor(self, actor: ray.actor.ActorHandle, ready_ref: ray.ObjectRef):
+        """Adds a pending actor to the pool.
+
+        This actor won't be pickable until it is marked as running via a
+        pending_to_running() call.
+
+        Args:
+            actor: The not-yet-ready actor to add as pending to the pool.
+            ready_ref: The ready future for the actor.
+        """
+        self._pending_actors[ready_ref] = actor
+
+    def pending_to_running(self, ready_ref: ray.ObjectRef) -> bool:
+        """Mark the actor corresponding to the provided ready future as running, making
+        the actor pickable.
+
+        Args:
+            ready_ref: The ready future for the actor that we wish to mark as running.
+
+        Returns:
+            Whether the actor was still pending. This can return False if the actor had
+            already been killed.
+        """
+        if ready_ref not in self._pending_actors:
+            # The actor has been removed from the pool before becoming running.
+            return False
+        actor = self._pending_actors.pop(ready_ref)
+        self._running_actors[actor] = _ActorState(
+            num_tasks_in_flight=0,
+            actor_location=ray.get(ready_ref),
+            is_restarting=False,
+        )
+        return True
+
+    def on_task_completed(self, actor: ray.actor.ActorHandle):
+        """Called when a task completes. Returns the provided actor to the pool."""
+        assert actor in self._running_actors
+        assert self._running_actors[actor].num_tasks_in_flight > 0
+        self._running_actors[actor].num_tasks_in_flight -= 1
+        self._total_num_tasks_in_flight -= 1
+        if not self._running_actors[actor].num_tasks_in_flight:
+            self._num_active_actors -= 1
+
+    def get_pending_actor_refs(self) -> List[ray.ObjectRef]:
+        return list(self._pending_actors.keys())
+
+    def get_running_actor_refs(self) -> List[ray.ObjectRef]:
+        return list(self._running_actors.keys())
+
+    def get_logical_ids(self) -> List[str]:
+        """Get the logical IDs for pending and running actors in the actor pool.
+
+        We can’t use Ray Core actor IDs because we need to identify actors by labels,
+        but labels must be set before creation, and actor IDs aren’t available until
+        after.
+        """
+        return list(self._actor_to_logical_id.values())
+
+    def get_logical_id_label_key(self) -> str:
+        """Get the label key for the logical actor ID.
+
+        Actors launched by this pool should have this label.
+        """
+        return self._LOGICAL_ACTOR_ID_LABEL_KEY
+
+    def num_idle_actors(self) -> int:
+        """Return the number of idle actors in the pool."""
+        return len(self._running_actors) - self._num_active_actors
+
+    def _remove_inactive_actor(self) -> bool:
+        """Kills a single pending or idle actor, if any actors are pending/idle.
+
+        Returns whether an inactive actor was actually released.
+        """
+        # We prioritize killing pending actors over idle actors to reduce actor starting
+        # churn.
+        released = self._try_remove_pending_actor()
+        if not released:
+            # If no pending actor was released, so kill actor.
+            released = self._try_remove_idle_actor()
+        return released
+
+    def _try_remove_pending_actor(self) -> bool:
+        if self._pending_actors:
+            # At least one pending actor, so kill first one.
+            ready_ref = next(iter(self._pending_actors.keys()))
+            actor = self._pending_actors.pop(ready_ref)
+            del self._actor_to_logical_id[actor]
+            return True
+        # No pending actors, so indicate to the caller that no actors were killed.
+        return False
+
+    def _try_remove_idle_actor(self) -> bool:
+        for actor, state in self._running_actors.items():
+            if state.num_tasks_in_flight == 0:
+                # At least one idle actor, so kill first one found.
+                # NOTE: This is a fire-and-forget op
+                self._release_running_actor(actor)
+                return True
+        # No idle actors, so indicate to the caller that no actors were killed.
+        return False
+
+    def shutdown(self, force: bool = False):
+        """Kills all actors, including running/active actors.
+
+        This is called once the operator is shutting down.
+        """
+        self._release_pending_actors(force=force)
+        self._release_running_actors(force=force)
+
+    def _release_pending_actors(self, force: bool):
+        # Release pending actors from the set of pending ones
+        pending = dict(self._pending_actors)
+        self._pending_actors.clear()
+
+        if force:
+            for _, actor in pending.items():
+                # NOTE: Actors can't be brought back after being ``ray.kill``-ed,
+                #       hence we're only doing that if this is a forced release
+                ray.kill(actor)
+
+    def _release_running_actors(self, force: bool):
+        running = list(self._running_actors.keys())
+
+        on_exit_refs = []
+
+        # First release actors and collect their shutdown hook object-refs
+        for actor in running:
+            ref = self._release_running_actor(actor)
+            if ref:
+                on_exit_refs.append(ref)
+
+        # Wait for all actors to shutdown gracefully before killing them
+        ray.wait(on_exit_refs, timeout=self._ACTOR_POOL_GRACEFUL_SHUTDOWN_TIMEOUT_S)
+
+        # NOTE: Actors can't be brought back after being ``ray.kill``-ed,
+        #       hence we're only doing that if this is a forced release
+        if force:
+            for actor in running:
+                ray.kill(actor)
+
+    def _release_running_actor(
+        self, actor: ray.actor.ActorHandle
+    ) -> Optional[ObjectRef]:
+        """Remove the given actor from the pool and trigger its `on_exit` callback.
+
+        This method returns a ``ref`` to the result
+        """
+        # NOTE: By default, we remove references to the actor and let ref counting
+        # garbage collect the actor, instead of using ray.kill.
+        #
+        # Otherwise, actor cannot be reconstructed for the purposes of produced
+        # object's lineage reconstruction.
+        if actor not in self._running_actors:
+            return None
+
+        # Update cached statistics before removing the actor
+        actor_state = self._running_actors[actor]
+
+        # Update total tasks in flight
+        self._total_num_tasks_in_flight -= actor_state.num_tasks_in_flight
+
+        # Update active actors count
+        if actor_state.num_tasks_in_flight > 0:
+            self._num_active_actors -= 1
+
+        # Update restarting actors count
+        if actor_state.is_restarting:
+            self._num_restarting_actors -= 1
+
+        if self._enable_actor_pool_on_exit_hook:
+            # Call `on_exit` to trigger `UDF.__del__` which may perform
+            # cleanup operations.
+            ref = actor.on_exit.remote()
+        else:
+            ref = None
+        del self._running_actors[actor]
+        del self._actor_to_logical_id[actor]
+
+        return ref
+
+    def get_actor_info(self) -> _ActorPoolInfo:
+        """Returns current snapshot of actors' being used in the pool"""
+        return _ActorPoolInfo(
+            running=self.num_alive_actors(),
+            pending=self.num_pending_actors(),
+            restarting=self.num_restarting_actors(),
+        )
+
+    def per_actor_resource_usage(self) -> ExecutionResources:
+        """Per actor resource usage."""
+        return self._per_actor_resource_usage
+
+    def get_pool_util(self) -> float:
+        if self.num_running_actors() == 0:
+            return 0.0
+        else:
+            # We compute utilization as a ration of
+            #  - Number of submitted tasks over
+            #  - Max number of tasks that Actor Pool could currently run
+            #
+            # This value could exceed 100%, since by default actors are allowed
+            # to queue tasks (to pipeline task execution by overlapping block
+            # fetching with the execution of the previous task)
+            return self.num_tasks_in_flight() / (
+                self._max_actor_concurrency * self.num_running_actors()
+            )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/aggregate_num_rows.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/aggregate_num_rows.py
new file mode 100644
index 0000000000000000000000000000000000000000..68084d2d0ad7d24ddb8eea70fe31f2f1b774f1d1
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/aggregate_num_rows.py
@@ -0,0 +1,64 @@
+import ray
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import PhysicalOperator, RefBundle
+from ray.data._internal.stats import StatsDict
+from ray.data.block import BlockAccessor
+from ray.data.context import DataContext
+
+
+class AggregateNumRows(PhysicalOperator):
+    """Count number of rows in input bundles.
+
+    This operator aggregates the number of rows in input bundles using the bundles'
+    block metadata. It outputs a single row with the specified column name.
+    """
+
+    def __init__(
+        self,
+        input_dependencies,
+        data_context: DataContext,
+        column_name: str,
+    ):
+        super().__init__(
+            "AggregateNumRows",
+            input_dependencies,
+            data_context,
+        )
+
+        self._column_name = column_name
+
+        self._num_rows = 0
+        self._has_outputted = False
+        self._estimated_num_output_bundles = 1
+        self._estimated_output_num_rows = 1
+
+    def has_next(self) -> bool:
+        return self._inputs_complete and not self._has_outputted
+
+    def _get_next_inner(self) -> RefBundle:
+        assert self._inputs_complete
+
+        builder = DelegatingBlockBuilder()
+        builder.add({self._column_name: self._num_rows})
+        block = builder.build()
+        block_ref = ray.put(block)
+
+        metadata = BlockAccessor.for_block(block).get_metadata()
+        schema = BlockAccessor.for_block(block).schema()
+        bundle = RefBundle([(block_ref, metadata)], owns_blocks=True, schema=schema)
+
+        self._has_outputted = True
+        return bundle
+
+    def get_stats(self) -> StatsDict:
+        return {}
+
+    def _add_input_inner(self, refs, input_index) -> None:
+        assert refs.num_rows() is not None
+        self._num_rows += refs.num_rows()
+
+    def throttling_disabled(self) -> bool:
+        return True
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        return True
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/base_physical_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/base_physical_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bd1a5c9f79448b4e0fd5d77e3c4de251880e07c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/base_physical_operator.py
@@ -0,0 +1,277 @@
+import abc
+from typing import List, Optional
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.physical_operator import _create_sub_pb
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
+from ray.data._internal.logical.interfaces import LogicalOperator
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+
+class InternalQueueOperatorMixin(PhysicalOperator, abc.ABC):
+    @abc.abstractmethod
+    def internal_input_queue_num_blocks(self) -> int:
+        """Returns Operator's internal input queue size (in blocks)"""
+        ...
+
+    @abc.abstractmethod
+    def internal_input_queue_num_bytes(self) -> int:
+        """Returns Operator's internal input queue size (in bytes)"""
+        ...
+
+    @abc.abstractmethod
+    def internal_output_queue_num_blocks(self) -> int:
+        """Returns Operator's internal output queue size (in blocks)"""
+        ...
+
+    @abc.abstractmethod
+    def internal_output_queue_num_bytes(self) -> int:
+        """Returns Operator's internal output queue size (in bytes)"""
+        ...
+
+    @abc.abstractmethod
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queue(s).
+
+        This should drain all buffered input bundles and update metrics appropriately
+        by calling on_input_dequeued().
+        """
+        ...
+
+    @abc.abstractmethod
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue(s).
+
+        This should drain all buffered output bundles and update metrics appropriately
+        by calling on_output_dequeued().
+        """
+        ...
+
+    def mark_execution_finished(self) -> None:
+        """Mark execution as finished and clear internal queues.
+
+        This default implementation calls the parent's mark_execution_finished()
+        and then clears internal input and output queues.
+        """
+        super().mark_execution_finished()
+        self.clear_internal_input_queue()
+        self.clear_internal_output_queue()
+
+
+class OneToOneOperator(PhysicalOperator):
+    """An operator that has one input and one output dependency.
+
+    This operator serves as the base for map, filter, limit, etc.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        target_max_block_size_override: Optional[int] = None,
+    ):
+        """Create a OneToOneOperator.
+        Args:
+            input_op: Operator generating input data for this op.
+            name: The name of this operator.
+            target_max_block_size_override: The target maximum number of bytes to
+                include in an output block.
+        """
+        super().__init__(name, [input_op], data_context, target_max_block_size_override)
+
+    @property
+    def input_dependency(self) -> PhysicalOperator:
+        return self.input_dependencies[0]
+
+
+class AllToAllOperator(
+    InternalQueueOperatorMixin, SubProgressBarMixin, PhysicalOperator
+):
+    """A blocking operator that executes once its inputs are complete.
+
+    This operator implements distributed sort / shuffle operations, etc.
+    """
+
+    def __init__(
+        self,
+        bulk_fn: AllToAllTransformFn,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        target_max_block_size_override: Optional[int] = None,
+        num_outputs: Optional[int] = None,
+        sub_progress_bar_names: Optional[List[str]] = None,
+        name: str = "AllToAll",
+    ):
+        """Create an AllToAllOperator.
+        Args:
+            bulk_fn: The blocking transformation function to run. The inputs are the
+                list of input ref bundles, and the outputs are the output ref bundles
+                and a stats dict.
+            input_op: Operator generating input data for this op.
+            data_context: The DataContext instance containing configuration settings.
+            target_max_block_size_override: The target maximum number of bytes to
+                include in an output block.
+            num_outputs: The number of expected output bundles for progress bar.
+            sub_progress_bar_names: The names of internal sub progress bars.
+            name: The name of this operator.
+        """
+        self._bulk_fn = bulk_fn
+        self._next_task_index = 0
+        self._num_outputs = num_outputs
+        self._output_rows = 0
+        self._sub_progress_bar_names = sub_progress_bar_names
+        self._sub_progress_bar_dict = None
+        self._input_buffer: List[RefBundle] = []
+        self._output_buffer: List[RefBundle] = []
+        self._stats: StatsDict = {}
+        super().__init__(name, [input_op], data_context, target_max_block_size_override)
+
+    def num_outputs_total(self) -> Optional[int]:
+        return (
+            self._num_outputs
+            if self._num_outputs
+            else self.input_dependencies[0].num_outputs_total()
+        )
+
+    def num_output_rows_total(self) -> Optional[int]:
+        return (
+            self._output_rows
+            if self._output_rows
+            else self.input_dependencies[0].num_output_rows_total()
+        )
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.completed()
+        assert input_index == 0, input_index
+        self._input_buffer.append(refs)
+        self._metrics.on_input_queued(refs)
+
+    def internal_input_queue_num_blocks(self) -> int:
+        return sum(len(bundle.block_refs) for bundle in self._input_buffer)
+
+    def internal_input_queue_num_bytes(self) -> int:
+        return sum(bundle.size_bytes() for bundle in self._input_buffer)
+
+    def internal_output_queue_num_blocks(self) -> int:
+        return sum(len(bundle.block_refs) for bundle in self._output_buffer)
+
+    def internal_output_queue_num_bytes(self) -> int:
+        return sum(bundle.size_bytes() for bundle in self._output_buffer)
+
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queue."""
+        while self._input_buffer:
+            bundle = self._input_buffer.pop()
+            self._metrics.on_input_dequeued(bundle)
+
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue."""
+        while self._output_buffer:
+            bundle = self._output_buffer.pop()
+            self._metrics.on_output_dequeued(bundle)
+
+    def all_inputs_done(self) -> None:
+        ctx = TaskContext(
+            task_idx=self._next_task_index,
+            op_name=self.name,
+            sub_progress_bar_dict=self._sub_progress_bar_dict,
+            target_max_block_size_override=self.target_max_block_size_override,
+        )
+        # NOTE: We don't account object store memory use from intermediate `bulk_fn`
+        # outputs (e.g., map outputs for map-reduce).
+        self._output_buffer, self._stats = self._bulk_fn(self._input_buffer, ctx)
+
+        while self._input_buffer:
+            refs = self._input_buffer.pop()
+            self._metrics.on_input_dequeued(refs)
+
+        for ref in self._output_buffer:
+            self._metrics.on_output_queued(ref)
+
+        self._next_task_index += 1
+
+        super().all_inputs_done()
+
+    def has_next(self) -> bool:
+        return len(self._output_buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        bundle = self._output_buffer.pop(0)
+        self._metrics.on_output_dequeued(bundle)
+        self._output_rows += bundle.num_rows()
+        return bundle
+
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    def get_transformation_fn(self) -> AllToAllTransformFn:
+        return self._bulk_fn
+
+    def progress_str(self) -> str:
+        return f"{self.num_output_rows_total() or 0} rows output"
+
+    def initialize_sub_progress_bars(self, position: int) -> int:
+        """Initialize all internal sub progress bars, and return the number of bars."""
+        if self._sub_progress_bar_names is not None:
+            self._sub_progress_bar_dict = {}
+            for name in self._sub_progress_bar_names:
+                bar, position = _create_sub_pb(
+                    name, self.num_output_rows_total(), position
+                )
+                self._sub_progress_bar_dict[name] = bar
+            return len(self._sub_progress_bar_dict)
+        else:
+            return 0
+
+    def close_sub_progress_bars(self):
+        """Close all internal sub progress bars."""
+        if self._sub_progress_bar_dict is not None:
+            for sub_bar in self._sub_progress_bar_dict.values():
+                sub_bar.close()
+
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        return self._sub_progress_bar_names
+
+    def set_sub_progress_bar(self, name, pg):
+        # not type-checking due to circular imports
+        if self._sub_progress_bar_dict is None:
+            self._sub_progress_bar_dict = {}
+        self._sub_progress_bar_dict[name] = pg
+
+    def supports_fusion(self):
+        return True
+
+    def implements_accurate_memory_accounting(self):
+        return True
+
+
+class NAryOperator(PhysicalOperator):
+    """An operator that has multiple input dependencies and one output.
+
+    This operator serves as the base for union, zip, etc.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: LogicalOperator,
+    ):
+        """Create a OneToOneOperator.
+        Args:
+            input_op: Operator generating input data for this op.
+            name: The name of this operator.
+        """
+        input_names = ", ".join([op._name for op in input_ops])
+        op_name = f"{self.__class__.__name__}({input_names})"
+        super().__init__(
+            op_name,
+            list(input_ops),
+            data_context,
+        )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/hash_aggregate.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/hash_aggregate.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6ff69dd8e517d3b2b03b86fc1dbe088eefea84c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/hash_aggregate.py
@@ -0,0 +1,229 @@
+import logging
+import math
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+from ray.data._internal.arrow_block import ArrowBlockAccessor
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.hash_shuffle import (
+    BlockTransformer,
+    HashShufflingOperatorBase,
+    StatefulShuffleAggregation,
+)
+from ray.data._internal.util import GiB, MiB
+from ray.data.aggregate import AggregateFn
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+
+logger = logging.getLogger(__name__)
+
+
+class ReducingShuffleAggregation(StatefulShuffleAggregation):
+    """Aggregation performing reduction of the shuffled sequence using provided
+    list of aggregating functions.
+
+    NOTE: That reductions are performed incrementally in a streaming fashion upon
+          accumulation of pre-configured buffer of rows to run aggregation on."""
+
+    _DEFAULT_BLOCKS_BUFFER_LIMIT = 1000
+
+    def __init__(
+        self,
+        aggregator_id: int,
+        key_columns: Optional[Tuple[str]],
+        aggregation_fns: Tuple[AggregateFn],
+    ):
+
+        super().__init__(aggregator_id)
+
+        assert key_columns is not None, "Shuffle aggregation requires key columns"
+
+        self._sort_key: "SortKey" = ReducingShuffleAggregation._get_sort_key(
+            key_columns
+        )
+        self._aggregation_fns: Tuple[AggregateFn] = aggregation_fns
+
+        self._aggregated_blocks: List[Block] = []
+
+    def accept(self, input_seq_id: int, partition_id: int, partition_shard: Block):
+        assert (
+            input_seq_id == 0
+        ), f"Single sequence is expected (got seq-id {input_seq_id})"
+
+        # Received partition shard is already partially aggregated, hence
+        # we simply add it to the list of aggregated blocks
+        #
+        # NOTE: We're not separating blocks by partition as it's ultimately not
+        #       relevant for the aggregations performed
+        self._aggregated_blocks.append(partition_shard)
+
+        # Aggregation is performed incrementally, rather
+        # than being deferred to the finalization stage
+        if len(self._aggregated_blocks) > self._DEFAULT_BLOCKS_BUFFER_LIMIT:
+            # NOTE: This method will reset partially aggregated blocks to hold
+            #       the new combined one
+            #
+            # TODO make aggregation async
+            self._combine_aggregated_blocks(should_finalize=False)
+
+    def finalize(self, partition_id: int) -> Block:
+        if len(self._aggregated_blocks) == 0:
+            return ArrowBlockAccessor._empty_table()
+
+        return self._combine_aggregated_blocks(should_finalize=True)
+
+    def clear(self, partition_id: int):
+        self._aggregated_blocks: List[Block] = []
+
+    def _combine_aggregated_blocks(self, *, should_finalize: bool) -> Block:
+        assert len(self._aggregated_blocks) > 0
+
+        block_accessor = BlockAccessor.for_block(self._aggregated_blocks[0])
+        combined_block, _ = block_accessor._combine_aggregated_blocks(
+            self._aggregated_blocks,
+            sort_key=self._sort_key,
+            aggs=self._aggregation_fns,
+            finalize=should_finalize,
+        )
+
+        # For combined block that's not yet finalized reset cached aggregated
+        # blocks to only hold newly combined one
+        if not should_finalize:
+            self._aggregated_blocks = [combined_block]
+
+        return combined_block
+
+    @staticmethod
+    def _get_sort_key(key_columns: Tuple[str]):
+        from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+        return SortKey(key=list(key_columns), descending=False)
+
+
+class HashAggregateOperator(HashShufflingOperatorBase):
+    def __init__(
+        self,
+        data_context: DataContext,
+        input_op: PhysicalOperator,
+        key_columns: Tuple[str],
+        aggregation_fns: Tuple[AggregateFn],
+        *,
+        num_partitions: Optional[int] = None,
+        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            name_factory=(
+                lambda num_partitions: f"HashAggregate(key_columns={key_columns}, "
+                f"num_partitions={num_partitions})"
+            ),
+            input_ops=[input_op],
+            data_context=data_context,
+            key_columns=[key_columns],
+            num_partitions=(
+                # NOTE: In case of global aggregations (ie with no key columns specified),
+                #       we override number of partitions to 1, since the whole dataset
+                #       will be reduced to just a single row
+                num_partitions
+                if len(key_columns) > 0
+                else 1
+            ),
+            partition_aggregation_factory=(
+                lambda aggregator_id, target_partition_ids: ReducingShuffleAggregation(
+                    aggregator_id,
+                    key_columns,
+                    aggregation_fns,
+                )
+            ),
+            input_block_transformer=_create_aggregating_transformer(
+                key_columns, aggregation_fns
+            ),
+            aggregator_ray_remote_args_override=aggregator_ray_remote_args_override,
+            shuffle_progress_bar_name="Shuffle",
+            finalize_progress_bar_name="Aggregation",
+        )
+
+    def _get_operator_num_cpus_override(self) -> float:
+        return self.data_context.hash_aggregate_operator_actor_num_cpus_override
+
+    @classmethod
+    def _estimate_aggregator_memory_allocation(
+        cls,
+        *,
+        num_aggregators: int,
+        num_partitions: int,
+        estimated_dataset_bytes: int,
+    ) -> int:
+        partition_byte_size_estimate = math.ceil(
+            estimated_dataset_bytes / num_partitions
+        )
+
+        # Estimate of object store memory required to accommodate all partitions
+        # handled by a single aggregator
+        aggregator_shuffle_object_store_memory_required: int = math.ceil(
+            estimated_dataset_bytes / num_aggregators
+        )
+        # Estimate of memory required to accommodate single partition as an output
+        # (inside Object Store)
+        output_object_store_memory_required: int = partition_byte_size_estimate
+
+        aggregator_total_memory_required: int = (
+            # Inputs (object store)
+            aggregator_shuffle_object_store_memory_required
+            +
+            # Output (object store)
+            output_object_store_memory_required
+        )
+
+        logger.info(
+            f"Estimated memory requirement for aggregating aggregator "
+            f"(partitions={num_partitions}, "
+            f"aggregators={num_aggregators}, "
+            f"dataset (estimate)={estimated_dataset_bytes / GiB:.1f}GiB): "
+            f"shuffle={aggregator_shuffle_object_store_memory_required / MiB:.1f}MiB, "
+            f"output={output_object_store_memory_required / MiB:.1f}MiB, "
+            f"total={aggregator_total_memory_required / MiB:.1f}MiB, "
+        )
+
+        return aggregator_total_memory_required
+
+
+def _create_aggregating_transformer(
+    key_columns: Tuple[str], aggregation_fns: Tuple[AggregateFn]
+) -> BlockTransformer:
+    """Method creates input block transformer performing partial aggregation of
+    the block applied prior to block being shuffled (to reduce amount of bytes shuffled)"""
+
+    sort_key = ReducingShuffleAggregation._get_sort_key(key_columns)
+
+    def _aggregate(block: Block) -> Block:
+        from ray.data._internal.planner.exchange.aggregate_task_spec import (
+            SortAggregateTaskSpec,
+        )
+
+        # TODO unify blocks schemas, to avoid validating every block
+        # Validate block's schema compatible with aggregations
+        for agg_fn in aggregation_fns:
+            agg_fn._validate(BlockAccessor.for_block(block).schema())
+
+        # Project block to only carry columns used in aggregation
+        pruned_block = SortAggregateTaskSpec._prune_unused_columns(
+            block,
+            sort_key,
+            aggregation_fns,
+        )
+
+        # NOTE: If columns to aggregate on have been provided,
+        #       sort the block on these before aggregation
+        if sort_key.get_columns():
+            target_block = BlockAccessor.for_block(pruned_block).sort(sort_key)
+        else:
+            target_block = pruned_block
+
+        return BlockAccessor.for_block(target_block)._aggregate(
+            sort_key, aggregation_fns
+        )
+
+    return _aggregate
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/hash_shuffle.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/hash_shuffle.py
new file mode 100644
index 0000000000000000000000000000000000000000..8168fb8b0a4c726183a87fb9003c0672def7396c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/hash_shuffle.py
@@ -0,0 +1,1604 @@
+import abc
+import functools
+import itertools
+import logging
+import math
+import random
+import threading
+import time
+from collections import defaultdict, deque
+from dataclasses import dataclass
+from typing import (
+    Any,
+    AsyncGenerator,
+    Callable,
+    DefaultDict,
+    Deque,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+import pyarrow as pa
+
+import ray
+from ray import ObjectRef
+from ray._private.ray_constants import env_integer
+from ray.actor import ActorHandle
+from ray.data._internal.arrow_block import ArrowBlockBuilder
+from ray.data._internal.arrow_ops.transform_pyarrow import (
+    _create_empty_table,
+    hash_partition,
+)
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    ExecutionResources,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    DataOpTask,
+    MetadataOpTask,
+    OpTask,
+    _create_sub_pb,
+    estimate_total_num_of_blocks,
+)
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
+from ray.data._internal.logical.interfaces import LogicalOperator
+from ray.data._internal.stats import OpRuntimeMetrics
+from ray.data._internal.table_block import TableBlockAccessor
+from ray.data._internal.util import GiB, MiB
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+    BlockStats,
+    BlockType,
+    to_stats,
+)
+from ray.data.context import (
+    DEFAULT_MAX_HASH_SHUFFLE_AGGREGATORS,
+    DEFAULT_TARGET_MAX_BLOCK_SIZE,
+    DataContext,
+)
+
+logger = logging.getLogger(__name__)
+
+
+BlockTransformer = Callable[[Block], Block]
+
+StatefulShuffleAggregationFactory = Callable[
+    [int, List[int]], "StatefulShuffleAggregation"
+]
+
+
+DEFAULT_HASH_SHUFFLE_AGGREGATOR_MAX_CONCURRENCY = env_integer(
+    "RAY_DATA_DEFAULT_HASH_SHUFFLE_AGGREGATOR_MAX_CONCURRENCY", 8
+)
+
+DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION = env_integer(
+    "RAY_DATA_DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION", 1 * GiB
+)
+
+
+class StatefulShuffleAggregation(abc.ABC):
+    """Interface for a stateful aggregation to be used by hash-based shuffling
+    operators (inheriting from `HashShufflingOperatorBase`) and subsequent
+    aggregation of the dataset.
+
+    Any stateful aggregation has to adhere to the following protocol:
+
+        - Individual input sequence(s) will be (hash-)partitioned into N
+        partitions each.
+
+        - Accepting individual partition shards: for any given partition (identified
+        by partition-id) of the input sequence (identified by input-id) aggregation
+        will be receiving corresponding partition shards as the input sequence is
+        being shuffled.
+
+        - Upon completion of the shuffling (ie once whole sequence is shuffled)
+        aggregation will receive a call to finalize the aggregation at which point
+        it's expected to produce and return resulting block.
+
+        - After successful finalization aggregation's `clear` method will be invoked
+         to clear any accumulated state to release resources.
+    """
+
+    def __init__(self, aggregator_id: int):
+        self._aggregator_id = aggregator_id
+
+    def accept(self, input_seq_id: int, partition_id: int, partition_shard: Block):
+        """Accepts corresponding partition shard for the partition identified by
+        - Input sequence id
+        - Partition id
+        """
+        raise NotImplementedError()
+
+    def finalize(self, partition_id: int) -> Block:
+        """Finalizes aggregation of partitions (identified by partition-id)
+        from all input sequences returning resulting block.
+        """
+
+        raise NotImplementedError()
+
+    def clear(self, partition_id: int):
+        """Clears out any accumulated state for provided partition-id.
+
+        NOTE: This method is invoked after aggregation is finalized for the given
+              partition."""
+        raise NotImplementedError()
+
+
+class Concat(StatefulShuffleAggregation):
+    """Trivial aggregation recombining dataset's individual partition
+    from the partition shards provided during shuffling stage. Returns
+    single combined `Block` (Pyarrow `Table`)
+    """
+
+    def __init__(
+        self,
+        aggregator_id: int,
+        target_partition_ids: List[int],
+        *,
+        should_sort: bool,
+        key_columns: Optional[Tuple[str]] = None,
+    ):
+        super().__init__(aggregator_id)
+
+        assert (
+            not should_sort or key_columns
+        ), f"Key columns have to be specified when `should_sort=True` (got {list(key_columns)})"
+
+        self._should_sort = should_sort
+        self._key_columns = key_columns
+
+        # Block builders for individual partitions (identified by partition index)
+        self._partition_block_builders: Dict[int, ArrowBlockBuilder] = {
+            partition_id: ArrowBlockBuilder() for partition_id in target_partition_ids
+        }
+
+    def accept(self, input_seq_id: int, partition_id: int, partition_shard: Block):
+        assert input_seq_id == 0, (
+            f"Concat is unary stateful aggregation, got sequence "
+            f"index of {input_seq_id}"
+        )
+        assert partition_id in self._partition_block_builders, (
+            f"Received shard from unexpected partition '{partition_id}' "
+            f"(expecting {self._partition_block_builders.keys()})"
+        )
+
+        self._partition_block_builders[partition_id].add_block(partition_shard)
+
+    def finalize(self, partition_id: int) -> Block:
+        block = self._partition_block_builders[partition_id].build()
+
+        if self._should_sort:
+            block = block.sort_by([(k, "ascending") for k in self._key_columns])
+
+        return block
+
+    def clear(self, partition_id: int):
+        self._partition_block_builders.pop(partition_id)
+
+
+@ray.remote
+def _shuffle_block(
+    block: Block,
+    input_index: int,
+    key_columns: List[str],
+    pool: "AggregatorPool",
+    block_transformer: Optional[BlockTransformer] = None,
+    send_empty_blocks: bool = False,
+    override_partition_id: Optional[int] = None,
+) -> Tuple[BlockMetadata, Dict[int, "_PartitionStats"]]:
+    """Shuffles provided block following the algorithm:
+
+    1. Hash-partitions provided block into N partitions (where N is determined by
+       the number of receiving aggregators)
+
+    2. Individual (non-empty) partitions are subsequently submitted to respective
+       aggregators
+
+    Args:
+        block: Incoming block (in the form of Pyarrow's `Table`) to be shuffled
+        input_index: Id of the input sequence block belongs to
+        key_columns: Columns to be used by hash-partitioning algorithm
+        pool: Hash-shuffling operator's pool of aggregators that are due to receive
+              corresponding partitions (of the block)
+        send_empty_blocks: If set to true, empty blocks will NOT be filtered and
+            still be fanned out to individual aggregators to distribute schemas
+            (only known once we receive incoming block)
+        override_partition_id: Target (overridden) partition id that input block will be
+            assigned to
+        block_transformer: Block transformer that will be applied to every block prior
+            to shuffling
+
+    Returns:
+        A tuple of
+            - Metadata for the block being shuffled
+            - Map of partition ids to partition shard stats produced from the
+            shuffled block
+    """
+    stats = BlockExecStats.builder()
+    assert (len(key_columns) > 0) ^ (override_partition_id is not None), (
+        f"Either list of key columns to hash-partition by (got {key_columns} or "
+        f"target partition id override (got {override_partition_id}) must be provided!"
+    )
+
+    # Apply block transformer prior to shuffling (if any)
+    if block_transformer is not None:
+        block = block_transformer(block)
+
+    # Make sure we're handling Arrow blocks
+    block: Block = TableBlockAccessor.try_convert_block_type(
+        block, block_type=BlockType.ARROW
+    )
+
+    if block.num_rows == 0:
+        empty = BlockAccessor.for_block(block).get_metadata(exec_stats=stats.build())
+        return (empty, {})
+
+    num_partitions = pool.num_partitions
+
+    assert isinstance(block, pa.Table), f"Expected Pyarrow's `Table`, got {type(block)}"
+
+    # In case when no target key columns have been provided shuffling is
+    # reduced to just forwarding whole block to the target aggregator
+    if key_columns:
+        block_partitions = hash_partition(
+            block, hash_cols=key_columns, num_partitions=num_partitions
+        )
+    else:
+        assert (
+            0 <= override_partition_id < num_partitions
+        ), f"Expected override partition id < {num_partitions} (got {override_partition_id})"
+
+        block_partitions = {override_partition_id: block}
+
+    partition_shards_stats = {}
+    awaitable_to_partition_map = {}
+
+    for partition_id in range(num_partitions):
+        partition_shard = block_partitions.get(partition_id)
+
+        if partition_shard is None:
+            # NOTE: Hash-based shuffle operator uses empty blocks to disseminate
+            #       schema to aggregators that otherwise might not receive it,
+            #       in cases when corresponding target partition is resulting into
+            #       empty one during hash-partitioning
+            if not send_empty_blocks:
+                continue
+
+            partition_shard = _create_empty_table(block.schema)
+
+        # Capture partition shard metadata
+        #
+        # NOTE: We're skipping over empty shards as these are used for schema
+        #       broadcasting and aren't relevant to keep track of
+        if partition_shard.num_rows > 0:
+            partition_shards_stats[partition_id] = _PartitionStats.for_table(
+                partition_shard
+            )
+
+        aggregator = pool.get_aggregator_for_partition(partition_id)
+        # Put target partition shard into the Object Store to make sure partition shards
+        # are managed t/h Object Store irrespective of their size
+        partition_ref = ray.put(partition_shard)
+        # NOTE: Shuffling task is only considered completed upon target aggregator
+        #       accepting its respective partition shard
+        awaitable = aggregator.submit.remote(input_index, partition_id, partition_ref)
+
+        awaitable_to_partition_map[awaitable] = partition_id
+
+    pending_submissions = list(awaitable_to_partition_map.keys())
+
+    i = 0
+
+    # Before completing shuffling task await for all the blocks
+    # to get accepted by corresponding aggregators
+    #
+    # NOTE: This synchronization is crucial to make sure aggregations are not
+    #       getting finalized before they receive corresponding partitions
+    while len(pending_submissions) > 0:
+        ready, unready = ray.wait(
+            pending_submissions, num_returns=len(pending_submissions), timeout=1
+        )
+
+        pending_submissions = unready
+        i += 1
+
+    original_block_metadata = BlockAccessor.for_block(block).get_metadata(
+        exec_stats=stats.build()
+    )
+
+    if logger.isEnabledFor(logging.DEBUG):
+        num_rows_series, byte_sizes_series = zip(
+            *[(s.num_rows, s.byte_size) for s in partition_shards_stats.values()]
+        )
+
+        quantiles = [0, 50, 100]
+        num_rows_quantiles = np.percentile(num_rows_series, quantiles)
+        byte_sizes_quantiles = np.percentile(byte_sizes_series, quantiles)
+
+        logger.debug(
+            f"Shuffled block (rows={original_block_metadata.num_rows}, "
+            f"bytes={original_block_metadata.size_bytes/MiB:.1f}MB) "
+            f"into {len(partition_shards_stats)} partitions ("
+            f"quantiles={'/'.join(map(str, quantiles))}, "
+            f"rows={'/'.join(map(str, num_rows_quantiles))}, "
+            f"bytes={'/'.join(map(str, byte_sizes_quantiles))})"
+        )
+
+    # Return metadata for the original, shuffled block
+    return original_block_metadata, partition_shards_stats
+
+
+@dataclass
+class _PartitionStats:
+    num_rows: int
+    byte_size: int
+
+    @staticmethod
+    def from_block_metadata(block_metadata: BlockMetadata) -> "_PartitionStats":
+        return _PartitionStats(
+            num_rows=block_metadata.num_rows,
+            byte_size=block_metadata.size_bytes,
+        )
+
+    @staticmethod
+    def for_table(table: pa.Table) -> "_PartitionStats":
+        return _PartitionStats(
+            num_rows=table.num_rows,
+            byte_size=table.nbytes,
+        )
+
+    @staticmethod
+    def combine(one: "_PartitionStats", other: "_PartitionStats") -> "_PartitionStats":
+        return _PartitionStats(
+            num_rows=one.num_rows + other.num_rows,
+            byte_size=one.byte_size + other.byte_size,
+        )
+
+
+class HashShuffleProgressBarMixin(SubProgressBarMixin):
+    @property
+    @abc.abstractmethod
+    def shuffle_name(self) -> str:
+        ...
+
+    @property
+    @abc.abstractmethod
+    def reduce_name(self) -> str:
+        ...
+
+    def _validate_sub_progress_bar_names(self):
+        assert self.shuffle_name is not None, "shuffle_name should not be None"
+        assert self.reduce_name is not None, "reduce_name should not be None"
+
+    def initialize_sub_progress_bars(self, position: int) -> int:
+        """Display all sub progress bars in the termainl, and return the number of bars."""
+        self._validate_sub_progress_bar_names()
+
+        # shuffle
+        progress_bars_created = 0
+        self.shuffle_bar = None
+        self.shuffle_bar, position = _create_sub_pb(
+            self.shuffle_name, self.num_output_rows_total(), position
+        )
+        progress_bars_created += 1
+        self.shuffle_metrics = OpRuntimeMetrics(self)
+
+        # reduce
+        self.reduce_bar = None
+        self.reduce_bar, position = _create_sub_pb(
+            self.reduce_name, self.num_output_rows_total(), position
+        )
+        progress_bars_created += 1
+        self.reduce_metrics = OpRuntimeMetrics(self)
+
+        return progress_bars_created
+
+    def close_sub_progress_bars(self):
+        """Close all internal sub progress bars."""
+        self.shuffle_bar.close()
+        self.reduce_bar.close()
+
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        self._validate_sub_progress_bar_names()
+
+        # shuffle
+        self.shuffle_bar = None
+        self.shuffle_metrics = OpRuntimeMetrics(self)
+
+        # reduce
+        self.reduce_bar = None
+        self.reduce_metrics = OpRuntimeMetrics(self)
+
+        return [self.shuffle_name, self.reduce_name]
+
+    def set_sub_progress_bar(self, name, pg):
+        # No type-hints due to circular imports. `name` should be a `str`
+        # and `pg` should be a `SubProgressBar`
+        if self.shuffle_name is not None and self.shuffle_name == name:
+            self.shuffle_bar = pg
+        elif self.reduce_name is not None and self.reduce_name == name:
+            self.reduce_bar = pg
+
+
+def _derive_max_shuffle_aggregators(
+    total_cluster_resources: ExecutionResources,
+    data_context: DataContext,
+) -> int:
+    # Motivation for derivation of max # of shuffle aggregators is based on the
+    # following observations:
+    #
+    #   - Shuffle operation is necessarily a terminal operation: it terminates current
+    #     shuffle stage (set of operators that can execute concurrently)
+    #   - Shuffle operation has very low computation footprint until all preceding
+    #     operation completes (ie until shuffle finalization)
+    #   - When shuffle is finalized only shuffle operator is executing (ie it has
+    #     all of the cluster resources available at its disposal)
+    #
+    # As such we establish that the max number of shuffle
+    # aggregators (workers):
+    #
+    #   - Should not exceed total # of CPUs (to fully utilize cluster resources
+    #   while avoiding thrashing these due to over-allocation)
+    #   - Should be capped at fixed size (128 by default)
+    return min(
+        math.ceil(total_cluster_resources.cpu),
+        data_context.max_hash_shuffle_aggregators
+        or DEFAULT_MAX_HASH_SHUFFLE_AGGREGATORS,
+    )
+
+
+class HashShufflingOperatorBase(PhysicalOperator, HashShuffleProgressBarMixin):
+    """Physical operator base-class for any operators requiring hash-based
+    shuffling.
+
+    Hash-based shuffling follows standard map-reduce architecture:
+
+        1. Every incoming block is mapped (using provided `input_block_transformer`,
+           if any)
+
+        2. After mapping, every block is hash-partitioned into `num_partitions`
+           partitions and distributed (shuffled) to corresponding aggregators.
+
+        3. Aggregators perform "reducing" stage, aggregating individual partitions
+           (using configured `StatefulAggregation`), and ultimately yield resutling
+           blocks.
+
+    NOTE: This operator can perform hash-based shuffling for multiple sequences
+          simultaneously (as required by Join operator for ex).
+    """
+
+    def __init__(
+        self,
+        name_factory: Callable[[int], str],
+        input_ops: List[PhysicalOperator],
+        data_context: DataContext,
+        *,
+        key_columns: List[Tuple[str]],
+        partition_aggregation_factory: StatefulShuffleAggregationFactory,
+        num_partitions: Optional[int] = None,
+        partition_size_hint: Optional[int] = None,
+        input_block_transformer: Optional[BlockTransformer] = None,
+        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
+        shuffle_progress_bar_name: Optional[str] = None,
+        finalize_progress_bar_name: Optional[str] = None,
+    ):
+        input_logical_ops = [
+            input_physical_op._logical_operators[0] for input_physical_op in input_ops
+        ]
+
+        estimated_input_blocks = [
+            input_op.estimated_num_outputs() for input_op in input_logical_ops
+        ]
+
+        # Derive target num partitions as either of
+        #   - Requested target number of partitions
+        #   - Max estimated target number of blocks generated by the input op(s)
+        #   - Default configured hash-shuffle parallelism (200)
+        target_num_partitions: int = (
+            num_partitions
+            or (max(estimated_input_blocks) if all(estimated_input_blocks) else None)
+            or data_context.default_hash_shuffle_parallelism
+        )
+
+        super().__init__(
+            name=name_factory(target_num_partitions),
+            input_dependencies=input_ops,
+            data_context=data_context,
+        )
+
+        assert partition_size_hint is None or partition_size_hint > 0
+
+        if shuffle_progress_bar_name is None:
+            shuffle_progress_bar_name = "Shuffle"
+        if finalize_progress_bar_name is None:
+            finalize_progress_bar_name = "Reduce"
+
+        self._shuffle_name = shuffle_progress_bar_name
+        self._reduce_name = finalize_progress_bar_name
+
+        assert len(key_columns) == len(input_ops), (
+            "Each input operation has to specify matching tuple of columns used as "
+            "its hashing keys"
+        )
+
+        self._key_column_names: List[Tuple[str]] = key_columns
+        self._num_partitions: int = target_num_partitions
+
+        # Determine max number of shuffle aggregators (defaults to
+        # `DataContext.min_parallelism`)
+        total_available_cluster_resources = _get_total_cluster_resources()
+        max_shuffle_aggregators = _derive_max_shuffle_aggregators(
+            total_available_cluster_resources, data_context
+        )
+
+        # Cap number of aggregators to not exceed max configured
+        num_aggregators = min(target_num_partitions, max_shuffle_aggregators)
+
+        # Target dataset's size estimated as either of
+        #   1. ``partition_size_hint`` multiplied by target number of partitions
+        #   2. Estimation of input ops' outputs bytes
+        if partition_size_hint is not None:
+            # TODO replace with dataset-byte-size hint
+            estimated_dataset_bytes = partition_size_hint * target_num_partitions
+        else:
+            estimated_dataset_bytes = _try_estimate_output_bytes(
+                input_logical_ops,
+            )
+
+        ray_remote_args = self._get_default_aggregator_ray_remote_args(
+            num_partitions=target_num_partitions,
+            num_aggregators=num_aggregators,
+            total_available_cluster_resources=total_available_cluster_resources,
+            estimated_dataset_bytes=estimated_dataset_bytes,
+        )
+
+        if aggregator_ray_remote_args_override is not None:
+            # Set default values missing for configs missing in the override
+            ray_remote_args.update(aggregator_ray_remote_args_override)
+
+        self._aggregator_pool: AggregatorPool = AggregatorPool(
+            num_partitions=target_num_partitions,
+            num_aggregators=num_aggregators,
+            aggregation_factory=partition_aggregation_factory,
+            aggregator_ray_remote_args=ray_remote_args,
+            data_context=data_context,
+        )
+
+        # We track the running usage total because iterating
+        # and summing over all shuffling tasks can be expensive
+        # if the # of shuffling tasks is large
+        self._shuffling_resource_usage = ExecutionResources.zero()
+
+        self._input_block_transformer = input_block_transformer
+
+        self._next_shuffle_tasks_idx: int = 0
+        # Shuffling tasks are mapped like following
+        #   - Input sequence id -> Task id -> Task
+        #
+        # NOTE: Input sequences correspond to the outputs of the input operators
+        self._shuffling_tasks: DefaultDict[
+            int, Dict[int, MetadataOpTask]
+        ] = defaultdict(dict)
+
+        self._next_aggregate_task_idx: int = 0
+        # Aggregating tasks are mapped like following
+        #   - Task id -> Task
+        #
+        # NOTE: Aggregating tasks are invariant of the # of input operators, as
+        #       aggregation is assumed to always produce a single sequence
+        self._finalizing_tasks: Dict[int, DataOpTask] = dict()
+
+        # This is a workaround to be able to distribute schemas to individual
+        # aggregators (keeps track which input sequences have already broadcasted
+        # their schemas)
+        self._has_schemas_broadcasted: DefaultDict[int, bool] = defaultdict(bool)
+        # Set of partitions still pending finalization
+        self._pending_finalization_partition_ids: Set[int] = set(
+            range(target_num_partitions)
+        )
+
+        self._output_queue: Deque[RefBundle] = deque()
+
+        self._output_blocks_stats: List[BlockStats] = list()
+        self._shuffled_blocks_stats: List[BlockStats] = list()
+
+        # Incremental individual partition metadata accumulated separately for
+        # individual input sequences during shuffling. Maps
+        #
+        #   input_sequence_id -> partition_id -> _PartitionStats
+        #
+        self._partitions_stats: DefaultDict[
+            int, Dict[int, _PartitionStats]
+        ] = defaultdict(dict)
+
+        self._health_monitoring_started: bool = False
+        self._health_monitoring_start_time: float = 0.0
+        self._pending_aggregators_refs: Optional[List[ObjectRef[ActorHandle]]] = None
+
+    def start(self, options: ExecutionOptions) -> None:
+        super().start(options)
+
+        self._aggregator_pool.start()
+
+    @property
+    def shuffle_name(self) -> str:
+        return self._shuffle_name
+
+    @property
+    def reduce_name(self) -> str:
+        return self._reduce_name
+
+    def _add_input_inner(self, input_bundle: RefBundle, input_index: int) -> None:
+
+        # TODO move to base class
+        self.shuffle_metrics.on_input_received(input_bundle)
+        self._do_add_input_inner(input_bundle, input_index)
+
+    def _do_add_input_inner(self, input_bundle: RefBundle, input_index: int):
+        input_blocks_refs: List[ObjectRef[Block]] = input_bundle.block_refs
+        input_blocks_metadata: List[BlockMetadata] = input_bundle.metadata
+
+        for block_ref, block_metadata in zip(input_blocks_refs, input_blocks_metadata):
+            # If operator hasn't propagated schemas (for this sequence) to its
+            # aggregator pool, it will need to do that upon receiving of the
+            # first block
+            should_broadcast_schemas = not self._has_schemas_broadcasted[input_index]
+            input_key_column_names = self._key_column_names[input_index]
+            # Compose shuffling task resource bundle
+            shuffle_task_resource_bundle = {
+                "num_cpus": 0.5,
+                "memory": self._estimate_shuffling_memory_req(
+                    block_metadata,
+                    target_max_block_size=(
+                        self._data_context.target_max_block_size
+                        or DEFAULT_TARGET_MAX_BLOCK_SIZE
+                    ),
+                ),
+            }
+
+            cur_shuffle_task_idx = self._next_shuffle_tasks_idx
+            self._next_shuffle_tasks_idx += 1
+
+            # NOTE: In cases when NO key-columns are provided for hash-partitioning
+            #       to be performed on (legitimate scenario for global aggregations),
+            #       shuffling is essentially reduced to round-robin'ing of the blocks
+            #       among the aggregators
+            override_partition_id = (
+                cur_shuffle_task_idx % self._num_partitions
+                if not input_key_column_names
+                else None
+            )
+
+            # Fan out provided input blocks to "shuffle" it
+            #   - Block is first hash-partitioned into N partitions
+            #   - Individual partitions then are submitted to the corresponding
+            #     aggregators
+            input_block_partition_shards_metadata_tuple_ref: ObjectRef[
+                Tuple[BlockMetadata, Dict[int, _PartitionStats]]
+            ] = _shuffle_block.options(
+                **shuffle_task_resource_bundle,
+                num_returns=1,
+                # Make sure tasks are retried indefinitely
+                max_retries=-1,
+            ).remote(
+                block_ref,
+                input_index,
+                input_key_column_names,
+                self._aggregator_pool,
+                block_transformer=self._input_block_transformer,
+                send_empty_blocks=should_broadcast_schemas,
+                override_partition_id=override_partition_id,
+            )
+
+            if should_broadcast_schemas:
+                self._has_schemas_broadcasted[input_index] = True
+
+            def _on_partitioning_done(cur_shuffle_task_idx: int):
+                task = self._shuffling_tasks[input_index].pop(cur_shuffle_task_idx)
+                self._shuffling_resource_usage = (
+                    self._shuffling_resource_usage.subtract(
+                        task.get_requested_resource_bundle()
+                    )
+                )
+                # Fetch input block and resulting partition shards block metadata and
+                # handle obtained metadata
+                #
+                # NOTE: We set timeout equal to 1m here as an upper-bound to make
+                #       sure that `ray.get(...)` invocation couldn't stall the pipeline
+                #       indefinitely
+                input_block_metadata, partition_shards_stats = ray.get(
+                    task.get_waitable(), timeout=60
+                )
+
+                self._handle_shuffled_block_metadata(
+                    input_index, input_block_metadata, partition_shards_stats
+                )
+
+                # Update Shuffle metrics on task output generated
+                blocks = [(task.get_waitable(), input_block_metadata)]
+                # NOTE: schema doesn't matter because we are creating a ref bundle
+                # for metrics recording purposes
+                out_bundle = RefBundle(blocks, schema=None, owns_blocks=False)
+                self.shuffle_metrics.on_output_taken(input_bundle)
+                self.shuffle_metrics.on_task_output_generated(
+                    cur_shuffle_task_idx, out_bundle
+                )
+                self.shuffle_metrics.on_task_finished(cur_shuffle_task_idx, None)
+
+                # Update Shuffle progress bar
+                self.shuffle_bar.update(increment=input_block_metadata.num_rows or 0)
+
+            # TODO update metrics
+            task = self._shuffling_tasks[input_index][
+                cur_shuffle_task_idx
+            ] = MetadataOpTask(
+                task_index=cur_shuffle_task_idx,
+                object_ref=input_block_partition_shards_metadata_tuple_ref,
+                task_done_callback=functools.partial(
+                    _on_partitioning_done, cur_shuffle_task_idx
+                ),
+                task_resource_bundle=ExecutionResources.from_resource_dict(
+                    shuffle_task_resource_bundle
+                ),
+            )
+            if task.get_requested_resource_bundle() is not None:
+                self._shuffling_resource_usage = self._shuffling_resource_usage.add(
+                    task.get_requested_resource_bundle()
+                )
+
+            #  Update Shuffle Metrics on task submission
+            self.shuffle_metrics.on_task_submitted(
+                cur_shuffle_task_idx,
+                RefBundle(
+                    [(block_ref, block_metadata)], schema=None, owns_blocks=False
+                ),
+            )
+
+            # Update Shuffle progress bar
+            _, _, num_rows = estimate_total_num_of_blocks(
+                cur_shuffle_task_idx + 1,
+                self.upstream_op_num_outputs(),
+                self.shuffle_metrics,
+                total_num_tasks=None,
+            )
+            self.shuffle_bar.update(total=num_rows)
+
+    def has_next(self) -> bool:
+        self._try_finalize()
+        return len(self._output_queue) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        bundle: RefBundle = self._output_queue.popleft()
+
+        # TODO move to base class
+        self.reduce_metrics.on_output_dequeued(bundle)
+        self.reduce_metrics.on_output_taken(bundle)
+
+        self._output_blocks_stats.extend(to_stats(bundle.metadata))
+
+        return bundle
+
+    def get_active_tasks(self) -> List[OpTask]:
+        # Collect shuffling tasks for every input sequence
+        shuffling_tasks = self._get_active_shuffling_tasks()
+
+        # Collect aggregating tasks for every input sequence
+        finalizing_tasks: List[DataOpTask] = list(self._finalizing_tasks.values())
+
+        return shuffling_tasks + finalizing_tasks
+
+    def _get_active_shuffling_tasks(self) -> List[MetadataOpTask]:
+        return list(
+            itertools.chain.from_iterable(
+                [
+                    input_shuffling_task_map.values()
+                    for input_shuffling_task_map in self._shuffling_tasks.values()
+                ]
+            )
+        )
+
+    def _is_shuffling_done(self):
+        return self._inputs_complete and all(
+            [
+                len(self._shuffling_tasks[input_seq_idx]) == 0
+                for input_seq_idx in range(len(self._input_dependencies))
+            ]
+        )
+
+    def _try_finalize(self):
+        # Skip if finalization of all partitions had been already scheduled
+        if self._is_finalized():
+            return
+
+        # Finalization can only proceed once
+        #   - All input sequences have been ingested
+        #   - All outstanding shuffling tasks have completed
+        if not self._is_shuffling_done():
+            return
+
+        def _on_bundle_ready(partition_id: int, bundle: RefBundle):
+            # Add finalized block to the output queue
+            self._output_queue.append(bundle)
+
+            # Update Finalize Metrics on task output generated
+            self.reduce_metrics.on_output_queued(bundle)
+            self.reduce_metrics.on_task_output_generated(
+                task_index=partition_id, output=bundle
+            )
+            self.reduce_metrics.on_task_finished(
+                task_index=partition_id, exception=None
+            )
+            _, num_outputs, num_rows = estimate_total_num_of_blocks(
+                partition_id + 1,
+                self.upstream_op_num_outputs(),
+                self.reduce_metrics,
+                total_num_tasks=self._num_partitions,
+            )
+            self._estimated_num_output_bundles = num_outputs
+            self._estimated_output_num_rows = num_rows
+
+            # Update Finalize progress bar
+            self.reduce_bar.update(
+                increment=bundle.num_rows() or 0, total=self.num_output_rows_total()
+            )
+
+        def _on_aggregation_done(partition_id: int, exc: Optional[Exception]):
+            if partition_id in self._finalizing_tasks:
+                self._finalizing_tasks.pop(partition_id)
+
+                if exc:
+                    logger.error(
+                        f"Aggregation of the {partition_id} partition "
+                        f"failed with: {exc}",
+                        exc_info=exc,
+                    )
+
+        # NOTE: Unless explicitly set finalization batch size defaults to the #
+        #       of shuffle aggregators
+        max_batch_size = (
+            self.data_context.max_hash_shuffle_finalization_batch_size
+            or self._aggregator_pool.num_aggregators
+        )
+
+        num_running_finalizing_tasks = len(self._finalizing_tasks)
+        num_remaining_partitions = len(self._pending_finalization_partition_ids)
+
+        # Finalization is executed in batches of no more than
+        # `DataContext.max_hash_shuffle_finalization_batch_size` tasks at a time.
+        #
+        # Batch size is used as a lever to limit memory pressure on the nodes
+        # where aggregators are run by limiting # of finalization tasks running
+        # concurrently
+        next_batch_size = min(
+            num_remaining_partitions,
+            max_batch_size - num_running_finalizing_tasks,
+        )
+
+        assert next_batch_size >= 0, (
+            f"Finalization batch size must be greater than 0 "
+            f"(got {next_batch_size}; "
+            f"remaining={num_remaining_partitions}, "
+            f"finalizing={num_running_finalizing_tasks}, "
+            f"max_batch_size={max_batch_size})"
+        )
+
+        if next_batch_size == 0:
+            return
+
+        # We're sampling randomly next set of partitions to be finalized
+        # to distribute finalization window uniformly across the nodes of the cluster
+        # and avoid effect of "sliding lense" effect where we finalize the batch of
+        # N *adjacent* partitions that may be co-located on the same node:
+        #
+        #   - Adjacent partitions i and i+1 are handled by adjacent
+        #   aggregators (since membership is determined as i % num_aggregators)
+        #
+        #   - Adjacent aggregators have high likelihood of running on the
+        #   same node (when num aggregators > num nodes)
+        #
+        # NOTE: This doesn't affect determinism, since this only impacts order
+        #       of finalization (hence not required to be seeded)
+        target_partition_ids = random.sample(
+            list(self._pending_finalization_partition_ids), next_batch_size
+        )
+
+        logger.debug(
+            f"Scheduling partitions {target_partition_ids} for finalization: "
+            f"{[self._get_partition_stats(pid) for pid in target_partition_ids]}"
+        )
+
+        for partition_id in target_partition_ids:
+            aggregator = self._aggregator_pool.get_aggregator_for_partition(
+                partition_id
+            )
+
+            # Estimate (heap) memory requirement to execute finalization task
+            # Compose shuffling task resource bundle
+            finalize_task_resource_bundle = {
+                # TODO currently not possible to specify the resources for an actor
+                # "memory": self._estimate_finalization_memory_req(partition_id),
+            }
+
+            # Request finalization of the partition
+            block_gen = aggregator.finalize.options(
+                **finalize_task_resource_bundle,
+            ).remote(partition_id)
+
+            self._finalizing_tasks[partition_id] = DataOpTask(
+                task_index=partition_id,
+                streaming_gen=block_gen,
+                output_ready_callback=functools.partial(_on_bundle_ready, partition_id),
+                task_done_callback=functools.partial(
+                    _on_aggregation_done, partition_id
+                ),
+                task_resource_bundle=(
+                    ExecutionResources.from_resource_dict(finalize_task_resource_bundle)
+                ),
+            )
+
+            # Pop partition id from remaining set
+            self._pending_finalization_partition_ids.remove(partition_id)
+
+            # Update Finalize Metrics on task submission
+            # NOTE: This is empty because the input is directly forwarded from the
+            # output of the shuffling stage, which we don't return.
+            empty_bundle = RefBundle([], schema=None, owns_blocks=False)
+            self.reduce_metrics.on_task_submitted(partition_id, empty_bundle)
+
+    def _do_shutdown(self, force: bool = False) -> None:
+        self._aggregator_pool.shutdown(force=True)
+        # NOTE: It's critical for Actor Pool to release actors before calling into
+        #       the base method that will attempt to cancel and join pending.
+        super()._do_shutdown(force)
+        # Release any pending refs
+        self._shuffling_tasks.clear()
+        self._finalizing_tasks.clear()
+
+    def _extra_metrics(self):
+        shuffle_name = f"{self._name}_shuffle"
+        finalize_name = f"{self._name}_finalize"
+
+        self.shuffle_metrics.as_dict()
+
+        return {
+            shuffle_name: self.shuffle_metrics.as_dict(),
+            finalize_name: self.reduce_metrics.as_dict(),
+        }
+
+    def get_stats(self):
+        shuffle_name = f"{self._name}_shuffle"
+        reduce_name = f"{self._name}_finalize"
+        return {
+            shuffle_name: self._shuffled_blocks_stats,
+            reduce_name: self._output_blocks_stats,
+        }
+
+    def current_processor_usage(self) -> ExecutionResources:
+        # Current processors resource usage is comprised by
+        #   - Base Aggregator actors resource utilization (captured by
+        #     `base_resource_usage` method)
+        #   - Active shuffling tasks
+        #   - Active finalizing tasks (actor tasks)
+        base_usage = self.base_resource_usage
+        running_usage = self._shuffling_resource_usage
+
+        # TODO add memory to resources being tracked
+        return base_usage.add(running_usage)
+
+    @property
+    def base_resource_usage(self) -> ExecutionResources:
+        # TODO add memory to resources being tracked
+        return ExecutionResources(
+            cpu=(
+                self._aggregator_pool.num_aggregators
+                * self._aggregator_pool._aggregator_ray_remote_args["num_cpus"]
+            ),
+            object_store_memory=0,
+            gpu=0,
+        )
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        return ExecutionResources(
+            # TODO fix (this hack is currently to force Ray to spin up more tasks when
+            #      shuffling to autoscale hardware capacity)
+            cpu=0.01,
+            # cpu=self._shuffle_block_ray_remote_args.get("num_cpus", 0),
+            # TODO estimate (twice avg block size)
+            object_store_memory=0,
+            gpu=0,
+        )
+
+    def completed(self) -> bool:
+        # TODO separate marking as completed from the check
+        return self._is_finalized() and super().completed()
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        return True
+
+    def _is_finalized(self):
+        return len(self._pending_finalization_partition_ids) == 0
+
+    def _handle_shuffled_block_metadata(
+        self,
+        input_seq_id: int,
+        input_block_metadata: BlockMetadata,
+        partition_shards_stats: Dict[int, _PartitionStats],
+    ):
+        # Keep track of the progress of shuffling incoming blocks
+        self._shuffled_blocks_stats.append(input_block_metadata.to_stats())
+
+        # Update incremental input sequence partitions metadata
+        for partition_id, new_partition_shard_stats in partition_shards_stats.items():
+            current_partition_stats: Optional[_PartitionStats] = self._partitions_stats[
+                input_seq_id
+            ].get(partition_id, None)
+
+            self._partitions_stats[input_seq_id][partition_id] = (
+                _PartitionStats.combine(
+                    current_partition_stats,
+                    new_partition_shard_stats,
+                )
+                if current_partition_stats
+                else new_partition_shard_stats
+            )
+
+    def _get_partition_stats(
+        self, partition_id: int
+    ) -> Dict[int, Optional[_PartitionStats]]:
+        return {
+            # NOTE: Some partitions might be empty (and hence missing) in some sequences
+            input_seq_id: partition_stats_map.get(partition_id)
+            for input_seq_id, partition_stats_map in self._partitions_stats.items()
+        }
+
+    @classmethod
+    def _estimate_shuffling_memory_req(
+        cls,
+        block_metadata: BlockMetadata,
+        target_max_block_size: int,
+    ):
+        estimated_block_bytes = (
+            block_metadata.size_bytes
+            if block_metadata.size_bytes is not None
+            else target_max_block_size
+        )
+
+        return estimated_block_bytes * 2
+
+    def _get_default_aggregator_ray_remote_args(
+        self,
+        *,
+        num_partitions: int,
+        num_aggregators: int,
+        total_available_cluster_resources: ExecutionResources,
+        estimated_dataset_bytes: Optional[int],
+    ):
+        assert num_partitions >= num_aggregators
+
+        if estimated_dataset_bytes is not None:
+            estimated_aggregator_memory_required = self._estimate_aggregator_memory_allocation(
+                num_aggregators=num_aggregators,
+                num_partitions=num_partitions,
+                # NOTE: If no partition size hint is provided we simply assume target
+                #       max block size specified as the best partition size estimate
+                estimated_dataset_bytes=estimated_dataset_bytes,
+            )
+        else:
+            # NOTE: In cases when we're unable to estimate dataset size,
+            #       we simply fallback to request the minimum of:
+            #       - conservative 50% of total available memory for a join operation.
+            #       - ``DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION`` worth of
+            #       memory for every Aggregator.
+
+            max_memory_per_aggregator = (
+                total_available_cluster_resources.memory / num_aggregators
+            )
+            modest_memory_per_aggregator = max_memory_per_aggregator / 2
+
+            estimated_aggregator_memory_required = min(
+                modest_memory_per_aggregator,
+                DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION,
+            )
+
+        remote_args = {
+            "num_cpus": self._get_aggregator_num_cpus(
+                total_available_cluster_resources,
+                estimated_aggregator_memory_required,
+                num_aggregators=num_aggregators,
+            ),
+            "memory": estimated_aggregator_memory_required,
+            # NOTE: By default aggregating actors should be spread across available
+            #       nodes to prevent any single node being overloaded with a "thundering
+            #       herd"
+            "scheduling_strategy": "SPREAD",
+            # Allow actor tasks to execute out of order by default to prevent head-of-line
+            # blocking scenario.
+            "allow_out_of_order_execution": True,
+        }
+
+        return remote_args
+
+    @abc.abstractmethod
+    def _get_operator_num_cpus_override(self) -> int:
+        pass
+
+    def _get_aggregator_num_cpus(
+        self,
+        total_available_cluster_resources: ExecutionResources,
+        estimated_aggregator_memory_required: int,
+        num_aggregators: int,
+    ) -> float:
+        """Estimates number of CPU resources to be provisioned for individual
+        Aggregators.
+
+        Due to semantic of the Aggregator's role (outlined below), their CPU
+        allocation is mostly playing a role of complimenting their memory allocation
+        such that it serves as a protection mechanism from over-allocation of the
+        tasks that do not specify their respective memory resources.
+        """
+
+        # First, check whether there is an override
+        if self._get_operator_num_cpus_override() is not None:
+            return self._get_operator_num_cpus_override()
+
+        # Note that
+        #
+        #  - Shuffle aggregators have modest computational footprint until
+        #    finalization stage
+        #  - Finalization stage actually always executes standalone, since it only
+        #    starts when all preceding operations complete
+        #
+        # Though we don't need to purposefully allocate any meaningful amount of
+        # CPU resources to the shuffle aggregators, we're still allocating nominal
+        # CPU resources to it such that to compliment its required memory allocation
+        # and therefore protect from potential OOMs in case other tasks getting
+        # scheduled onto the same node, but not specifying their respective memory
+        # requirements.
+        #
+        # CPU allocation is determined like following
+        #
+        #   CPUs = Total memory required / 4 GiB (standard ratio in the conventional clouds)
+        #
+        # But no more than
+        #   - 25% of total available CPUs but
+        #   - No more than 4 CPUs per aggregator
+        #
+        cap = min(4.0, total_available_cluster_resources.cpu * 0.25 / num_aggregators)
+
+        target_num_cpus = min(cap, estimated_aggregator_memory_required / (4 * GiB))
+
+        # Round resource to 2d decimal point (for readability)
+        return round(target_num_cpus, 2)
+
+    @classmethod
+    def _estimate_aggregator_memory_allocation(
+        cls,
+        *,
+        num_aggregators: int,
+        num_partitions: int,
+        estimated_dataset_bytes: int,
+    ) -> int:
+        raise NotImplementedError()
+
+    @classmethod
+    def _gen_op_name(cls, num_partitions: int) -> str:
+        raise NotImplementedError()
+
+
+class HashShuffleOperator(HashShufflingOperatorBase):
+    def __init__(
+        self,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        *,
+        key_columns: Tuple[str],
+        num_partitions: Optional[int] = None,
+        should_sort: bool = False,
+        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            name_factory=(
+                lambda num_partitions: f"Shuffle(key_columns={key_columns}, num_partitions={num_partitions})"
+            ),
+            input_ops=[input_op],
+            data_context=data_context,
+            key_columns=[key_columns],
+            num_partitions=num_partitions,
+            aggregator_ray_remote_args_override=aggregator_ray_remote_args_override,
+            partition_aggregation_factory=(
+                lambda aggregator_id, target_partition_ids: Concat(
+                    aggregator_id,
+                    target_partition_ids,
+                    should_sort=should_sort,
+                    key_columns=key_columns,
+                )
+            ),
+            shuffle_progress_bar_name="Shuffle",
+        )
+
+    def _get_operator_num_cpus_override(self) -> float:
+        return self.data_context.hash_shuffle_operator_actor_num_cpus_override
+
+    @classmethod
+    def _estimate_aggregator_memory_allocation(
+        cls,
+        *,
+        num_aggregators: int,
+        num_partitions: int,
+        estimated_dataset_bytes: int,
+    ) -> int:
+        partition_byte_size_estimate = math.ceil(
+            estimated_dataset_bytes / num_partitions
+        )
+
+        # Estimate of object store memory required to accommodate all partitions
+        # handled by a single aggregator
+        aggregator_shuffle_object_store_memory_required: int = math.ceil(
+            estimated_dataset_bytes / num_aggregators
+        )
+        # Estimate of memory required to accommodate single partition as an output
+        # (inside Object Store)
+        output_object_store_memory_required: int = partition_byte_size_estimate
+
+        aggregator_total_memory_required: int = (
+            # Inputs (object store)
+            aggregator_shuffle_object_store_memory_required
+            +
+            # Output (object store)
+            output_object_store_memory_required
+        )
+
+        logger.info(
+            f"Estimated memory requirement for shuffling aggregator "
+            f"(partitions={num_partitions}, "
+            f"aggregators={num_aggregators}, "
+            f"dataset (estimate)={estimated_dataset_bytes / GiB:.1f}GiB): "
+            f"shuffle={aggregator_shuffle_object_store_memory_required / MiB:.1f}MiB, "
+            f"output={output_object_store_memory_required / MiB:.1f}MiB, "
+            f"total={aggregator_total_memory_required / MiB:.1f}MiB, "
+        )
+
+        return aggregator_total_memory_required
+
+
+@dataclass
+class AggregatorHealthInfo:
+    """Health information about aggregators for issue detection."""
+
+    started_at: float
+    ready_aggregators: int
+    total_aggregators: int
+    has_unready_aggregators: bool
+    wait_time: float
+    required_resources: ExecutionResources
+
+
+class AggregatorPool:
+    def __init__(
+        self,
+        num_partitions: int,
+        num_aggregators: int,
+        aggregation_factory: StatefulShuffleAggregationFactory,
+        aggregator_ray_remote_args: Dict[str, Any],
+        data_context: DataContext,
+    ):
+        assert (
+            num_partitions >= 1
+        ), f"Number of partitions has to be >= 1 (got {num_partitions})"
+
+        self._data_context = data_context
+        self._num_partitions = num_partitions
+        self._num_aggregators: int = num_aggregators
+        self._aggregator_partition_map: Dict[
+            int, List[int]
+        ] = self._allocate_partitions(
+            num_partitions=num_partitions,
+        )
+
+        self._aggregators: List[ray.actor.ActorHandle] = []
+
+        self._aggregation_factory_ref: ObjectRef[
+            StatefulShuffleAggregationFactory
+        ] = ray.put(aggregation_factory)
+
+        self._aggregator_ray_remote_args: Dict[
+            str, Any
+        ] = self._derive_final_shuffle_aggregator_ray_remote_args(
+            aggregator_ray_remote_args,
+            self._aggregator_partition_map,
+        )
+
+    def start(self):
+        # Check cluster resources before starting aggregators
+        self._check_cluster_resources()
+
+        logger.debug(
+            f"Starting {self._num_aggregators} shuffle aggregators with remote "
+            f"args: {self._aggregator_ray_remote_args}"
+        )
+
+        for aggregator_id in range(self._num_aggregators):
+            target_partition_ids = self._aggregator_partition_map[aggregator_id]
+
+            assert len(target_partition_ids) > 0
+
+            aggregator = HashShuffleAggregator.options(
+                **self._aggregator_ray_remote_args
+            ).remote(aggregator_id, target_partition_ids, self._aggregation_factory_ref)
+
+            self._aggregators.append(aggregator)
+
+        # Start issue detector actor
+        self.start_health_monitoring()
+
+    def _check_cluster_resources(self) -> None:
+        """Check if cluster has enough resources to schedule all aggregators.
+        Raises:
+            ValueError: If cluster doesn't have sufficient resources.
+        """
+        try:
+            cluster_resources = ray.cluster_resources()
+            available_resources = ray.available_resources()
+        except Exception as e:
+            logger.warning(f"Failed to get cluster resources: {e}")
+            return
+
+        # Calculate required resources for all aggregators
+        required_cpus = (
+            self._aggregator_ray_remote_args.get("num_cpus", 1) * self._num_aggregators
+        )
+        required_memory = (
+            self._aggregator_ray_remote_args.get("memory", 0) * self._num_aggregators
+        )
+
+        # Check CPU resources
+        total_cpus = cluster_resources.get("CPU", 0)
+        available_cpus = available_resources.get("CPU", 0)
+
+        if required_cpus > total_cpus:
+            logger.warning(
+                f"Insufficient CPU resources in cluster for hash shuffle operation. "
+                f"Required: {required_cpus} CPUs for {self._num_aggregators} aggregators, "
+                f"but cluster only has {total_cpus} total CPUs. "
+                f"Consider either increasing the cluster size or reducing the number of aggregators via `DataContext.max_hash_shuffle_aggregators`."
+            )
+
+        if required_cpus > available_cpus:
+            logger.warning(
+                f"Limited available CPU resources for hash shuffle operation. "
+                f"Required: {required_cpus} CPUs, available: {available_cpus} CPUs. "
+                f"Aggregators may take longer to start due to contention for resources."
+            )
+
+        # Check memory resources if specified
+        if required_memory > 0:
+            total_memory = cluster_resources.get("memory", 0)
+            available_memory = available_resources.get("memory", 0)
+
+            if required_memory > total_memory:
+                logger.warning(
+                    f"Insufficient memory resources in cluster for hash shuffle operation. "
+                    f"Required: {required_memory / GiB:.1f} GiB for {self._num_aggregators} aggregators, "
+                    f"but cluster only has {total_memory / GiB:.1f} GiB total memory. "
+                    f"Consider reducing the number of partitions or increasing cluster size."
+                )
+
+            if required_memory > available_memory:
+                logger.warning(
+                    f"Limited available memory resources for hash shuffle operation. "
+                    f"Required: {required_memory / GiB:.1f} GiB, available: {available_memory / GiB:.1f} GiB. "
+                    f"Aggregators may take longer to start due to resource contention."
+                )
+
+            logger.debug(
+                f"Resource check passed for hash shuffle operation: "
+                f"required CPUs={required_cpus}, available CPUs={available_cpus}, "
+                f"required memory={required_memory / GiB:.1f} GiB, available memory={available_memory / GiB:.1f} GiB"
+            )
+
+    @property
+    def num_partitions(self):
+        return self._num_partitions
+
+    @property
+    def num_aggregators(self):
+        return self._num_aggregators
+
+    def get_aggregator_for_partition(self, partition_id: int) -> ActorHandle:
+        return self._aggregators[self._get_aggregator_id_for_partition(partition_id)]
+
+    def _allocate_partitions(self, *, num_partitions: int):
+        assert num_partitions >= self._num_aggregators
+
+        aggregator_to_partition_map: DefaultDict[int, List[int]] = defaultdict(list)
+
+        for partition_id in range(num_partitions):
+            aggregator_id = self._get_aggregator_id_for_partition(partition_id)
+            aggregator_to_partition_map[aggregator_id].append(partition_id)
+
+        return aggregator_to_partition_map
+
+    def _get_aggregator_id_for_partition(self, partition_id: int) -> int:
+        assert partition_id < self._num_partitions
+
+        return partition_id % self._num_aggregators
+
+    @staticmethod
+    def _derive_final_shuffle_aggregator_ray_remote_args(
+        aggregator_ray_remote_args: Dict[str, Any],
+        aggregator_partition_map: Dict[int, List[int]],
+    ):
+        max_partitions_per_aggregator = max(
+            [len(ps) for ps in aggregator_partition_map.values()]
+        )
+
+        # Cap shuffle aggregator concurrency at the smaller of
+        #   - Max number of partitions per aggregator
+        #   - Threshold (8 by default)
+        max_concurrency = min(
+            max_partitions_per_aggregator,
+            DEFAULT_HASH_SHUFFLE_AGGREGATOR_MAX_CONCURRENCY,
+        )
+
+        assert (
+            max_concurrency >= 1
+        ), f"{max_partitions_per_aggregator=}, {DEFAULT_MAX_HASH_SHUFFLE_AGGREGATORS}"
+
+        # NOTE: ShuffleAggregator is configured as threaded actor to allow for
+        #       multiple requests to be handled "concurrently" (par GIL) --
+        #       while it's not a real concurrency in its fullest of senses, having
+        #       multiple Python threads allows as to parallelize all activities not
+        #       requiring the GIL (inside Ray Core) such that concurrent request
+        #       handling tasks are only blocked on GIL and are ready to execute as
+        #       soon as it's released.
+        finalized_remote_args = {
+            "max_concurrency": max_concurrency,
+            **aggregator_ray_remote_args,
+        }
+
+        return finalized_remote_args
+
+    def shutdown(self, force: bool):
+        # Shutdown aggregators
+        if force:
+            for actor in self._aggregators:
+                # NOTE: Actors can't be brought back after being ``ray.kill``-ed,
+                #       hence we're only doing that if this is a forced release
+                ray.kill(actor)
+
+        self._aggregators.clear()
+
+    def check_aggregator_health(self) -> Optional[AggregatorHealthInfo]:
+        """Get health information about aggregators for issue detection.
+
+        Returns:
+            AggregatorHealthInfo with health info or None if monitoring hasn't started.
+        """
+        if not self._health_monitoring_started:
+            return None
+
+        if self._pending_aggregators_refs is None:
+            # Initialize readiness refs
+            self._pending_aggregators_refs = [
+                aggregator.__ray_ready__.remote() for aggregator in self._aggregators
+            ]
+
+        # Use ray.wait to check readiness in non-blocking fashion
+        _, unready_refs = ray.wait(
+            self._pending_aggregators_refs,
+            num_returns=len(self._pending_aggregators_refs),
+            timeout=0,  # Non-blocking
+        )
+
+        # Update readiness refs to only track the unready ones
+        self._pending_aggregators_refs = unready_refs
+
+        current_time = time.time()
+        ready_aggregators = self._num_aggregators - len(unready_refs)
+        required_cpus = (
+            self._aggregator_ray_remote_args.get("num_cpus", 1) * self._num_aggregators
+        )
+        required_memory = (
+            self._aggregator_ray_remote_args.get("memory", 0) * self._num_aggregators
+        )
+
+        return AggregatorHealthInfo(
+            started_at=self._health_monitoring_start_time,
+            ready_aggregators=ready_aggregators,
+            total_aggregators=self._num_aggregators,
+            has_unready_aggregators=len(unready_refs) > 0,
+            wait_time=current_time - self._health_monitoring_start_time,
+            required_resources=ExecutionResources(
+                cpu=required_cpus, memory=required_memory
+            ),
+        )
+
+    def start_health_monitoring(self):
+        """Start health monitoring (without separate actor)."""
+        self._health_monitoring_started = True
+        self._health_monitoring_start_time = time.time()
+        self._pending_aggregators_refs = None
+
+
+@ray.remote(
+    # Make sure tasks are retried indefinitely
+    max_task_retries=-1
+)
+class HashShuffleAggregator:
+    """Actor handling of the assigned partitions during hash-shuffle operation
+
+    NOTE: This actor might have ``max_concurrency`` > 1 (depending on the number of
+          assigned partitions, and has to be thread-safe!
+    """
+
+    def __init__(
+        self,
+        aggregator_id: int,
+        target_partition_ids: List[int],
+        agg_factory: StatefulShuffleAggregationFactory,
+    ):
+        self._lock = threading.Lock()
+        self._agg: StatefulShuffleAggregation = agg_factory(
+            aggregator_id, target_partition_ids
+        )
+
+    def submit(self, input_seq_id: int, partition_id: int, partition_shard: Block):
+        with self._lock:
+            self._agg.accept(input_seq_id, partition_id, partition_shard)
+
+    def finalize(
+        self, partition_id: int
+    ) -> AsyncGenerator[Union[Block, "BlockMetadataWithSchema"], None]:
+        with self._lock:
+            # Finalize given partition id
+            exec_stats_builder = BlockExecStats.builder()
+            block = self._agg.finalize(partition_id)
+            exec_stats = exec_stats_builder.build()
+            # Clear any remaining state (to release resources)
+            self._agg.clear(partition_id)
+
+        # TODO break down blocks to target size
+        yield block
+        yield BlockMetadataWithSchema.from_block(block, stats=exec_stats)
+
+
+def _get_total_cluster_resources() -> ExecutionResources:
+    """Retrieves total available cluster resources:
+
+    1. If AutoscalerV2 is used, then corresponding max configured resources of
+        the corresponding `ClusterConfig` is returned.
+    2. In case `ClusterConfig` is not set then falls back to currently available
+        cluster resources (retrieved by `ray.cluster_resources()`)
+
+    """
+    return ExecutionResources.from_resource_dict(
+        ray._private.state.state.get_max_resources_from_cluster_config()
+        or ray.cluster_resources()
+    )
+
+
+# TODO rebase on generic operator output estimation
+def _try_estimate_output_bytes(
+    input_logical_ops: List[LogicalOperator],
+) -> Optional[int]:
+    inferred_op_output_bytes = [
+        op.infer_metadata().size_bytes for op in input_logical_ops
+    ]
+
+    # Return sum of input ops estimated output byte sizes,
+    # if all are well defined
+    if all(nbs is not None for nbs in inferred_op_output_bytes):
+        return sum(inferred_op_output_bytes)
+
+    return None
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/input_data_buffer.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/input_data_buffer.py
new file mode 100644
index 0000000000000000000000000000000000000000..552eced217836985b7a7129c94e657c09a498b7f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/input_data_buffer.py
@@ -0,0 +1,99 @@
+from typing import Callable, List, Optional
+
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+
+class InputDataBuffer(PhysicalOperator):
+    """Defines the input data for the operator DAG.
+
+    For example, this may hold cached blocks from a previous Dataset execution, or
+    the arguments for read tasks.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        input_data: Optional[List[RefBundle]] = None,
+        input_data_factory: Optional[Callable[[int], List[RefBundle]]] = None,
+    ):
+        """Create an InputDataBuffer.
+
+        Args:
+            data_context: :class:`~ray.data.context.DataContext`
+                object to use injestion.
+            input_data: The list of bundles to output from this operator.
+            input_data_factory: The factory to get input data, if input_data is None.
+        """
+        super().__init__("Input", [], data_context)
+        if input_data is not None:
+            assert input_data_factory is None
+            # Copy the input data to avoid mutating the original list.
+            self._input_data = input_data[:]
+            self._is_input_initialized = True
+            self._initialize_metadata()
+        else:
+            # Initialize input lazily when execution is started.
+            assert input_data_factory is not None
+            self._input_data_factory = input_data_factory
+            self._is_input_initialized = False
+        self._input_data_index = 0
+        self.mark_execution_finished()
+
+    def start(self, options: ExecutionOptions) -> None:
+        if not self._is_input_initialized:
+            self._input_data = self._input_data_factory(
+                self.target_max_block_size_override
+                or self.data_context.target_max_block_size
+            )
+            self._is_input_initialized = True
+            self._initialize_metadata()
+        # InputDataBuffer does not take inputs from other operators,
+        # so we record input metrics here
+        for bundle in self._input_data:
+            self._metrics.on_input_received(bundle)
+        super().start(options)
+
+    def has_next(self) -> bool:
+        return self._input_data_index < len(self._input_data)
+
+    def _get_next_inner(self) -> RefBundle:
+        # We can't pop the input data. If we do, Ray might garbage collect the block
+        # references, and Ray won't be able to reconstruct downstream objects.
+        bundle = self._input_data[self._input_data_index]
+        self._input_data_index += 1
+        return bundle
+
+    def get_stats(self) -> StatsDict:
+        return {}
+
+    def _add_input_inner(self, refs, input_index) -> None:
+        raise ValueError("Inputs are not allowed for this operator.")
+
+    def _initialize_metadata(self):
+        assert self._input_data is not None and self._is_input_initialized
+        self._estimated_num_output_bundles = len(self._input_data)
+
+        block_metadata = []
+        total_rows = 0
+        for bundle in self._input_data:
+            block_metadata.extend(bundle.metadata)
+            bundle_num_rows = bundle.num_rows()
+            if total_rows is not None and bundle_num_rows is not None:
+                total_rows += bundle_num_rows
+            else:
+                # total row is unknown
+                total_rows = None
+        if total_rows:
+            self._estimated_num_output_rows = total_rows
+        self._stats = {
+            "input": block_metadata,
+        }
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        return True
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/join.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/join.py
new file mode 100644
index 0000000000000000000000000000000000000000..74f5897bd7892afa2ba18b0dec5aee3254451270
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/join.py
@@ -0,0 +1,497 @@
+import logging
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Type
+
+from ray._private.arrow_utils import get_pyarrow_version
+from ray.air.util.transform_pyarrow import _is_pa_extension_type
+from ray.data._internal.arrow_block import ArrowBlockAccessor, ArrowBlockBuilder
+from ray.data._internal.arrow_ops.transform_pyarrow import (
+    MIN_PYARROW_VERSION_RUN_END_ENCODED_TYPES,
+    MIN_PYARROW_VERSION_VIEW_TYPES,
+)
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.hash_shuffle import (
+    HashShufflingOperatorBase,
+    StatefulShuffleAggregation,
+)
+from ray.data._internal.logical.operators.join_operator import JoinType
+from ray.data._internal.util import GiB, MiB
+from ray.data.block import Block
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    import pyarrow as pa
+
+
+@dataclass(frozen=True)
+class _DatasetPreprocessingResult:
+    """Result of join preprocessing containing split tables.
+
+    Separates tables into supported (directly joinable) and unsupported
+    (requires indexing) column projections.
+    """
+
+    supported_projection: "pa.Table"
+    unsupported_projection: "pa.Table"
+
+
+_JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP = {
+    JoinType.INNER: "inner",
+    JoinType.LEFT_OUTER: "left outer",
+    JoinType.RIGHT_OUTER: "right outer",
+    JoinType.FULL_OUTER: "full outer",
+    JoinType.LEFT_SEMI: "left semi",
+    JoinType.RIGHT_SEMI: "right semi",
+    JoinType.LEFT_ANTI: "left anti",
+    JoinType.RIGHT_ANTI: "right anti",
+}
+
+
+logger = logging.getLogger(__name__)
+
+
+class JoiningShuffleAggregation(StatefulShuffleAggregation):
+    """Aggregation performing distributed joining of the 2 sequences,
+    by utilising hash-based shuffling.
+
+    Hash-based shuffling applied to 2 input sequences and employing the same
+    partitioning scheme allows to
+
+        - Accumulate identical keys from both sequences into the same
+        (numerical) partition. In other words, all keys such that
+
+            hash(key) % num_partitions = partition_id
+
+        - Perform join on individual partitions independently (from other partitions)
+
+    For actual joining Pyarrow native joining functionality is utilised, providing
+    incredible performance while allowing keep the data from being deserialized.
+    """
+
+    def __init__(
+        self,
+        *,
+        aggregator_id: int,
+        join_type: JoinType,
+        left_key_col_names: Tuple[str],
+        right_key_col_names: Tuple[str],
+        target_partition_ids: List[int],
+        data_context: DataContext,
+        left_columns_suffix: Optional[str] = None,
+        right_columns_suffix: Optional[str] = None,
+    ):
+        super().__init__(aggregator_id)
+
+        assert (
+            len(left_key_col_names) > 0
+        ), "At least 1 column to join on has to be provided"
+        assert len(right_key_col_names) == len(
+            left_key_col_names
+        ), "Number of column for both left and right join operands has to match"
+
+        assert join_type in _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP, (
+            f"Join type is not currently supported (got: {join_type}; "  # noqa: C416
+            f"supported: {[jt for jt in JoinType]})"  # noqa: C416
+        )
+
+        self._left_key_col_names: Tuple[str] = left_key_col_names
+        self._right_key_col_names: Tuple[str] = right_key_col_names
+        self._join_type: JoinType = join_type
+
+        self._left_columns_suffix: Optional[str] = left_columns_suffix
+        self._right_columns_suffix: Optional[str] = right_columns_suffix
+
+        # Partition builders for the partition corresponding to
+        # left and right input sequences respectively
+        self._left_input_seq_partition_builders: Dict[int, ArrowBlockBuilder] = {
+            partition_id: ArrowBlockBuilder() for partition_id in target_partition_ids
+        }
+
+        self._right_input_seq_partition_builders: Dict[int, ArrowBlockBuilder] = {
+            partition_id: ArrowBlockBuilder() for partition_id in target_partition_ids
+        }
+        self.data_context = data_context
+
+    def accept(self, input_seq_id: int, partition_id: int, partition_shard: Block):
+        assert 0 <= input_seq_id < 2
+
+        partition_builder = self._get_partition_builder(
+            input_seq_id=input_seq_id,
+            partition_id=partition_id,
+        )
+
+        partition_builder.add_block(partition_shard)
+
+    def finalize(self, partition_id: int) -> Block:
+
+        left_on, right_on = list(self._left_key_col_names), list(
+            self._right_key_col_names
+        )
+
+        preprocess_result_l, preprocess_result_r = self._preprocess(
+            left_on, right_on, partition_id
+        )
+
+        # Perform the join on supported columns
+        arrow_join_type = _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP[self._join_type]
+
+        # Perform the join on supported columns
+        supported = preprocess_result_l.supported_projection.join(
+            preprocess_result_r.supported_projection,
+            join_type=arrow_join_type,
+            keys=left_on,
+            right_keys=right_on,
+            left_suffix=self._left_columns_suffix,
+            right_suffix=self._right_columns_suffix,
+        )
+
+        # Add back unsupported columns (join type logic is in should_index_* variables)
+        supported = self._postprocess(
+            supported,
+            preprocess_result_l.unsupported_projection,
+            preprocess_result_r.unsupported_projection,
+        )
+
+        return supported
+
+    def _preprocess(
+        self,
+        left_on: List[str],
+        right_on: List[str],
+        partition_id: int,
+    ) -> Tuple[_DatasetPreprocessingResult, _DatasetPreprocessingResult]:
+        import pyarrow as pa
+
+        left_seq_partition: pa.Table = self._get_partition_builder(
+            input_seq_id=0, partition_id=partition_id
+        ).build()
+
+        right_seq_partition: pa.Table = self._get_partition_builder(
+            input_seq_id=1, partition_id=partition_id
+        ).build()
+
+        # Get supported columns
+        supported_l, unsupported_l = self._split_unsupported_columns(left_seq_partition)
+        supported_r, unsupported_r = self._split_unsupported_columns(
+            right_seq_partition
+        )
+
+        # Handle joins on unsupported columns
+        conflicting_columns: Set[str] = set(unsupported_l.column_names) & set(left_on)
+        if conflicting_columns:
+            raise ValueError(
+                f"Cannot join on columns with unjoinable types. "
+                f"Left join key columns {conflicting_columns} have unjoinable types "
+                f"(map, union, list, struct, etc.) which cannot be used for join operations."
+            )
+
+        conflicting_columns: Set[str] = set(unsupported_r.column_names) & set(right_on)
+        if conflicting_columns:
+            raise ValueError(
+                f"Cannot join on columns with unjoinable types. "
+                f"Right join key columns {conflicting_columns} have unjoinable types "
+                f"(map, union, list, struct, etc.) which cannot be used for join operations."
+            )
+
+        # Index if we have unsupported columns
+        should_index_l = self._should_index_side("left", supported_l, unsupported_l)
+        should_index_r = self._should_index_side("right", supported_r, unsupported_r)
+
+        # Add index columns for back-referencing if we have unsupported columns
+        if should_index_l:
+            supported_l = self._append_index_column(
+                table=supported_l, col_name=self._index_name("left")
+            )
+        if should_index_r:
+            supported_r = self._append_index_column(
+                table=supported_r, col_name=self._index_name("right")
+            )
+
+        left = _DatasetPreprocessingResult(
+            supported_projection=supported_l,
+            unsupported_projection=unsupported_l,
+        )
+        right = _DatasetPreprocessingResult(
+            supported_projection=supported_r,
+            unsupported_projection=unsupported_r,
+        )
+        return left, right
+
+    def _postprocess(
+        self,
+        supported: "pa.Table",
+        unsupported_l: "pa.Table",
+        unsupported_r: "pa.Table",
+    ) -> "pa.Table":
+        # Index if we have unsupported columns
+        should_index_l = self._index_name("left") in supported.schema.names
+        should_index_r = self._index_name("right") in supported.schema.names
+
+        # Add back unsupported columns (join type logic is in should_index_* variables)
+        if should_index_l:
+            supported = self._add_back_unsupported_columns(
+                joined_table=supported,
+                unsupported_table=unsupported_l,
+                index_col_name=self._index_name("left"),
+            )
+
+        if should_index_r:
+            supported = self._add_back_unsupported_columns(
+                joined_table=supported,
+                unsupported_table=unsupported_r,
+                index_col_name=self._index_name("right"),
+            )
+
+        return supported
+
+    def _index_name(self, suffix: str) -> str:
+        return f"__rd_index_level_{suffix}__"
+
+    def clear(self, partition_id: int):
+        self._left_input_seq_partition_builders.pop(partition_id)
+        self._right_input_seq_partition_builders.pop(partition_id)
+
+    def _get_partition_builder(self, *, input_seq_id: int, partition_id: int):
+        if input_seq_id == 0:
+            partition_builder = self._left_input_seq_partition_builders[partition_id]
+        elif input_seq_id == 1:
+            partition_builder = self._right_input_seq_partition_builders[partition_id]
+        else:
+            raise ValueError(
+                f"Unexpected inpt sequence id of '{input_seq_id}' (expected 0 or 1)"
+            )
+        return partition_builder
+
+    def _should_index_side(
+        self, side: str, supported_table: "pa.Table", unsupported_table: "pa.Table"
+    ) -> bool:
+        """
+        Determine whether to create an index column for a given side of the join.
+
+        A column is "supported" if it is "joinable", and "unsupported" otherwise.
+        A supported_table is a table with only "supported" columns. Index columns are
+        needed when we have both supported and unsupported columns in a table, and
+        that table's columns will appear in the final result.
+
+        Args:
+            side: "left" or "right" to indicate which side of the join
+            supported_table: Table containing ONLY joinable columns
+            unsupported_table: Table containing ONLY unjoinable columns
+
+        Returns:
+            True if an index column should be created for this side
+        """
+        # Must have both supported and unsupported columns to need indexing.
+        # We cannot rely on row_count because it can return a non-zero row count
+        # for an empty-schema.
+        if not supported_table.schema or not unsupported_table.schema:
+            return False
+
+        # For semi/anti joins, only index the side that appears in the result
+        if side == "left":
+            # Left side appears in result for all joins except right_semi/right_anti
+            return self._join_type not in [JoinType.RIGHT_SEMI, JoinType.RIGHT_ANTI]
+        else:  # side == "right"
+            # Right side appears in result for all joins except left_semi/left_anti
+            return self._join_type not in [JoinType.LEFT_SEMI, JoinType.LEFT_ANTI]
+
+    def _split_unsupported_columns(
+        self, table: "pa.Table"
+    ) -> Tuple["pa.Table", "pa.Table"]:
+        """
+        Split a PyArrow table into two tables based on column joinability.
+
+        Separates columns into supported types and unsupported types that cannot be
+        directly joined on but should be preserved in results.
+
+        Args:
+            table: Input PyArrow table to split
+
+        Returns:
+            Tuple of (supported_table, unsupported_table) where:
+            - supported_table contains columns with primitive/joinable types
+            - unsupported_table contains columns with complex/unjoinable types
+        """
+        supported, unsupported = [], []
+        for idx in range(len(table.columns)):
+            col: "pa.ChunkedArray" = table.column(idx)
+            col_type: "pa.DataType" = col.type
+
+            if _is_pa_extension_type(col_type) or self._is_pa_join_not_supported(
+                col_type
+            ):
+                unsupported.append(idx)
+            else:
+                supported.append(idx)
+
+        return table.select(supported), table.select(unsupported)
+
+    def _add_back_unsupported_columns(
+        self,
+        joined_table: "pa.Table",
+        unsupported_table: "pa.Table",
+        index_col_name: str,
+    ) -> "pa.Table":
+        # Extract the index column array and drop the column from the joined table
+        i = joined_table.schema.get_field_index(index_col_name)
+        indices = joined_table.column(i)
+        joined_table = joined_table.remove_column(i)
+
+        # Project the unsupported columns using the indices and combine with joined table
+        projected = ArrowBlockAccessor(unsupported_table).take(indices)
+        return ArrowBlockAccessor(joined_table).hstack(projected)
+
+    def _append_index_column(self, table: "pa.Table", col_name: str) -> "pa.Table":
+        import numpy as np
+        import pyarrow as pa
+
+        index_col = pa.array(np.arange(table.num_rows))
+        return table.append_column(col_name, index_col)
+
+    def _is_pa_join_not_supported(self, type: "pa.DataType") -> bool:
+        """
+        The latest pyarrow versions do not support joins where the
+        tables contain the following types below (lists,
+        structs, maps, unions, extension types, etc.)
+
+        Args:
+            type: The input type of column.
+
+        Returns:
+            True if the type cannot be present (non join-key) during joins.
+            False if the type can be present.
+        """
+        import pyarrow as pa
+
+        pyarrow_version = get_pyarrow_version()
+        is_v12 = pyarrow_version >= MIN_PYARROW_VERSION_RUN_END_ENCODED_TYPES
+        is_v16 = pyarrow_version >= MIN_PYARROW_VERSION_VIEW_TYPES
+
+        return (
+            pa.types.is_map(type)
+            or pa.types.is_union(type)
+            or pa.types.is_list(type)
+            or pa.types.is_struct(type)
+            or pa.types.is_null(type)
+            or pa.types.is_large_list(type)
+            or pa.types.is_fixed_size_list(type)
+            or (is_v12 and pa.types.is_run_end_encoded(type))
+            or (
+                is_v16
+                and (
+                    pa.types.is_binary_view(type)
+                    or pa.types.is_string_view(type)
+                    or pa.types.is_list_view(type)
+                )
+            )
+        )
+
+
+class JoinOperator(HashShufflingOperatorBase):
+    def __init__(
+        self,
+        data_context: DataContext,
+        left_input_op: PhysicalOperator,
+        right_input_op: PhysicalOperator,
+        left_key_columns: Tuple[str],
+        right_key_columns: Tuple[str],
+        join_type: JoinType,
+        *,
+        num_partitions: Optional[int] = None,
+        left_columns_suffix: Optional[str] = None,
+        right_columns_suffix: Optional[str] = None,
+        partition_size_hint: Optional[int] = None,
+        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
+        shuffle_aggregation_type: Optional[Type[StatefulShuffleAggregation]] = None,
+    ):
+        if shuffle_aggregation_type is not None:
+            if not issubclass(shuffle_aggregation_type, StatefulShuffleAggregation):
+                raise TypeError(
+                    f"shuffle_aggregation_type must be a subclass of StatefulShuffleAggregation, "
+                    f"got {shuffle_aggregation_type}"
+                )
+
+        aggregation_class = shuffle_aggregation_type or JoiningShuffleAggregation
+        super().__init__(
+            name_factory=(
+                lambda num_partitions: f"Join(num_partitions={num_partitions})"
+            ),
+            input_ops=[left_input_op, right_input_op],
+            data_context=data_context,
+            key_columns=[left_key_columns, right_key_columns],
+            num_partitions=num_partitions,
+            partition_size_hint=partition_size_hint,
+            partition_aggregation_factory=(
+                lambda aggregator_id, target_partition_ids: aggregation_class(
+                    aggregator_id=aggregator_id,
+                    join_type=join_type,
+                    left_key_col_names=left_key_columns,
+                    right_key_col_names=right_key_columns,
+                    target_partition_ids=target_partition_ids,
+                    data_context=data_context,
+                    left_columns_suffix=left_columns_suffix,
+                    right_columns_suffix=right_columns_suffix,
+                )
+            ),
+            aggregator_ray_remote_args_override=aggregator_ray_remote_args_override,
+            shuffle_progress_bar_name="Shuffle",
+            finalize_progress_bar_name="Join",
+        )
+
+    def _get_operator_num_cpus_override(self) -> float:
+        return self.data_context.join_operator_actor_num_cpus_override
+
+    @classmethod
+    def _estimate_aggregator_memory_allocation(
+        cls,
+        *,
+        num_aggregators: int,
+        num_partitions: int,
+        estimated_dataset_bytes: int,
+    ) -> int:
+        partition_byte_size_estimate = math.ceil(
+            estimated_dataset_bytes / num_partitions
+        )
+
+        # Estimate of object store memory required to accommodate all partitions
+        # handled by a single aggregator
+        aggregator_shuffle_object_store_memory_required: int = math.ceil(
+            estimated_dataset_bytes / num_aggregators
+        )
+        # Estimate of memory required to perform actual (in-memory) join
+        # operation (inclusive of 50% overhead allocated for Pyarrow join
+        # implementation)
+        #
+        # NOTE:
+        #   - 2x due to budgeted 100% overhead of Arrow's in-memory join
+        join_memory_required: int = math.ceil(partition_byte_size_estimate * 2)
+        # Estimate of memory required to accommodate single partition as an output
+        # (inside Object Store)
+        #
+        # NOTE: x2 due to 2 sequences involved in joins
+        output_object_store_memory_required: int = partition_byte_size_estimate
+
+        aggregator_total_memory_required: int = (
+            # Inputs (object store)
+            aggregator_shuffle_object_store_memory_required
+            +
+            # Join (heap)
+            join_memory_required
+            +
+            # Output (object store)
+            output_object_store_memory_required
+        )
+
+        logger.info(
+            f"Estimated memory requirement for joining aggregator "
+            f"(partitions={num_partitions}, "
+            f"aggregators={num_aggregators}, "
+            f"dataset (estimate)={estimated_dataset_bytes / GiB:.1f}GiB): "
+            f"shuffle={aggregator_shuffle_object_store_memory_required / MiB:.1f}MiB, "
+            f"joining={join_memory_required / MiB:.1f}MiB, "
+            f"output={output_object_store_memory_required / MiB:.1f}MiB, "
+            f"total={aggregator_total_memory_required / MiB:.1f}MiB, "
+        )
+
+        return aggregator_total_memory_required
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/limit_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/limit_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6fe8ff56b407cae844e0dfef76b1a9de7345761
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/limit_operator.py
@@ -0,0 +1,136 @@
+import copy
+from collections import deque
+from typing import Deque, List, Optional, Tuple
+
+import ray
+from ray.data._internal.execution.interfaces import PhysicalOperator, RefBundle
+from ray.data._internal.execution.operators.base_physical_operator import (
+    OneToOneOperator,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data.block import Block, BlockAccessor, BlockMetadata, BlockStats
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+
+class LimitOperator(OneToOneOperator):
+    """Physical operator for limit."""
+
+    def __init__(
+        self,
+        limit: int,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+    ):
+        self._limit = limit
+        self._consumed_rows = 0
+        self._buffer: Deque[RefBundle] = deque()
+        self._name = f"limit={limit}"
+        self._output_blocks_stats: List[BlockStats] = []
+        self._cur_output_bundles = 0
+        super().__init__(self._name, input_op, data_context)
+        if self._limit <= 0:
+            self.mark_execution_finished()
+
+    def _limit_reached(self) -> bool:
+        return self._consumed_rows >= self._limit
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.completed()
+        assert input_index == 0, input_index
+        if self._limit_reached():
+            return
+        out_blocks: List[ObjectRef[Block]] = []
+        out_metadata: List[BlockMetadata] = []
+        for block, metadata in refs.blocks:
+            num_rows = metadata.num_rows
+            assert num_rows is not None
+            if self._consumed_rows + num_rows <= self._limit:
+                out_blocks.append(block)
+                out_metadata.append(metadata)
+                self._output_blocks_stats.append(metadata.to_stats())
+                self._consumed_rows += num_rows
+            else:
+                # Slice the last block.
+                def slice_fn(block, metadata, num_rows) -> Tuple[Block, BlockMetadata]:
+                    block = BlockAccessor.for_block(block).slice(
+                        0, num_rows, copy=False
+                    )
+                    metadata = copy.deepcopy(metadata)
+                    metadata.num_rows = num_rows
+                    metadata.size_bytes = BlockAccessor.for_block(block).size_bytes()
+                    return block, metadata
+
+                block, metadata_ref = cached_remote_fn(
+                    slice_fn, num_cpus=0, num_returns=2
+                ).remote(
+                    block,
+                    metadata,
+                    self._limit - self._consumed_rows,
+                )
+                out_blocks.append(block)
+                metadata = ray.get(metadata_ref)
+                out_metadata.append(metadata)
+                self._output_blocks_stats.append(metadata.to_stats())
+                self._consumed_rows = self._limit
+                break
+        self._cur_output_bundles += 1
+        out_refs = RefBundle(
+            list(zip(out_blocks, out_metadata)),
+            owns_blocks=refs.owns_blocks,
+            schema=refs.schema,
+        )
+        self._buffer.append(out_refs)
+        self._metrics.on_output_queued(out_refs)
+        if self._limit_reached():
+            self.mark_execution_finished()
+
+        # We cannot estimate if we have only consumed empty blocks,
+        # or if the input dependency's total number of output bundles is unknown.
+        num_inputs = self.input_dependencies[0].num_outputs_total()
+        if self._consumed_rows > 0 and num_inputs is not None:
+            # Estimate number of output bundles
+            # Check the case where _limit > # of input rows
+            estimated_total_output_rows = min(
+                self._limit, self._consumed_rows / self._cur_output_bundles * num_inputs
+            )
+            # _consumed_rows / _limit is roughly equal to
+            # _cur_output_bundles / total output blocks
+            self._estimated_num_output_bundles = round(
+                estimated_total_output_rows
+                / self._consumed_rows
+                * self._cur_output_bundles
+            )
+
+    def has_next(self) -> bool:
+        return len(self._buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        output = self._buffer.popleft()
+        self._metrics.on_output_dequeued(output)
+        return output
+
+    def get_stats(self) -> StatsDict:
+        return {self._name: self._output_blocks_stats}
+
+    def num_outputs_total(self) -> Optional[int]:
+        # Before execution is completed, we don't know how many output
+        # bundles we will have. We estimate based off the consumption so far.
+        if self._execution_finished:
+            return self._cur_output_bundles
+        return self._estimated_num_output_bundles
+
+    def num_output_rows_total(self) -> Optional[int]:
+        # The total number of rows is simply the limit or the number
+        # of input rows, whichever is smaller
+        input_num_rows = self.input_dependencies[0].num_output_rows_total()
+        if input_num_rows is None:
+            return None
+        return min(self._limit, input_num_rows)
+
+    def throttling_disabled(self) -> bool:
+        return True
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        return True
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/map_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/map_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..d2c628dbcad6f41a07d0a8e111196dc7cbb1fa60
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/map_operator.py
@@ -0,0 +1,1006 @@
+import copy
+import functools
+import itertools
+import logging
+from abc import ABC, abstractmethod
+from collections import defaultdict, deque
+from typing import (
+    Any,
+    Callable,
+    Deque,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
+
+import ray
+from ray import ObjectRef
+from ray._raylet import ObjectRefGenerator
+from ray.data._internal.compute import (
+    ActorPoolStrategy,
+    ComputeStrategy,
+    TaskPoolStrategy,
+)
+from ray.data._internal.execution.interfaces import (
+    BlockSlice,
+    ExecutionOptions,
+    ExecutionResources,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    DataOpTask,
+    MetadataOpTask,
+    OpTask,
+    estimate_total_num_of_blocks,
+)
+from ray.data._internal.execution.interfaces.ref_bundle import (
+    _iter_sliced_blocks,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+    OneToOneOperator,
+)
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import MemoryProfiler
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadataWithSchema,
+    BlockStats,
+    _take_first_non_empty_schema,
+    to_stats,
+)
+from ray.data.context import DataContext
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+logger = logging.getLogger(__name__)
+
+
+class BaseRefBundler(ABC):
+    """Interface for the rebundling behavior of the MapOperator."""
+
+    @abstractmethod
+    def num_blocks(self) -> int:
+        """Return the total number of blocks buffered inside the bundler."""
+        pass
+
+    @abstractmethod
+    def add_bundle(self, bundle: RefBundle):
+        """Add a new input bundle to the bundler."""
+        pass
+
+    @abstractmethod
+    def has_bundle(self) -> bool:
+        """Return whether the bundler currently holds a full bundle ready to emit."""
+        pass
+
+    @abstractmethod
+    def size_bytes(self) -> int:
+        """Estimate the total size in bytes of buffered bundles."""
+        pass
+
+    @abstractmethod
+    def get_next_bundle(
+        self,
+    ) -> Tuple[List[RefBundle], RefBundle]:
+        """Pop and return the next bundled input ready for task submission."""
+        pass
+
+    @abstractmethod
+    def done_adding_bundles(self):
+        """Signal that no additional bundles will be added to the bundler so the bundler can be finalized."""
+        pass
+
+
+class MapOperator(InternalQueueOperatorMixin, OneToOneOperator, ABC):
+    """A streaming operator that maps input bundles 1:1 to output bundles.
+
+    This operator implements the distributed map operation, supporting both task
+    and actor compute strategies.
+    """
+
+    MAP_UDF_WARN_SIZE_THRESHOLD = 100 * 1024**2
+    """
+    Warn if the size of the map UDF exceeds this threshold.
+    """
+
+    def __init__(
+        self,
+        map_transformer: MapTransformer,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        name: str,
+        target_max_block_size_override: Optional[int],
+        min_rows_per_bundle: Optional[int],
+        ref_bundler: Optional[BaseRefBundler],
+        supports_fusion: bool,
+        map_task_kwargs: Optional[Dict[str, Any]],
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]],
+        ray_remote_args: Optional[Dict[str, Any]],
+    ):
+        # NOTE: This constructor should not be called directly; use MapOperator.create()
+        # instead.
+        # NOTE: This constructor must be called by subclasses.
+        if map_task_kwargs is None:
+            map_task_kwargs = {}
+
+        self._map_transformer = map_transformer
+        self._supports_fusion = supports_fusion
+        self._map_task_kwargs = map_task_kwargs
+        self._ray_remote_args = _canonicalize_ray_remote_args(ray_remote_args or {})
+        self._ray_remote_args_fn = ray_remote_args_fn
+        self._ray_remote_args_factory_actor_locality = None
+        self._remote_args_for_metrics = copy.deepcopy(self._ray_remote_args)
+
+        # Bundles block references up to the min_rows_per_bundle target.
+        self._block_ref_bundler = ref_bundler or BlockRefBundler(min_rows_per_bundle)
+
+        # Queue for task outputs, either ordered or unordered (this is set by start()).
+        self._output_queue: Optional[_OutputQueue] = None
+        # Output metadata, added to on get_next().
+        self._output_blocks_stats: List[BlockStats] = []
+        # All active `DataOpTask`s.
+        self._data_tasks: Dict[int, DataOpTask] = {}
+        self._next_data_task_idx = 0
+        # All active `MetadataOpTask`s.
+        self._metadata_tasks: Dict[int, MetadataOpTask] = {}
+        self._next_metadata_task_idx = 0
+        # Keep track of all finished streaming generators.
+        super().__init__(name, input_op, data_context, target_max_block_size_override)
+
+        # If set, then all output blocks will be split into
+        # this many sub-blocks. This is to avoid having
+        # too-large blocks, which may reduce parallelism for
+        # the subsequent operator.
+        self._additional_split_factor = None
+        # Callback functions that generate additional task kwargs
+        # for the map task.
+        self._map_task_kwargs_fns: List[Callable[[], Dict[str, Any]]] = []
+
+    def add_map_task_kwargs_fn(self, map_task_kwargs_fn: Callable[[], Dict[str, Any]]):
+        """Add a callback function that generates additional kwargs for the map tasks.
+        In the map tasks, the kwargs can be accessible via `TaskContext.kwargs`.
+        """
+        self._map_task_kwargs_fns.append(map_task_kwargs_fn)
+
+    def get_map_task_kwargs(self) -> Dict[str, Any]:
+        """Get the kwargs for the map task.
+        Subclasses should pass the returned kwargs to the map tasks.
+        In the map tasks, the kwargs can be accessible via `TaskContext.kwargs`.
+        """
+        kwargs = self._map_task_kwargs.copy()
+        for fn in self._map_task_kwargs_fns:
+            kwargs.update(fn())
+        return kwargs
+
+    def get_additional_split_factor(self) -> int:
+        if self._additional_split_factor is None:
+            return 1
+        return self._additional_split_factor
+
+    def set_additional_split_factor(self, k: int):
+        self._additional_split_factor = k
+
+    def internal_input_queue_num_blocks(self) -> int:
+        return self._block_ref_bundler.num_blocks()
+
+    def internal_input_queue_num_bytes(self) -> int:
+        return self._block_ref_bundler.size_bytes()
+
+    def internal_output_queue_num_blocks(self) -> int:
+        return self._output_queue.num_blocks()
+
+    def internal_output_queue_num_bytes(self) -> int:
+        return self._output_queue.size_bytes()
+
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queue (block ref bundler)."""
+        while self._block_ref_bundler.has_bundle():
+            (input_bundles, _) = self._block_ref_bundler.get_next_bundle()
+            for input_bundle in input_bundles:
+                self._metrics.on_input_dequeued(input_bundle)
+
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue."""
+        while self._output_queue.has_next():
+            bundle = self._output_queue.get_next()
+            self._metrics.on_output_dequeued(bundle)
+
+    @property
+    def name(self) -> str:
+        name = super().name
+        if self._additional_split_factor is not None:
+            name += f"->SplitBlocks({self._additional_split_factor})"
+        return name
+
+    @classmethod
+    def create(
+        cls,
+        map_transformer: MapTransformer,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        target_max_block_size_override: Optional[int] = None,
+        name: str = "Map",
+        # TODO(ekl): slim down ComputeStrategy to only specify the compute
+        # config and not contain implementation code.
+        compute_strategy: Optional[ComputeStrategy] = None,
+        min_rows_per_bundle: Optional[int] = None,
+        ref_bundler: Optional[BaseRefBundler] = None,
+        supports_fusion: bool = True,
+        map_task_kwargs: Optional[Dict[str, Any]] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+        per_block_limit: Optional[int] = None,
+    ) -> "MapOperator":
+        """Create a MapOperator.
+
+        This factory creates the MapOperator pool implementation that corresponds to the
+        compute argument:
+            - If None or TaskPoolStrategy -> TaskPoolMapOperator
+            - If ActorPoolStrategy -> ActorPoolMapOperator
+
+        Args:
+            transform_fn: The function to apply to each ref bundle input.
+            input_op: Operator generating input data for this op.
+            init_fn: The callable class to instantiate if using ActorPoolMapOperator.
+            name: The name of this operator.
+            compute_strategy: Customize the compute strategy for this op.
+            target_max_block_size_override: Override for target max-block-size.
+            min_rows_per_bundle: The number of rows to gather per batch passed to the
+                transform_fn, or None to use the block size. Setting the batch size is
+                important for the performance of GPU-accelerated transform functions.
+                The actual rows passed may be less if the dataset is small.
+            ref_bundler: The ref bundler to use for this operator.
+            supports_fusion: Whether this operator supports fusion with other operators.
+            map_task_kwargs: A dictionary of kwargs to pass to the map task. You can
+                access these kwargs through the `TaskContext.kwargs` dictionary.
+            ray_remote_args_fn: A function that returns a dictionary of remote args
+                passed to each map worker. The purpose of this argument is to generate
+                dynamic arguments for each actor/task, and will be called each time
+                prior to initializing the worker. Args returned from this dict will
+                always override the args in ``ray_remote_args``. Note: this is an
+                advanced, experimental feature.
+            ray_remote_args: Customize the :func:`ray.remote` args for this op's tasks.
+            per_block_limit: Maximum number of rows to process per block, for early termination.
+        """
+        if (ref_bundler is not None and min_rows_per_bundle is not None) or (
+            min_rows_per_bundle is not None and ref_bundler is not None
+        ):
+            raise ValueError(
+                "min_rows_per_bundle and ref_bundler cannot be used together"
+            )
+
+        if compute_strategy is None:
+            compute_strategy = TaskPoolStrategy()
+
+        # Apply per-block limit to the map transformer if set
+        if per_block_limit is not None:
+            map_transformer = _wrap_transformer_with_limit(
+                map_transformer, per_block_limit
+            )
+
+        if isinstance(compute_strategy, TaskPoolStrategy):
+            from ray.data._internal.execution.operators.task_pool_map_operator import (
+                TaskPoolMapOperator,
+            )
+
+            return TaskPoolMapOperator(
+                map_transformer,
+                input_op,
+                data_context,
+                name=name,
+                target_max_block_size_override=target_max_block_size_override,
+                min_rows_per_bundle=min_rows_per_bundle,
+                ref_bundler=ref_bundler,
+                max_concurrency=compute_strategy.size,
+                supports_fusion=supports_fusion,
+                map_task_kwargs=map_task_kwargs,
+                ray_remote_args_fn=ray_remote_args_fn,
+                ray_remote_args=ray_remote_args,
+            )
+        elif isinstance(compute_strategy, ActorPoolStrategy):
+            from ray.data._internal.execution.operators.actor_pool_map_operator import (
+                ActorPoolMapOperator,
+            )
+
+            return ActorPoolMapOperator(
+                map_transformer,
+                input_op,
+                data_context,
+                target_max_block_size_override=target_max_block_size_override,
+                compute_strategy=compute_strategy,
+                name=name,
+                min_rows_per_bundle=min_rows_per_bundle,
+                ref_bundler=ref_bundler,
+                supports_fusion=supports_fusion,
+                map_task_kwargs=map_task_kwargs,
+                ray_remote_args_fn=ray_remote_args_fn,
+                ray_remote_args=ray_remote_args,
+            )
+        else:
+            raise ValueError(f"Unsupported execution strategy {compute_strategy}")
+
+    def start(self, options: "ExecutionOptions"):
+        super().start(options)
+        # Create output queue with desired ordering semantics.
+        if options.preserve_order:
+            self._output_queue = _OrderedOutputQueue()
+        else:
+            self._output_queue = _UnorderedOutputQueue()
+
+        if options.locality_with_output:
+            if isinstance(options.locality_with_output, list):
+                locs = options.locality_with_output
+            else:
+                locs = [ray.get_runtime_context().get_node_id()]
+
+            class RoundRobinAssign:
+                def __init__(self, locs):
+                    self.locs = locs
+                    self.i = 0
+
+                def __call__(self, args):
+                    args = copy.deepcopy(args)
+                    args["scheduling_strategy"] = NodeAffinitySchedulingStrategy(
+                        self.locs[self.i],
+                        soft=True,
+                        _spill_on_unavailable=True,
+                    )
+                    self.i += 1
+                    self.i %= len(self.locs)
+                    return args
+
+            self._ray_remote_args_factory_actor_locality = RoundRobinAssign(locs)
+
+        map_transformer = self._map_transformer
+        # Apply additional block split if needed.
+        if self.get_additional_split_factor() > 1:
+            split_factor = self.get_additional_split_factor()
+            split_transformer = MapTransformer(
+                [
+                    BlockMapTransformFn(
+                        lambda blocks, ctx: _split_blocks(blocks, split_factor),
+                        # NOTE: Disable block-shaping to avoid it overriding
+                        #       splitting
+                        disable_block_shaping=True,
+                    )
+                ]
+            )
+            map_transformer = map_transformer.fuse(split_transformer)
+        # Put the function def in the object store to avoid repeated serialization
+        # in case it's large (i.e., closure captures large objects).
+        self._map_transformer_ref = ray.put(map_transformer)
+        self._warn_large_udf()
+
+    def _warn_large_udf(self):
+        """Print a warning if the UDF is too large."""
+        udf_size = ray.experimental.get_local_object_locations(
+            [self._map_transformer_ref]
+        )[self._map_transformer_ref]["object_size"]
+        if udf_size > self.MAP_UDF_WARN_SIZE_THRESHOLD:
+            logger.warning(
+                f"The UDF of operator {self.name} is too large "
+                f"(size = {memory_string(udf_size)}). "
+                "Check if the UDF has accidentally captured large objects. "
+                "Load the large objects in the __init__ method "
+                "or pass them as ObjectRefs instead."
+            )
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int):
+        assert input_index == 0, input_index
+
+        # Add RefBundle to the bundler.
+        self._block_ref_bundler.add_bundle(refs)
+        self._metrics.on_input_queued(refs)
+
+        if self._block_ref_bundler.has_bundle():
+            # The ref bundler combines one or more RefBundles into a new larger
+            # RefBundle. Rather than dequeuing the new RefBundle, which was never
+            # enqueued in the first place, we dequeue the original RefBundles.
+            (input_refs, bundled_input) = self._block_ref_bundler.get_next_bundle()
+            for bundle in input_refs:
+                self._metrics.on_input_dequeued(bundle)
+
+            # If the bundler has a full bundle, add it to the operator's task submission
+            # queue
+            self._add_bundled_input(bundled_input)
+
+    def _get_dynamic_ray_remote_args(
+        self, input_bundle: Optional[RefBundle] = None
+    ) -> Dict[str, Any]:
+        ray_remote_args = copy.deepcopy(self._ray_remote_args)
+
+        # max_calls isn't supported in `.options()`, so we remove it when generating dynamic ray_remote_args
+        ray_remote_args.pop("max_calls", None)
+
+        # Override parameters from user provided remote args function.
+        if self._ray_remote_args_fn:
+            new_remote_args = self._ray_remote_args_fn()
+            for k, v in new_remote_args.items():
+                ray_remote_args[k] = v
+        # For tasks with small args, we will use SPREAD by default to optimize for
+        # compute load-balancing. For tasks with large args, we will use DEFAULT to
+        # allow the Ray locality scheduler a chance to optimize task placement.
+        if "scheduling_strategy" not in ray_remote_args:
+            ctx = self.data_context
+            if input_bundle and input_bundle.size_bytes() > ctx.large_args_threshold:
+                ray_remote_args[
+                    "scheduling_strategy"
+                ] = ctx.scheduling_strategy_large_args
+                # Takes precedence over small args case. This is to let users know
+                # when the large args case is being triggered.
+                self._remote_args_for_metrics = copy.deepcopy(ray_remote_args)
+            else:
+                ray_remote_args["scheduling_strategy"] = ctx.scheduling_strategy
+                # Only save to metrics if we haven't already done so.
+                if "scheduling_strategy" not in self._remote_args_for_metrics:
+                    self._remote_args_for_metrics = copy.deepcopy(ray_remote_args)
+        # This should take precedence over previously set scheduling strategy, as it
+        # implements actor-based locality overrides.
+        if self._ray_remote_args_factory_actor_locality:
+            return self._ray_remote_args_factory_actor_locality(ray_remote_args)
+        return ray_remote_args
+
+    @abstractmethod
+    def _add_bundled_input(self, refs: RefBundle):
+        """Add a pre-bundled upstream output to this operator.
+
+        Unlike the add_input() arg, this RefBundle has already been further bundled by
+        _block_ref_bundler up to the target size, meaning that this bundle is ready for
+        task submission.
+
+        This must be implemented by subclasses.
+
+        Args:
+            refs: The fully-bundled ref bundle that should be added as input.
+        """
+        raise NotImplementedError
+
+    def _submit_data_task(
+        self,
+        gen: ObjectRefGenerator,
+        inputs: RefBundle,
+        task_done_callback: Optional[Callable[[], None]] = None,
+    ):
+        """Submit a new data-handling task."""
+        # TODO(hchen):
+        # 1. Move this to the base PhyscialOperator class.
+        # 2. This method should only take a block-processing function as input,
+        #    instead of a streaming generator. The logic of submitting ray tasks
+        #    can also be capsulated in the base class.
+        task_index = self._next_data_task_idx
+        self._next_data_task_idx += 1
+        self._metrics.on_task_submitted(task_index, inputs)
+
+        def _output_ready_callback(
+            task_index,
+            output: RefBundle,
+        ):
+            # Since output is streamed, it should only contain one block.
+            assert len(output) == 1
+            self._metrics.on_task_output_generated(task_index, output)
+
+            # Notify output queue that the task has produced an new output.
+            self._output_queue.notify_task_output_ready(task_index, output)
+            self._metrics.on_output_queued(output)
+
+        def _task_done_callback(task_index: int, exception: Optional[Exception]):
+            self._metrics.on_task_finished(task_index, exception)
+
+            # Estimate number of tasks and rows from inputs received and tasks
+            # submitted so far
+            (
+                _,
+                self._estimated_num_output_bundles,
+                self._estimated_output_num_rows,
+            ) = estimate_total_num_of_blocks(
+                self._next_data_task_idx, self.upstream_op_num_outputs(), self._metrics
+            )
+
+            self._data_tasks.pop(task_index)
+            # Notify output queue that this task is complete.
+            self._output_queue.notify_task_completed(task_index)
+            if task_done_callback:
+                task_done_callback()
+
+        self._data_tasks[task_index] = DataOpTask(
+            task_index,
+            gen,
+            lambda output: _output_ready_callback(task_index, output),
+            functools.partial(_task_done_callback, task_index),
+        )
+
+    def _submit_metadata_task(
+        self, result_ref: ObjectRef, task_done_callback: Callable[[], None]
+    ):
+        """Submit a new metadata-handling task."""
+        # TODO(hchen): Move this to the base PhyscialOperator class.
+        task_index = self._next_metadata_task_idx
+        self._next_metadata_task_idx += 1
+
+        def _task_done_callback():
+            self._metadata_tasks.pop(task_index)
+            task_done_callback()
+
+        self._metadata_tasks[task_index] = MetadataOpTask(
+            task_index, result_ref, _task_done_callback
+        )
+
+    def get_active_tasks(self) -> List[OpTask]:
+        return list(self._metadata_tasks.values()) + list(self._data_tasks.values())
+
+    def all_inputs_done(self):
+        self._block_ref_bundler.done_adding_bundles()
+        if self._block_ref_bundler.has_bundle():
+            # Handle any leftover bundles in the bundler.
+            (
+                _,
+                bundled_input,
+            ) = self._block_ref_bundler.get_next_bundle()
+            self._add_bundled_input(bundled_input)
+        super().all_inputs_done()
+
+    def has_next(self) -> bool:
+        assert self._started
+        return self._output_queue.has_next()
+
+    def _get_next_inner(self) -> RefBundle:
+        assert self._started
+        bundle = self._output_queue.get_next()
+        self._metrics.on_output_dequeued(bundle)
+        self._output_blocks_stats.extend(to_stats(bundle.metadata))
+        return bundle
+
+    @abstractmethod
+    def progress_str(self) -> str:
+        raise NotImplementedError
+
+    def _extra_metrics(self) -> Dict[str, Any]:
+        return {"ray_remote_args": dict(sorted(self._remote_args_for_metrics.items()))}
+
+    def get_stats(self) -> StatsDict:
+        return {self._name: self._output_blocks_stats}
+
+    def get_map_transformer(self) -> MapTransformer:
+        return self._map_transformer
+
+    def _do_shutdown(self, force: bool = False):
+        # Invoke base-class sequence
+        super()._do_shutdown(force)
+        # Release refs
+        self._data_tasks.clear()
+        self._metadata_tasks.clear()
+
+    @abstractmethod
+    def current_processor_usage(self) -> ExecutionResources:
+        raise NotImplementedError
+
+    @abstractmethod
+    def pending_processor_usage(self) -> ExecutionResources:
+        raise NotImplementedError
+
+    @abstractmethod
+    def incremental_resource_usage(self) -> ExecutionResources:
+        raise NotImplementedError
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        return True
+
+    def supports_fusion(self) -> bool:
+        return self._supports_fusion
+
+    def num_active_tasks(self) -> int:
+        # Override `num_active_tasks` to only include data tasks and exclude
+        # metadata tasks, which are used by the actor-pool map operator to
+        # check if a newly created actor is ready.
+        # The reasons are because:
+        # 1. `PhysicalOperator.completed` checks `num_active_tasks`. The operator
+        #   should be considered completed if there are still pending actors.
+        # 2. The number of active tasks in the progress bar will be more accurate
+        #   to reflect the actual data processing tasks.
+        return len(self._data_tasks)
+
+
+def _map_task(
+    map_transformer: MapTransformer,
+    data_context: DataContext,
+    ctx: TaskContext,
+    *blocks: Block,
+    slices: Optional[List[BlockSlice]] = None,
+    **kwargs: Dict[str, Any],
+) -> Iterator[Union[Block, "BlockMetadataWithSchema"]]:
+    """Remote function for a single operator task.
+
+    Args:
+        fn: The callable that takes Iterator[Block] as input and returns
+            Iterator[Block] as output.
+        blocks: The concrete block values from the task ref bundle.
+        slices: List of block slices for this task to process.
+
+    Returns:
+        A generator of blocks, followed by the list of BlockMetadata for the blocks
+        as the last generator return.
+    """
+    logger.debug(
+        "Executing map task of operator %s with task index %d",
+        ctx.op_name,
+        ctx.task_idx,
+    )
+    DataContext._set_current(data_context)
+    ctx.kwargs.update(kwargs)
+    TaskContext.set_current(ctx)
+    stats = BlockExecStats.builder()
+    map_transformer.override_target_max_block_size(ctx.target_max_block_size_override)
+    block_iter: Iterable[Block]
+    if slices:
+        block_iter = _iter_sliced_blocks(blocks, slices)
+    else:
+        block_iter = iter(blocks)
+
+    with MemoryProfiler(data_context.memory_usage_poll_interval_s) as profiler:
+        for b_out in map_transformer.apply_transform(block_iter, ctx):
+            # TODO(Clark): Add input file propagation from input blocks.
+            m_out = BlockAccessor.for_block(b_out).get_metadata()
+            s_out = BlockAccessor.for_block(b_out).schema()
+            m_out.exec_stats = stats.build()
+            m_out.exec_stats.udf_time_s = map_transformer.udf_time()
+            m_out.exec_stats.task_idx = ctx.task_idx
+            m_out.exec_stats.max_uss_bytes = profiler.estimate_max_uss()
+            meta_with_schema = BlockMetadataWithSchema(metadata=m_out, schema=s_out)
+            yield b_out
+            yield meta_with_schema
+            stats = BlockExecStats.builder()
+            profiler.reset()
+
+    TaskContext.reset_current()
+
+
+class BlockRefBundler(BaseRefBundler):
+    """Rebundles RefBundles to get them close to a particular number of rows."""
+
+    def __init__(self, min_rows_per_bundle: Optional[int]):
+        """Creates a BlockRefBundler.
+
+        Args:
+            min_rows_per_bundle: The target number of rows per bundle. Note that we
+                bundle up to this target, but only exceed it if not doing so would
+                result in an empty bundle.
+        """
+        assert (
+            min_rows_per_bundle is None or min_rows_per_bundle >= 0
+        ), "Min rows per bundle has to be non-negative"
+
+        self._min_rows_per_bundle = min_rows_per_bundle
+        self._bundle_buffer: List[RefBundle] = []
+        self._bundle_buffer_size = 0
+        self._bundle_buffer_size_bytes = 0
+        self._finalized = False
+
+    def num_blocks(self):
+        return sum(len(b.block_refs) for b in self._bundle_buffer)
+
+    def add_bundle(self, bundle: RefBundle):
+        """Add a bundle to the bundler."""
+        self._bundle_buffer.append(bundle)
+        self._bundle_buffer_size += self._get_bundle_size(bundle)
+        self._bundle_buffer_size_bytes += bundle.size_bytes()
+
+    def has_bundle(self) -> bool:
+        """Returns whether the bundler has a bundle."""
+        return self._bundle_buffer and (
+            self._min_rows_per_bundle is None
+            or self._bundle_buffer_size >= self._min_rows_per_bundle
+            or (self._finalized and self._bundle_buffer_size >= 0)
+        )
+
+    def size_bytes(self) -> int:
+        return self._bundle_buffer_size_bytes
+
+    def get_next_bundle(
+        self,
+    ) -> Tuple[List[RefBundle], RefBundle]:
+        """Gets the next bundle.
+
+        Returns:
+            A two-tuple. The first element is a list of bundles that were combined into
+            the output bundle. The second element is the output bundle.
+        """
+        assert self.has_bundle()
+
+        if self._min_rows_per_bundle is None:
+            # Short-circuit if no bundle row target was defined.
+            assert len(self._bundle_buffer) == 1
+            bundle = self._bundle_buffer[0]
+            self._bundle_buffer = []
+            self._bundle_buffer_size = 0
+            self._bundle_buffer_size_bytes = 0
+            return [bundle], bundle
+
+        remainder = []
+        output_buffer = []
+        output_buffer_size = 0
+
+        for idx, bundle in enumerate(self._bundle_buffer):
+            bundle_size = self._get_bundle_size(bundle)
+
+            # Add bundle to the output buffer so long as either
+            #   - Output buffer size is still 0
+            #   - Output buffer doesn't exceeds the `_min_rows_per_bundle` threshold
+            if (
+                output_buffer_size < self._min_rows_per_bundle
+                or output_buffer_size == 0
+            ):
+                output_buffer.append(bundle)
+                output_buffer_size += bundle_size
+            else:
+                remainder = self._bundle_buffer[idx:]
+
+        self._bundle_buffer = remainder
+        self._bundle_buffer_size = sum(
+            self._get_bundle_size(bundle) for bundle in remainder
+        )
+        self._bundle_buffer_size_bytes = sum(
+            bundle.size_bytes() for bundle in remainder
+        )
+
+        return list(output_buffer), _merge_ref_bundles(*output_buffer)
+
+    def done_adding_bundles(self):
+        """Indicate that no more RefBundles will be added to this bundler."""
+        self._finalized = True
+
+    @staticmethod
+    def _get_bundle_size(bundle: RefBundle):
+        return bundle.num_rows() if bundle.num_rows() is not None else float("inf")
+
+
+def _merge_ref_bundles(*bundles: RefBundle) -> RefBundle:
+    """Merge N ref bundles into a single bundle of multiple blocks."""
+    # Check that at least one bundle is non-null.
+    bundles = [bundle for bundle in bundles if bundle is not None]
+    assert len(bundles) > 0
+    blocks = list(
+        itertools.chain(block for bundle in bundles for block in bundle.blocks)
+    )
+    owns_blocks = all(bundle.owns_blocks for bundle in bundles)
+    schema = _take_first_non_empty_schema(bundle.schema for bundle in bundles)
+    return RefBundle(blocks, owns_blocks=owns_blocks, schema=schema)
+
+
+class _OutputQueue(ABC):
+    """Interface for swapping between different output order modes."""
+
+    @abstractmethod
+    def notify_task_output_ready(self, task_index: int, output: RefBundle):
+        """Called when a task's output is ready."""
+        pass
+
+    def notify_task_completed(self, task_index: int):
+        """Called when a previously pending task completes."""
+        pass
+
+    @abstractmethod
+    def has_next(self) -> bool:
+        pass
+
+    @abstractmethod
+    def get_next(self) -> RefBundle:
+        pass
+
+    @abstractmethod
+    def num_blocks(self) -> int:
+        pass
+
+    @abstractmethod
+    def size_bytes(self) -> int:
+        pass
+
+
+class _OrderedOutputQueue(_OutputQueue):
+    """An queue that returns finished tasks in submission order."""
+
+    def __init__(self):
+        self._task_outputs: Dict[int, Deque[RefBundle]] = defaultdict(lambda: deque())
+        self._current_output_index: int = 0
+        self._completed_tasks: Set[int] = set()
+        self._size_bytes: int = 0
+        self._num_blocks: int = 0
+
+    def notify_task_output_ready(self, task_index: int, output: RefBundle):
+        self._task_outputs[task_index].append(output)
+        self._size_bytes += output.size_bytes()
+        self._num_blocks += len(output.blocks)
+
+    def _move_to_next_task(self):
+        """Move the outut index to the next task.
+
+        This method should only be called when the current task is complete and all
+        outputs have been taken.
+        """
+        assert len(self._task_outputs[self._current_output_index]) == 0
+        assert self._current_output_index in self._completed_tasks
+        del self._task_outputs[self._current_output_index]
+        self._completed_tasks.remove(self._current_output_index)
+        self._current_output_index += 1
+
+    def notify_task_completed(self, task_index: int):
+        assert task_index >= self._current_output_index
+        self._completed_tasks.add(task_index)
+        if task_index == self._current_output_index:
+            if len(self._task_outputs[task_index]) == 0:
+                self._move_to_next_task()
+
+    def has_next(self) -> bool:
+        return len(self._task_outputs[self._current_output_index]) > 0
+
+    def get_next(self) -> RefBundle:
+        next_bundle = self._task_outputs[self._current_output_index].popleft()
+        self._size_bytes -= next_bundle.size_bytes()
+        self._num_blocks -= len(next_bundle.blocks)
+        if len(self._task_outputs[self._current_output_index]) == 0:
+            if self._current_output_index in self._completed_tasks:
+                self._move_to_next_task()
+        return next_bundle
+
+    def num_blocks(self) -> int:
+        return self._num_blocks
+
+    def size_bytes(self) -> int:
+        return self._size_bytes
+
+
+class _UnorderedOutputQueue(_OutputQueue):
+    """An queue that does not guarantee output order of finished tasks."""
+
+    def __init__(self):
+        self._queue: Deque[RefBundle] = deque()
+        self._num_blocks: int = 0
+        self._size_bytes: int = 0
+
+    def notify_task_output_ready(self, _: int, output: RefBundle):
+        self._queue.append(output)
+        self._num_blocks += len(output.blocks)
+        self._size_bytes += output.size_bytes()
+
+    def has_next(self) -> bool:
+        return len(self._queue) > 0
+
+    def get_next(self) -> RefBundle:
+        next_bundle = self._queue.popleft()
+        self._num_blocks -= len(next_bundle.blocks)
+        self._size_bytes -= next_bundle.size_bytes()
+        return next_bundle
+
+    def num_blocks(self) -> int:
+        return self._num_blocks
+
+    def size_bytes(self) -> int:
+        return self._size_bytes
+
+
+def _canonicalize_ray_remote_args(ray_remote_args: Dict[str, Any]) -> Dict[str, Any]:
+    """Enforce rules on ray remote args for map tasks.
+
+    Namely, args must explicitly specify either CPU or GPU, not both. Disallowing
+    mixed resources avoids potential starvation and deadlock issues during scheduling,
+    and should not be a serious limitation for users.
+    """
+    ray_remote_args = ray_remote_args.copy()
+
+    # TODO: Might be better to log this warning at composition-time rather than at
+    # execution. Validating inputs early is a good practice.
+    if ray_remote_args.get("num_cpus") and ray_remote_args.get("num_gpus"):
+        logger.warning(
+            "Specifying both num_cpus and num_gpus for map tasks is experimental, "
+            "and may result in scheduling or stability issues. "
+            "Please report any issues to the Ray team: "
+            "https://github.com/ray-project/ray/issues/new/choose"
+        )
+
+    if "num_cpus" not in ray_remote_args and "num_gpus" not in ray_remote_args:
+        ray_remote_args["num_cpus"] = 1
+
+    return ray_remote_args
+
+
+def _splitrange(n, k):
+    """Calculates array lens of np.array_split().
+
+    This is the equivalent of
+    `[len(x) for x in np.array_split(range(n), k)]`.
+    """
+    base = n // k
+    output = [base] * k
+    rem = n - sum(output)
+    for i in range(len(output)):
+        if rem > 0:
+            output[i] += 1
+            rem -= 1
+    assert rem == 0, (rem, output, n, k)
+    assert sum(output) == n, (output, n, k)
+    return output
+
+
+def _split_blocks(blocks: Iterable[Block], split_factor: float) -> Iterable[Block]:
+    for block in blocks:
+        block = BlockAccessor.for_block(block)
+        offset = 0
+        split_sizes = _splitrange(block.num_rows(), split_factor)
+        for size in split_sizes:
+            if size <= 0:
+                continue
+            yield block.slice(offset, offset + size, copy=False)
+            offset += size
+
+
+def _wrap_transformer_with_limit(
+    map_transformer: MapTransformer, per_block_limit: int
+) -> MapTransformer:
+    """Wrap a MapTransformer with per-block limit functionality."""
+
+    # Create a new limit transform function that goes at the end
+    limit_transform_fn = _create_per_block_limit_transform_fn(per_block_limit)
+
+    # Add the limit transform as the last step
+    # Appending at the end so that the cap applies to the final output
+    # blocks after all prior transforms.
+    existing_transform_fns = map_transformer.get_transform_fns()
+    new_transform_fns = existing_transform_fns + [limit_transform_fn]
+
+    # Create new transformer with the limit added
+    # TODO: Modify `add_transform_fns` to do this operation internally instead of modifying in place.
+    new_transformer = MapTransformer(
+        new_transform_fns,
+        init_fn=map_transformer._init_fn,
+        output_block_size_option_override=map_transformer._output_block_size_option_override,
+    )
+
+    return new_transformer
+
+
+def _per_block_limit_fn(
+    input: Iterable[Block], ctx: TaskContext, per_block_limit: int
+) -> Iterable[Block]:
+    """Apply per-block limit to the input blocks."""
+    from ray.data.block import BlockAccessor
+
+    # This is used to track the number of rows processed within this task.
+    processed_rows = 0
+
+    for block in input:
+        if processed_rows >= per_block_limit:
+            # We've hit the limit, stop processing
+            break
+
+        block_accessor = BlockAccessor.for_block(block)
+        block_rows = block_accessor.num_rows()
+
+        if processed_rows + block_rows <= per_block_limit:
+            # Entire block fits within limit
+            processed_rows += block_rows
+            yield block
+        else:
+            # Need to truncate this block
+            remaining_rows = per_block_limit - processed_rows
+            truncated_block = block_accessor.slice(0, remaining_rows, copy=False)
+            processed_rows += remaining_rows
+            yield truncated_block
+
+
+def _create_per_block_limit_transform_fn(per_block_limit: int) -> BlockMapTransformFn:
+    """Create a transform function that applies per-block row limits."""
+    limit_fn = functools.partial(_per_block_limit_fn, per_block_limit=per_block_limit)
+    return BlockMapTransformFn(limit_fn)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/map_transformer.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/map_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..73cf8a50c8b717f0c4518e930f9dd96a0c66a33d
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/map_transformer.py
@@ -0,0 +1,428 @@
+import time
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar, Union
+
+from ray.data._internal.block_batching.block_batching import batch_blocks
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.output_buffer import BlockOutputBuffer, OutputBlockSizeOption
+from ray.data.block import BatchFormat, Block, BlockAccessor, DataBatch
+
+# Allowed input/output data types for a MapTransformFn.
+Row = Dict[str, Any]
+MapTransformFnData = Union[Block, Row, DataBatch]
+
+# Function signature of a MapTransformFn.
+IN = TypeVar("IN")
+OUT = TypeVar("OUT")
+MapTransformCallable = Callable[[Iterable[IN], TaskContext], Iterable[OUT]]
+
+
+class MapTransformFnDataType(Enum):
+    """An enum that represents the input/output data type of a MapTransformFn."""
+
+    Block = 0
+    Row = 1
+    Batch = 2
+
+
+class MapTransformFn(ABC):
+    """Represents a single transform function in a MapTransformer."""
+
+    def __init__(
+        self,
+        input_type: MapTransformFnDataType,
+        *,
+        is_udf: bool = False,
+        output_block_size_option: Optional[OutputBlockSizeOption] = None,
+    ):
+        """
+        Args:
+            input_type: Expected type of the input data.
+            is_udf: Whether this transformation is UDF or not.
+            output_block_size_option: (Optional) Output block size configuration.
+        """
+        self._input_type = input_type
+        self._output_block_size_option = output_block_size_option
+        self._is_udf = is_udf
+
+    @abstractmethod
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        pass
+
+    @abstractmethod
+    def _apply_transform(
+        self, ctx: TaskContext, inputs: Iterable[MapTransformFnData]
+    ) -> Iterable[MapTransformFnData]:
+        pass
+
+    def _pre_process(self, blocks: Iterable[Block]) -> Iterable[MapTransformFnData]:
+        return blocks
+
+    def _shape_blocks(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        buffer = BlockOutputBuffer(self._output_block_size_option)
+
+        # This method supports following modes of shaping of the output blocks:
+        #
+        #   1. Incremental: block is accumulated up to configured
+        #      ``_output_block_size_option``
+        #
+        #   2. *Non-incremental* (aka 1 block in / 1 block out): when
+        #      no ``OutputBlockSizeOption`` is provided this method will absorb
+        #      the whole input sequence and produce single block as an output
+        #
+        if self._input_type == MapTransformFnDataType.Block:
+            append = buffer.add_block
+        elif self._input_type == MapTransformFnDataType.Batch:
+            append = buffer.add_batch
+        else:
+            assert self._input_type == MapTransformFnDataType.Row
+            append = buffer.add
+
+        # Iterate over input sequence appending results to the
+        # buffer, while yielding incrementally
+        for result in results:
+            append(result)
+            # Try yielding incrementally
+            while buffer.has_next():
+                yield buffer.next()
+        # Finalize buffer
+        buffer.finalize()
+        # Yield remaining blocks from it
+        while buffer.has_next():
+            yield buffer.next()
+
+    def __call__(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> Iterable[Block]:
+        batches = self._pre_process(blocks)
+        results = self._apply_transform(ctx, batches)
+        yield from self._post_process(results)
+
+    @abstractmethod
+    def _can_skip_block_sizing(self):
+        pass
+
+    @property
+    def output_block_size_option(self):
+        return self._output_block_size_option
+
+    def override_target_max_block_size(self, target_max_block_size: Optional[int]):
+        self._output_block_size_option = OutputBlockSizeOption.of(
+            target_max_block_size=target_max_block_size
+        )
+
+    @property
+    def target_max_block_size(self):
+        if self._output_block_size_option is None:
+            return None
+        else:
+            return self._output_block_size_option.target_max_block_size
+
+    @property
+    def target_num_rows_per_block(self):
+        if self._output_block_size_option is None:
+            return None
+        else:
+            return self._output_block_size_option.target_num_rows_per_block
+
+
+class MapTransformer:
+    """Encapsulates the data transformation logic of a physical MapOperator.
+
+    A MapTransformer may consist of one or more steps, each of which is represented
+    as a MapTransformFn. The first MapTransformFn must take blocks as input, and
+    the last MapTransformFn must output blocks. The intermediate data types can
+    be blocks, rows, or batches.
+    """
+
+    def __init__(
+        self,
+        transform_fns: List[MapTransformFn],
+        *,
+        init_fn: Optional[Callable[[], None]] = None,
+        output_block_size_option_override: Optional[OutputBlockSizeOption] = None,
+    ):
+        """
+        Args:
+            transform_fns: A list of `MapTransformFn`s that will be executed sequentially
+                to transform data.
+            init_fn: A function that will be called before transforming data.
+                Used for the actor-based map operator.
+            output_block_size_option_override: (Optional) Output block size configuration.
+        """
+
+        self._transform_fns = []
+        self._init_fn = init_fn if init_fn is not None else lambda: None
+        self._output_block_size_option_override = output_block_size_option_override
+        self._udf_time = 0
+
+        # Add transformations
+        self.add_transform_fns(transform_fns)
+
+    def add_transform_fns(self, transform_fns: List[MapTransformFn]) -> None:
+        """Set the transform functions."""
+        assert len(transform_fns) > 0
+        self._transform_fns = self._combine_transformations(
+            self._transform_fns, transform_fns
+        )
+
+    def get_transform_fns(self) -> List[MapTransformFn]:
+        """Get the transform functions."""
+        return self._transform_fns
+
+    def override_target_max_block_size(self, target_max_block_size: Optional[int]):
+        self._output_block_size_option_override = OutputBlockSizeOption.of(
+            target_max_block_size=target_max_block_size
+        )
+
+    @property
+    def target_max_block_size_override(self) -> Optional[int]:
+        if self._output_block_size_option_override is None:
+            return None
+        else:
+            return self._output_block_size_option_override.target_max_block_size
+
+    def init(self) -> None:
+        """Initialize the transformer.
+
+        Should be called before applying the transform.
+        """
+        self._init_fn()
+
+    def _udf_timed_iter(
+        self, input: Iterable[MapTransformFnData]
+    ) -> Iterable[MapTransformFnData]:
+        while True:
+            try:
+                start = time.perf_counter()
+                output = next(input)
+                self._udf_time += time.perf_counter() - start
+                yield output
+            except StopIteration:
+                break
+
+    def apply_transform(
+        self,
+        input_blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> Iterable[Block]:
+        """Apply the transform functions to the input blocks."""
+
+        # NOTE: We only need to configure last transforming function to do
+        #       appropriate block sizing
+        last_transform = self._transform_fns[-1]
+
+        if self.target_max_block_size_override is not None:
+            last_transform.override_target_max_block_size(
+                self.target_max_block_size_override
+            )
+
+        iter = input_blocks
+        # Apply the transform functions sequentially to the input iterable.
+        for transform_fn in self._transform_fns:
+            iter = transform_fn(iter, ctx)
+            if transform_fn._is_udf:
+                iter = self._udf_timed_iter(iter)
+
+        return iter
+
+    def fuse(self, other: "MapTransformer") -> "MapTransformer":
+        """Fuse two `MapTransformer`s together."""
+        assert (
+            self.target_max_block_size_override == other.target_max_block_size_override
+            or (
+                self.target_max_block_size_override is None
+                or other.target_max_block_size_override is None
+            )
+        )
+        # Define them as standalone variables to avoid fused_init_fn capturing the
+        # entire `MapTransformer` object.
+        self_init_fn = self._init_fn
+        other_init_fn = other._init_fn
+
+        def fused_init_fn():
+            self_init_fn()
+            other_init_fn()
+
+        combined_transform_fns = self._combine_transformations(
+            self._transform_fns,
+            other._transform_fns,
+        )
+
+        transformer = MapTransformer(
+            combined_transform_fns,
+            init_fn=fused_init_fn,
+            output_block_size_option_override=OutputBlockSizeOption.of(
+                target_max_block_size=(
+                    self.target_max_block_size_override
+                    or other.target_max_block_size_override
+                ),
+            ),
+        )
+
+        return transformer
+
+    @classmethod
+    def _combine_transformations(
+        cls, ones: List[MapTransformFn], others: List[MapTransformFn]
+    ) -> list[Any]:
+        return ones + others
+
+    def udf_time(self) -> float:
+        return self._udf_time
+
+
+# Below are subclasses of MapTransformFn.
+
+
+class RowMapTransformFn(MapTransformFn):
+    """A rows-to-rows MapTransformFn."""
+
+    def __init__(
+        self,
+        row_fn: MapTransformCallable[Row, Row],
+        *,
+        is_udf: bool = False,
+        output_block_size_option: OutputBlockSizeOption,
+    ):
+        super().__init__(
+            input_type=MapTransformFnDataType.Row,
+            is_udf=is_udf,
+            output_block_size_option=output_block_size_option,
+        )
+
+        self._row_fn = row_fn
+
+    def _pre_process(self, blocks: Iterable[Block]) -> Iterable[MapTransformFnData]:
+        for block in blocks:
+            block = BlockAccessor.for_block(block)
+            for row in block.iter_rows(public_row_format=True):
+                yield row
+
+    def _apply_transform(
+        self, ctx: TaskContext, inputs: Iterable[MapTransformFnData]
+    ) -> Iterable[MapTransformFnData]:
+        yield from self._row_fn(inputs, ctx)
+
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        return self._shape_blocks(results)
+
+    def _can_skip_block_sizing(self):
+        return False
+
+    def __repr__(self) -> str:
+        return f"RowMapTransformFn({self._row_fn})"
+
+
+class BatchMapTransformFn(MapTransformFn):
+    """A batch-to-batch MapTransformFn."""
+
+    def __init__(
+        self,
+        batch_fn: MapTransformCallable[DataBatch, DataBatch],
+        *,
+        is_udf: bool = False,
+        batch_size: Optional[int] = None,
+        batch_format: Optional[BatchFormat] = None,
+        zero_copy_batch: bool = True,
+        output_block_size_option: Optional[OutputBlockSizeOption] = None,
+    ):
+        super().__init__(
+            input_type=MapTransformFnDataType.Batch,
+            is_udf=is_udf,
+            output_block_size_option=output_block_size_option,
+        )
+
+        self._batch_size = batch_size
+        self._batch_format = batch_format
+        self._zero_copy_batch = zero_copy_batch
+        self._ensure_copy = not zero_copy_batch and batch_size is not None
+
+        self._batch_fn = batch_fn
+
+    def _pre_process(self, blocks: Iterable[Block]) -> Iterable[MapTransformFnData]:
+        # TODO make batch-udf zero-copy by default
+        ensure_copy = not self._zero_copy_batch and self._batch_size is not None
+
+        return batch_blocks(
+            blocks=iter(blocks),
+            stats=None,
+            batch_size=self._batch_size,
+            batch_format=self._batch_format,
+            ensure_copy=ensure_copy,
+        )
+
+    def _apply_transform(
+        self, ctx: TaskContext, batches: Iterable[MapTransformFnData]
+    ) -> Iterable[MapTransformFnData]:
+        yield from self._batch_fn(batches, ctx)
+
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        return self._shape_blocks(results)
+
+    def _can_skip_block_sizing(self):
+        return self._output_block_size_option is None and self._batch_format in (
+            BatchFormat.ARROW,
+            BatchFormat.PANDAS,
+        )
+
+    def __repr__(self) -> str:
+        return f"BatchMapTransformFn({self._batch_fn=}, {self._batch_format=}, {self._batch_size=}, {self._zero_copy_batch=})"
+
+
+class BlockMapTransformFn(MapTransformFn):
+    """A block-to-block MapTransformFn."""
+
+    def __init__(
+        self,
+        block_fn: MapTransformCallable[Block, Block],
+        *,
+        is_udf: bool = False,
+        disable_block_shaping: bool = False,
+        output_block_size_option: Optional[OutputBlockSizeOption] = None,
+    ):
+        """
+        Initializes the object with a transformation function, accompanying options, and
+        configuration for handling blocks during processing.
+
+        Args:
+            block_fn: Callable function to apply a transformation to a block.
+            is_udf: Specifies if the transformation function is a user-defined
+                function (defaults to ``False``).
+            disable_block_shaping: Disables block-shaping, making transformer to
+                produce blocks as is.
+            output_block_size_option: (Optional) Configure output block sizing.
+        """
+
+        super().__init__(
+            input_type=MapTransformFnDataType.Block,
+            is_udf=is_udf,
+            output_block_size_option=output_block_size_option,
+        )
+
+        self._block_fn = block_fn
+        self._disable_block_shaping = disable_block_shaping
+
+    def _apply_transform(
+        self, ctx: TaskContext, blocks: Iterable[Block]
+    ) -> Iterable[Block]:
+        yield from self._block_fn(blocks, ctx)
+
+    def _post_process(self, results: Iterable[MapTransformFnData]) -> Iterable[Block]:
+        # Short-circuit for block transformations for which no
+        # block-shaping is required
+        if self._disable_block_shaping:
+            return results
+
+        return self._shape_blocks(results)
+
+    def _can_skip_block_sizing(self):
+        return self._output_block_size_option is None
+
+    def __repr__(self) -> str:
+        return (
+            f"BlockMapTransformFn({self._block_fn=}, {self._output_block_size_option=})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/output_splitter.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/output_splitter.py
new file mode 100644
index 0000000000000000000000000000000000000000..628350b0c366003faba4ea203c8bfed6e70edae5
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/output_splitter.py
@@ -0,0 +1,361 @@
+import math
+import time
+from collections import deque
+from typing import Any, Collection, Dict, List, Optional, Tuple
+
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    NodeIdStr,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+)
+from ray.data._internal.execution.util import locality_string
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+
+
+class OutputSplitter(InternalQueueOperatorMixin, PhysicalOperator):
+    """An operator that splits the given data into `n` output splits.
+
+    The output bundles of this operator will have a `bundle.output_split_idx` attr
+    set to an integer from [0..n-1]. This operator tries to divide the rows evenly
+    across output splits. If the `equal` option is set, the operator will furthermore
+    guarantee an exact split of rows across outputs, truncating the Dataset.
+
+    Implementation wise, this operator keeps an internal buffer of bundles. The buffer
+    has a minimum size calculated to enable a good locality hit rate, as well as ensure
+    we can satisfy the `equal` requirement.
+
+    OutputSplitter does not provide any ordering guarantees.
+    """
+
+    def __init__(
+        self,
+        input_op: PhysicalOperator,
+        n: int,
+        equal: bool,
+        data_context: DataContext,
+        locality_hints: Optional[List[NodeIdStr]] = None,
+    ):
+        super().__init__(
+            f"split({n}, equal={equal})",
+            [input_op],
+            data_context,
+        )
+        self._equal = equal
+        # Buffer of bundles not yet assigned to output splits.
+        self._buffer: List[RefBundle] = []
+        # The outputted bundles with output_split attribute set.
+        self._output_queue: deque[RefBundle] = deque()
+        # The number of rows output to each output split so far.
+        self._num_output: List[int] = [0 for _ in range(n)]
+        # The time of the overhead for the output splitter (operator level)
+        self._output_splitter_overhead_time = 0
+
+        if locality_hints is not None:
+            if n != len(locality_hints):
+                raise ValueError(
+                    "Locality hints list must have length `n`: "
+                    f"len({locality_hints}) != {n}"
+                )
+        self._locality_hints = locality_hints
+        if locality_hints:
+            # To optimize locality, we should buffer a certain number of elements
+            # internally before dispatch to allow the locality algorithm a good chance
+            # of selecting a preferred location. We use a small multiple of `n` since
+            # it's reasonable to buffer a couple blocks per consumer.
+            self._min_buffer_size = 2 * n
+        else:
+            self._min_buffer_size = 0
+        self._locality_hits = 0
+        self._locality_misses = 0
+
+    def num_outputs_total(self) -> Optional[int]:
+        # OutputSplitter does not change the number of blocks,
+        # so we can return the number of blocks from the input op.
+        return self.input_dependencies[0].num_outputs_total()
+
+    def num_output_rows_total(self) -> Optional[int]:
+        # The total number of rows is the same as the number of input rows.
+        return self.input_dependencies[0].num_output_rows_total()
+
+    def start(self, options: ExecutionOptions) -> None:
+        if options.preserve_order:
+            # If preserve_order is set, we need to ignore locality hints to ensure determinism.
+            self._locality_hints = None
+            self._min_buffer_size = 0
+
+        super().start(options)
+
+    def throttling_disabled(self) -> bool:
+        """Disables resource-based throttling.
+
+        It doesn't make sense to throttle the inputs to this operator, since all that
+        would do is lower the buffer size and prevent us from emitting outputs /
+        reduce the locality hit rate.
+        """
+        return True
+
+    def has_next(self) -> bool:
+        return len(self._output_queue) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        output = self._output_queue.popleft()
+        self._metrics.on_output_dequeued(output)
+        return output
+
+    def get_stats(self) -> StatsDict:
+        return {"split": []}  # TODO(ekl) add split metrics?
+
+    def _extra_metrics(self) -> Dict[str, Any]:
+        stats = {}
+        for i, num in enumerate(self._num_output):
+            stats[f"num_output_{i}"] = num
+        stats["output_splitter_overhead_time"] = self._output_splitter_overhead_time
+        return stats
+
+    def _add_input_inner(self, bundle, input_index) -> None:
+        if bundle.num_rows() is None:
+            raise ValueError("OutputSplitter requires bundles with known row count")
+        self._buffer.append(bundle)
+        self._metrics.on_input_queued(bundle)
+        self._dispatch_bundles()
+
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+        if not self._equal:
+            self._dispatch_bundles(dispatch_all=True)
+            assert not self._buffer, "Should have dispatched all bundles."
+            return
+
+        # Otherwise:
+        # Need to finalize distribution of buffered data to output splits.
+        buffer_size = sum(b.num_rows() for b in self._buffer)
+        max_n = max(self._num_output)
+
+        # First calculate the min rows to add per output to equalize them.
+        allocation = [max_n - n for n in self._num_output]
+        remainder = buffer_size - sum(allocation)
+        # Invariant: buffer should always be large enough to equalize.
+        assert remainder >= 0, (remainder, buffer_size, allocation)
+
+        # Equally distribute remaining rows in buffer to outputs.
+        x = remainder // len(allocation)
+        allocation = [a + x for a in allocation]
+
+        # Execute the split.
+        for i, count in enumerate(allocation):
+            bundles = self._split_from_buffer(count)
+            for b in bundles:
+                b.output_split_idx = i
+                self._output_queue.append(b)
+                self._metrics.on_output_queued(b)
+        self._buffer = []
+
+    def internal_input_queue_num_blocks(self) -> int:
+        return sum(len(b.block_refs) for b in self._buffer)
+
+    def internal_input_queue_num_bytes(self) -> int:
+        return sum(b.size_bytes() for b in self._buffer)
+
+    def internal_output_queue_num_blocks(self) -> int:
+        return sum(len(b.block_refs) for b in self._output_queue)
+
+    def internal_output_queue_num_bytes(self) -> int:
+        return sum(b.size_bytes() for b in self._output_queue)
+
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queue."""
+        while self._buffer:
+            bundle = self._buffer.pop()
+            self._metrics.on_input_dequeued(bundle)
+
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue."""
+        while self._output_queue:
+            bundle = self._output_queue.popleft()
+            self._metrics.on_output_dequeued(bundle)
+
+    def progress_str(self) -> str:
+        if self._locality_hints:
+            return locality_string(self._locality_hits, self._locality_misses)
+        else:
+            return "[locality disabled]"
+
+    def _dispatch_bundles(self, dispatch_all: bool = False) -> None:
+        start_time = time.perf_counter()
+        # Dispatch all dispatchable bundles from the internal buffer.
+        # This may not dispatch all bundles when equal=True.
+        while self._buffer and (
+            dispatch_all or len(self._buffer) >= self._min_buffer_size
+        ):
+            target_index = self._select_output_index()
+            target_bundle = self._pop_bundle_to_dispatch(target_index)
+            if self._can_safely_dispatch(target_index, target_bundle.num_rows()):
+                target_bundle.output_split_idx = target_index
+                self._num_output[target_index] += target_bundle.num_rows()
+                self._output_queue.append(target_bundle)
+                self._metrics.on_output_queued(target_bundle)
+                if self._locality_hints:
+                    preferred_loc = self._locality_hints[target_index]
+                    if preferred_loc in self._get_locations(target_bundle):
+                        self._locality_hits += 1
+                    else:
+                        self._locality_misses += 1
+            else:
+                # Put it back and abort.
+                self._buffer.insert(0, target_bundle)
+                self._metrics.on_input_queued(target_bundle)
+                break
+        self._output_splitter_overhead_time += time.perf_counter() - start_time
+
+    def _select_output_index(self) -> int:
+        # Greedily dispatch to the consumer with the least data so far.
+        i, _ = min(enumerate(self._num_output), key=lambda t: t[1])
+        return i
+
+    def _pop_bundle_to_dispatch(self, target_index: int) -> RefBundle:
+        if self._locality_hints:
+            preferred_loc = self._locality_hints[target_index]
+            for bundle in self._buffer:
+                if preferred_loc in self._get_locations(bundle):
+                    self._buffer.remove(bundle)
+                    self._metrics.on_input_dequeued(bundle)
+                    return bundle
+
+        bundle = self._buffer.pop(0)
+        self._metrics.on_input_dequeued(bundle)
+        return bundle
+
+    def _can_safely_dispatch(self, target_index: int, nrow: int) -> bool:
+        if not self._equal:
+            # If not in equals mode, dispatch away with no buffer requirements.
+            return True
+        output_distribution = self._num_output.copy()
+        output_distribution[target_index] += nrow
+        buffer_requirement = self._calculate_buffer_requirement(output_distribution)
+        buffer_size = sum(b.num_rows() for b in self._buffer)
+        return buffer_size >= buffer_requirement
+
+    def _calculate_buffer_requirement(self, output_distribution: List[int]) -> int:
+        # Calculate the new number of rows that we'd need to equalize the row
+        # distribution after the bundle dispatch.
+        max_n = max(output_distribution)
+        return sum([max_n - n for n in output_distribution])
+
+    def _split_from_buffer(self, nrow: int) -> List[RefBundle]:
+        output = []
+        acc = 0
+        while acc < nrow:
+            b = self._buffer.pop()
+            self._metrics.on_input_dequeued(b)
+            if acc + b.num_rows() <= nrow:
+                output.append(b)
+                acc += b.num_rows()
+            else:
+                left, right = _split(b, nrow - acc)
+                output.append(left)
+                acc += left.num_rows()
+                self._buffer.append(right)
+                self._metrics.on_input_queued(right)
+                assert acc == nrow, (acc, nrow)
+
+        assert sum(b.num_rows() for b in output) == nrow, (acc, nrow)
+        return output
+
+    @staticmethod
+    def _get_locations(bundle: RefBundle) -> Collection[NodeIdStr]:
+        """Fetches list of node ids holding the objects of the given bundle.
+
+        This method may be overridden for testing.
+
+        Returns:
+            A list of node ids where the objects in the bundle are located
+        """
+        preferred_locations = bundle.get_preferred_object_locations()
+
+        return preferred_locations.keys()
+
+    def implements_accurate_memory_accounting(self) -> bool:
+        return True
+
+
+def _split(bundle: RefBundle, left_size: int) -> Tuple[RefBundle, RefBundle]:
+    left_blocks, left_meta = [], []
+    right_blocks, right_meta = [], []
+    acc = 0
+    for b, m in bundle.blocks:
+        if acc >= left_size:
+            right_blocks.append(b)
+            right_meta.append(m)
+        elif acc + m.num_rows <= left_size:
+            left_blocks.append(b)
+            left_meta.append(m)
+            acc += m.num_rows
+        else:
+            # Trouble case: split it up.
+            lm, rm = _split_meta(m, left_size - acc)
+            lb, rb = _split_block(b, left_size - acc)
+            left_meta.append(lm)
+            right_meta.append(rm)
+            left_blocks.append(lb)
+            right_blocks.append(rb)
+            acc += lm.num_rows
+            assert acc == left_size
+    left = RefBundle(
+        list(zip(left_blocks, left_meta)),
+        owns_blocks=bundle.owns_blocks,
+        schema=bundle.schema,
+    )
+    right = RefBundle(
+        list(zip(right_blocks, right_meta)),
+        owns_blocks=bundle.owns_blocks,
+        schema=bundle.schema,
+    )
+    assert left.num_rows() == left_size
+    assert left.num_rows() + right.num_rows() == bundle.num_rows()
+    return left, right
+
+
+def _split_meta(
+    m: BlockMetadata, left_size: int
+) -> Tuple[BlockMetadata, BlockMetadata]:
+    left_bytes = int(math.floor(m.size_bytes * (left_size / m.num_rows)))
+    left = BlockMetadata(
+        num_rows=left_size,
+        size_bytes=left_bytes,
+        input_files=m.input_files,
+        exec_stats=None,
+    )
+    right = BlockMetadata(
+        num_rows=m.num_rows - left_size,
+        size_bytes=m.size_bytes - left_bytes,
+        input_files=m.input_files,
+        exec_stats=None,
+    )
+    return left, right
+
+
+def _split_block(
+    b: ObjectRef[Block], left_size: int
+) -> Tuple[ObjectRef[Block], ObjectRef[Block]]:
+    split_single_block = cached_remote_fn(_split_single_block)
+    left, right = split_single_block.options(num_cpus=0, num_returns=2).remote(
+        b, left_size
+    )
+    return left, right
+
+
+def _split_single_block(b: Block, left_size: int) -> Tuple[Block, Block]:
+    acc = BlockAccessor.for_block(b)
+    left = acc.slice(0, left_size)
+    right = acc.slice(left_size, acc.num_rows())
+    assert BlockAccessor.for_block(left).num_rows() == left_size
+    assert BlockAccessor.for_block(right).num_rows() == (acc.num_rows() - left_size)
+    return left, right
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/sub_progress.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/sub_progress.py
new file mode 100644
index 0000000000000000000000000000000000000000..8fe0e1935f38c48ce0e96d0474f0f55b9f06aabd
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/sub_progress.py
@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+
+class SubProgressBarMixin(ABC):
+    """Abstract class for operators that support sub-progress bars"""
+
+    @abstractmethod
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        """
+        Returns list of sub-progress bar names
+
+        This is used to create the sub-progress bars in the progress manager.
+        Note that sub-progress bars will be created in the order returned by
+        this method.
+        """
+        ...
+
+    @abstractmethod
+    def set_sub_progress_bar(self, name, pg):
+        """
+        Sets sub-progress bars
+
+        name: name of sub-progress bar
+        pg: SubProgressBar instance (progress_manager.py)
+        """
+        # Skipping type-checking for circular imports
+        ...
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/task_pool_map_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/task_pool_map_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..623e4b859ba2d9924f57c13a05057b5654449f6d
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/task_pool_map_operator.py
@@ -0,0 +1,177 @@
+import warnings
+from typing import Any, Callable, Dict, Optional
+
+from ray.data._internal.execution.interfaces import (
+    ExecutionResources,
+    PhysicalOperator,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.operators.map_operator import (
+    BaseRefBundler,
+    MapOperator,
+    _map_task,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data.context import DataContext
+
+
+class TaskPoolMapOperator(MapOperator):
+    """A MapOperator implementation that executes tasks on a task pool."""
+
+    def __init__(
+        self,
+        map_transformer: MapTransformer,
+        input_op: PhysicalOperator,
+        data_context: DataContext,
+        name: str = "TaskPoolMap",
+        target_max_block_size_override: Optional[int] = None,
+        min_rows_per_bundle: Optional[int] = None,
+        ref_bundler: Optional[BaseRefBundler] = None,
+        max_concurrency: Optional[int] = None,
+        supports_fusion: bool = True,
+        map_task_kwargs: Optional[Dict[str, Any]] = None,
+        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+    ):
+        """Create an TaskPoolMapOperator instance.
+
+        Args:
+            transform_fn: The function to apply to each ref bundle input.
+            input_op: Operator generating input data for this op.
+            name: The name of this operator.
+            target_max_block_size_override: Override for target max-block-size.
+            min_rows_per_bundle: The number of rows to gather per batch passed to the
+                transform_fn, or None to use the block size. Setting the batch size is
+                important for the performance of GPU-accelerated transform functions.
+                The actual rows passed may be less if the dataset is small.
+            ref_bundler: The ref bundler to use for this operator.
+            max_concurrency: The maximum number of Ray tasks to use concurrently,
+                or None to use as many tasks as possible.
+            supports_fusion: Whether this operator supports fusion with other operators.
+            map_task_kwargs: A dictionary of kwargs to pass to the map task. You can
+                access these kwargs through the `TaskContext.kwargs` dictionary.
+            ray_remote_args_fn: A function that returns a dictionary of remote args
+                passed to each map worker. The purpose of this argument is to generate
+                dynamic arguments for each actor/task, and will be called each time
+                prior to initializing the worker. Args returned from this dict will
+                always override the args in ``ray_remote_args``. Note: this is an
+                advanced, experimental feature.
+            ray_remote_args: Customize the :func:`ray.remote` args for this op's tasks.
+        """
+        super().__init__(
+            map_transformer,
+            input_op,
+            data_context,
+            name,
+            target_max_block_size_override,
+            min_rows_per_bundle,
+            ref_bundler,
+            supports_fusion,
+            map_task_kwargs,
+            ray_remote_args_fn,
+            ray_remote_args,
+        )
+
+        if max_concurrency is not None and max_concurrency <= 0:
+            raise ValueError(f"max_concurrency have to be > 0 (got {max_concurrency})")
+
+        self._max_concurrency = max_concurrency
+
+        # NOTE: Unlike static Ray remote args, dynamic arguments extracted from the
+        #       blocks themselves are going to be passed inside `fn.options(...)`
+        #       invocation
+        ray_remote_static_args = {
+            **(self._ray_remote_args or {}),
+            "num_returns": "streaming",
+            "_labels": {self._OPERATOR_ID_LABEL_KEY: self.id},
+        }
+
+        self._map_task = cached_remote_fn(_map_task, **ray_remote_static_args)
+
+    def _add_bundled_input(self, bundle: RefBundle):
+        # Submit the task as a normal Ray task.
+        ctx = TaskContext(
+            task_idx=self._next_data_task_idx,
+            op_name=self.name,
+            target_max_block_size_override=self.target_max_block_size_override,
+        )
+
+        dynamic_ray_remote_args = self._get_dynamic_ray_remote_args(input_bundle=bundle)
+        dynamic_ray_remote_args["name"] = self.name
+
+        if (
+            "_generator_backpressure_num_objects" not in dynamic_ray_remote_args
+            and self.data_context._max_num_blocks_in_streaming_gen_buffer is not None
+        ):
+            # The `_generator_backpressure_num_objects` parameter should be
+            # `2 * _max_num_blocks_in_streaming_gen_buffer` because we yield
+            # 2 objects for each block: the block and the block metadata.
+            dynamic_ray_remote_args["_generator_backpressure_num_objects"] = (
+                2 * self.data_context._max_num_blocks_in_streaming_gen_buffer
+            )
+
+        data_context = self.data_context
+
+        gen = self._map_task.options(**dynamic_ray_remote_args).remote(
+            self._map_transformer_ref,
+            data_context,
+            ctx,
+            *bundle.block_refs,
+            slices=bundle.slices,
+            **self.get_map_task_kwargs(),
+        )
+        self._submit_data_task(gen, bundle)
+
+    def progress_str(self) -> str:
+        return ""
+
+    def current_processor_usage(self) -> ExecutionResources:
+        num_active_workers = self.num_active_tasks()
+        return ExecutionResources(
+            cpu=self._ray_remote_args.get("num_cpus", 0) * num_active_workers,
+            gpu=self._ray_remote_args.get("num_gpus", 0) * num_active_workers,
+        )
+
+    def pending_processor_usage(self) -> ExecutionResources:
+        return ExecutionResources()
+
+    def incremental_resource_usage(self) -> ExecutionResources:
+        return self.per_task_resource_allocation().copy(
+            object_store_memory=(
+                self._metrics.obj_store_mem_max_pending_output_per_task or 0
+            ),
+        )
+
+    def per_task_resource_allocation(self) -> ExecutionResources:
+        return ExecutionResources(
+            cpu=self._ray_remote_args.get("num_cpus", 0),
+            gpu=self._ray_remote_args.get("num_gpus", 0),
+            memory=self._ray_remote_args.get("memory", 0),
+        )
+
+    def min_scheduling_resources(
+        self: "PhysicalOperator",
+    ) -> ExecutionResources:
+        return self.incremental_resource_usage()
+
+    def get_max_concurrency_limit(self) -> Optional[int]:
+        return self._max_concurrency
+
+    def all_inputs_done(self):
+        super().all_inputs_done()
+
+        if (
+            self._max_concurrency is not None
+            and self._metrics.num_inputs_received < self._max_concurrency
+        ):
+            warnings.warn(
+                f"The maximum number of concurrent tasks for '{self.name}' is set to "
+                f"{self._max_concurrency}, but the operator only received "
+                f"{self._metrics.num_inputs_received} input(s). This means that the "
+                f"operator can launch at most {self._metrics.num_inputs_received} "
+                "task(s), which is less than the concurrency limit. You might be able "
+                "to increase the number of concurrent tasks by configuring "
+                "`override_num_blocks` earlier in the pipeline."
+            )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/union_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/union_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4850c2d32c61d575f875998ea36b8d650e087b0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/union_operator.py
@@ -0,0 +1,139 @@
+import collections
+from typing import List, Optional
+
+from ray.data._internal.execution.bundle_queue import BundleQueue, FIFOBundleQueue
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+    NAryOperator,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext
+
+
+class UnionOperator(InternalQueueOperatorMixin, NAryOperator):
+    """An operator that combines output blocks from
+    two or more input operators into a single output."""
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: PhysicalOperator,
+    ):
+        """Create a UnionOperator.
+
+        Args:
+            input_ops: Operators generating input data for this operator to union.
+        """
+
+        # By default, union does not preserve the order of output blocks.
+        # To preserve the order, configure ExecutionOptions accordingly.
+        self._preserve_order = False
+
+        # Intermediary buffers used to store blocks from each input dependency.
+        # Only used when `self._prserve_order` is True.
+        self._input_buffers: List[BundleQueue] = [
+            FIFOBundleQueue() for _ in range(len(input_ops))
+        ]
+
+        # The index of the input dependency that is currently the source of
+        # the output buffer. New inputs from this input dependency will be added
+        # directly to the output buffer. Only used when `self._preserve_order` is True.
+        self._input_idx_to_output = 0
+
+        self._output_buffer: collections.deque[RefBundle] = collections.deque()
+        self._stats: StatsDict = {"Union": []}
+        super().__init__(data_context, *input_ops)
+
+    def start(self, options: ExecutionOptions):
+        # Whether to preserve the order of the input data (both the
+        # order of the input operators and the order of the blocks within).
+        self._preserve_order = options.preserve_order
+        super().start(options)
+
+    def num_outputs_total(self) -> Optional[int]:
+        num_outputs = 0
+        for input_op in self.input_dependencies:
+            input_num_outputs = input_op.num_outputs_total()
+            if input_num_outputs is None:
+                return None
+            num_outputs += input_num_outputs
+        return num_outputs
+
+    def num_output_rows_total(self) -> Optional[int]:
+        total_rows = 0
+        for input_op in self.input_dependencies:
+            input_num_rows = input_op.num_output_rows_total()
+            if input_num_rows is None:
+                return None
+            total_rows += input_num_rows
+        return total_rows
+
+    def internal_input_queue_num_blocks(self) -> int:
+        return sum(q.num_blocks() for q in self._input_buffers)
+
+    def internal_input_queue_num_bytes(self) -> int:
+        return sum(q.estimate_size_bytes() for q in self._input_buffers)
+
+    def internal_output_queue_num_blocks(self) -> int:
+        return sum(len(q.blocks) for q in self._output_buffer)
+
+    def internal_output_queue_num_bytes(self) -> int:
+        return sum(q.size_bytes() for q in self._output_buffer)
+
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queues."""
+        for input_buffer in self._input_buffers:
+            while input_buffer:
+                bundle = input_buffer.get_next()
+                self._metrics.on_input_dequeued(bundle)
+
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue."""
+        while self._output_buffer:
+            bundle = self._output_buffer.popleft()
+            self._metrics.on_output_dequeued(bundle)
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.completed()
+        assert 0 <= input_index <= len(self._input_dependencies), input_index
+
+        if not self._preserve_order:
+            self._output_buffer.append(refs)
+            self._metrics.on_output_queued(refs)
+        else:
+            self._input_buffers[input_index].add(refs)
+            self._metrics.on_input_queued(refs)
+
+    def all_inputs_done(self) -> None:
+        super().all_inputs_done()
+
+        if not self._preserve_order:
+            return
+
+        assert len(self._output_buffer) == 0, len(self._output_buffer)
+        for input_buffer in self._input_buffers:
+            while input_buffer:
+                refs = input_buffer.get_next()
+                self._metrics.on_input_dequeued(refs)
+                self._output_buffer.append(refs)
+                self._metrics.on_output_queued(refs)
+
+    def has_next(self) -> bool:
+        # Check if the output buffer still contains at least one block.
+        return len(self._output_buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        refs = self._output_buffer.popleft()
+        self._metrics.on_output_dequeued(refs)
+        return refs
+
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    def implements_accurate_memory_accounting(self):
+        return True
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/operators/zip_operator.py b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/zip_operator.py
new file mode 100644
index 0000000000000000000000000000000000000000..a024aa47a7fba84f9f0675d1d44aaf8abce019cf
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/operators/zip_operator.py
@@ -0,0 +1,329 @@
+import collections
+import itertools
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import ray
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import PhysicalOperator, RefBundle
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+    NAryOperator,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.split import _split_at_indices
+from ray.data._internal.stats import StatsDict
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockPartition,
+    to_stats,
+)
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+
+    from ray.data.block import BlockMetadataWithSchema
+
+
+class ZipOperator(InternalQueueOperatorMixin, NAryOperator):
+    """An operator that zips its inputs together.
+
+    NOTE: the implementation is bulk for now, which materializes all its inputs in
+    object store, before starting execution. Should re-implement it as a streaming
+    operator in the future.
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        *input_ops: PhysicalOperator,
+    ):
+        """Create a ZipOperator.
+
+        Args:
+            input_ops: Operators generating input data for this operator to zip.
+        """
+        assert len(input_ops) >= 2
+        self._input_buffers: List[collections.deque[RefBundle]] = [
+            collections.deque() for _ in range(len(input_ops))
+        ]
+        self._output_buffer: collections.deque[RefBundle] = collections.deque()
+        self._stats: StatsDict = {}
+        super().__init__(
+            data_context,
+            *input_ops,
+        )
+
+    def num_outputs_total(self) -> Optional[int]:
+        num_outputs = None
+        for input_op in self.input_dependencies:
+            input_num_outputs = input_op.num_outputs_total()
+            if input_num_outputs is None:
+                continue
+            if num_outputs is None:
+                num_outputs = input_num_outputs
+            else:
+                num_outputs = max(num_outputs, input_num_outputs)
+        return num_outputs
+
+    def num_output_rows_total(self) -> Optional[int]:
+        num_rows = None
+        for input_op in self.input_dependencies:
+            input_num_rows = input_op.num_output_rows_total()
+            if input_num_rows is None:
+                continue
+            if num_rows is None:
+                num_rows = input_num_rows
+            else:
+                num_rows = max(num_rows, input_num_rows)
+        return num_rows
+
+    def internal_input_queue_num_blocks(self) -> int:
+        return sum(
+            len(bundle.block_refs) for buf in self._input_buffers for bundle in buf
+        )
+
+    def internal_input_queue_num_bytes(self) -> int:
+        return sum(bundle.size_bytes() for buf in self._input_buffers for bundle in buf)
+
+    def internal_output_queue_num_blocks(self) -> int:
+        return sum(len(bundle.block_refs) for bundle in self._output_buffer)
+
+    def internal_output_queue_num_bytes(self) -> int:
+        return sum(bundle.size_bytes() for bundle in self._output_buffer)
+
+    def clear_internal_input_queue(self) -> None:
+        """Clear internal input queues."""
+        for input_buffer in self._input_buffers:
+            while input_buffer:
+                bundle = input_buffer.popleft()
+                self._metrics.on_input_dequeued(bundle)
+
+    def clear_internal_output_queue(self) -> None:
+        """Clear internal output queue."""
+        while self._output_buffer:
+            bundle = self._output_buffer.popleft()
+            self._metrics.on_output_dequeued(bundle)
+
+    def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
+        assert not self.completed()
+        assert 0 <= input_index <= len(self._input_dependencies), input_index
+        self._input_buffers[input_index].append(refs)
+        self._metrics.on_input_queued(refs)
+
+    def all_inputs_done(self) -> None:
+        assert len(self._output_buffer) == 0, len(self._output_buffer)
+
+        # Start with the first input buffer
+        while self._input_buffers[0]:
+            refs = self._input_buffers[0].popleft()
+            self._output_buffer.append(refs)
+            self._metrics.on_input_dequeued(refs)
+
+        # Process each additional input buffer
+        for input_buffer in self._input_buffers[1:]:
+            self._output_buffer, self._stats = self._zip(
+                self._output_buffer, input_buffer
+            )
+
+            # Clear the input buffer AFTER using it in _zip
+            while input_buffer:
+                refs = input_buffer.popleft()
+                self._metrics.on_input_dequeued(refs)
+
+        # Mark outputs as ready
+        for ref in self._output_buffer:
+            self._metrics.on_output_queued(ref)
+
+        super().all_inputs_done()
+
+    def has_next(self) -> bool:
+        return len(self._output_buffer) > 0
+
+    def _get_next_inner(self) -> RefBundle:
+        refs = self._output_buffer.popleft()
+        self._metrics.on_output_dequeued(refs)
+        return refs
+
+    def get_stats(self) -> StatsDict:
+        return self._stats
+
+    def implements_accurate_memory_accounting(self):
+        return True
+
+    def _zip(
+        self,
+        left_input: collections.deque[RefBundle],
+        right_input: collections.deque[RefBundle],
+    ) -> Tuple[collections.deque[RefBundle], StatsDict]:
+        """Zip the RefBundles from `left_input` and `right_input` together.
+
+        Zip is done in 2 steps: aligning blocks, and zipping blocks from
+        both sides.
+
+        Aligning blocks (optional): check the blocks from `left_input` and
+        `right_input` are aligned or not, i.e. if having different number of blocks, or
+        having different number of rows in some blocks. If not aligned, repartition the
+        smaller input with `_split_at_indices` to align with larger input.
+
+        Zipping blocks: after blocks from both sides are aligned, zip
+        blocks from both sides together in parallel.
+        """
+        left_blocks_with_metadata = []
+        for bundle in left_input:
+            for block, meta in bundle.blocks:
+                left_blocks_with_metadata.append((block, meta))
+        right_blocks_with_metadata = []
+        for bundle in right_input:
+            for block, meta in bundle.blocks:
+                right_blocks_with_metadata.append((block, meta))
+
+        left_block_rows, left_block_bytes = self._calculate_blocks_rows_and_bytes(
+            left_blocks_with_metadata
+        )
+        right_block_rows, right_block_bytes = self._calculate_blocks_rows_and_bytes(
+            right_blocks_with_metadata
+        )
+
+        # Check that both sides have the same number of rows.
+        # TODO(Clark): Support different number of rows via user-directed
+        # dropping/padding.
+        total_left_rows = sum(left_block_rows)
+        total_right_rows = sum(right_block_rows)
+        if total_left_rows != total_right_rows:
+            raise ValueError(
+                "Cannot zip datasets of different number of rows: "
+                f"{total_left_rows}, {total_right_rows}"
+            )
+
+        # Whether the left and right input sides are inverted
+        input_side_inverted = False
+        if sum(right_block_bytes) > sum(left_block_bytes):
+            # Make sure that right side is smaller, so we minimize splitting
+            # work when aligning both sides.
+            # TODO(Clark): Improve this heuristic for minimizing splitting work,
+            # e.g. by generating the splitting plans for each route (via
+            # _generate_per_block_split_indices) and choosing the plan that splits
+            # the least cumulative bytes.
+            left_blocks_with_metadata, right_blocks_with_metadata = (
+                right_blocks_with_metadata,
+                left_blocks_with_metadata,
+            )
+            left_block_rows, right_block_rows = right_block_rows, left_block_rows
+            input_side_inverted = True
+
+        # Get the split indices that will align both sides.
+        indices = list(itertools.accumulate(left_block_rows))
+        indices.pop(-1)
+
+        # Split other at the alignment indices, such that for every block from
+        # left side, we have a list of blocks from right side that have the same
+        # cumulative number of rows as that left block.
+        # NOTE: _split_at_indices has a no-op fastpath if the blocks are already
+        # aligned.
+        aligned_right_blocks_with_metadata = _split_at_indices(
+            right_blocks_with_metadata,
+            indices,
+            block_rows=right_block_rows,
+        )
+        del right_blocks_with_metadata
+
+        left_blocks = [b for b, _ in left_blocks_with_metadata]
+        right_blocks_list = aligned_right_blocks_with_metadata[0]
+        del left_blocks_with_metadata, aligned_right_blocks_with_metadata
+
+        zip_one_block = cached_remote_fn(_zip_one_block, num_returns=2)
+
+        output_blocks = []
+        output_metadata_schema = []
+        for left_block, right_blocks in zip(left_blocks, right_blocks_list):
+            # For each block from left side, zip it together with 1 or more blocks from
+            # right side. We're guaranteed to have that left_block has the same number
+            # of rows as right_blocks has cumulatively.
+            res, meta_with_schema = zip_one_block.remote(
+                left_block, *right_blocks, inverted=input_side_inverted
+            )
+            output_blocks.append(res)
+            output_metadata_schema.append(meta_with_schema)
+
+        # Early release memory.
+        del left_blocks, right_blocks_list
+
+        # TODO(ekl) it might be nice to have a progress bar here.
+        output_metadata_schema: List[BlockMetadataWithSchema] = ray.get(
+            output_metadata_schema
+        )
+
+        output_refs: collections.deque[RefBundle] = collections.deque()
+        input_owned = all(b.owns_blocks for b in left_input)
+        for block, meta_with_schema in zip(output_blocks, output_metadata_schema):
+            output_refs.append(
+                RefBundle(
+                    [
+                        (
+                            block,
+                            meta_with_schema.metadata,
+                        )
+                    ],
+                    owns_blocks=input_owned,
+                    schema=meta_with_schema.schema,
+                )
+            )
+        stats = {self._name: to_stats(output_metadata_schema)}
+
+        # Clean up inputs.
+        for ref in left_input:
+            ref.destroy_if_owned()
+        for ref in right_input:
+            ref.destroy_if_owned()
+
+        return output_refs, stats
+
+    def _calculate_blocks_rows_and_bytes(
+        self,
+        blocks_with_metadata: BlockPartition,
+    ) -> Tuple[List[int], List[int]]:
+        """Calculate the number of rows and size in bytes for a list of blocks with
+        metadata.
+        """
+        get_num_rows_and_bytes = cached_remote_fn(_get_num_rows_and_bytes)
+        block_rows = []
+        block_bytes = []
+        for block, metadata in blocks_with_metadata:
+            if metadata.num_rows is None or metadata.size_bytes is None:
+                # Need to fetch number of rows or size in bytes, so just fetch both.
+                num_rows, size_bytes = ray.get(get_num_rows_and_bytes.remote(block))
+                # Cache on the block metadata.
+                metadata.num_rows = num_rows
+                metadata.size_bytes = size_bytes
+            block_rows.append(metadata.num_rows)
+            block_bytes.append(metadata.size_bytes)
+        return block_rows, block_bytes
+
+
+def _zip_one_block(
+    block: Block, *other_blocks: Block, inverted: bool = False
+) -> Tuple[Block, "BlockMetadataWithSchema"]:
+    """Zip together `block` with `other_blocks`."""
+    stats = BlockExecStats.builder()
+    # Concatenate other blocks.
+    # TODO(Clark): Extend BlockAccessor.zip() to work with N other blocks,
+    # so we don't need to do this concatenation.
+    builder = DelegatingBlockBuilder()
+    for other_block in other_blocks:
+        builder.add_block(other_block)
+    other_block = builder.build()
+    if inverted:
+        # Swap blocks if ordering was inverted during block alignment splitting.
+        block, other_block = other_block, block
+    # Zip block and other blocks.
+    result = BlockAccessor.for_block(block).zip(other_block)
+    from ray.data.block import BlockMetadataWithSchema
+
+    return result, BlockMetadataWithSchema.from_block(result, stats=stats.build())
+
+
+def _get_num_rows_and_bytes(block: Block) -> Tuple[int, int]:
+    block = BlockAccessor.for_block(block)
+    return block.num_rows(), block.size_bytes()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/progress_manager.py b/lib/python3.12/site-packages/ray/data/_internal/execution/progress_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..e179bfd25434caee69e1b8ab3869c5ca2c5ec44f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/progress_manager.py
@@ -0,0 +1,495 @@
+import logging
+import math
+import sys
+import time
+import uuid
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, List, Optional
+
+from ray.data._internal.execution.interfaces.physical_operator import PhysicalOperator
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.streaming_executor_state import OpState, Topology
+from ray.data._internal.progress_bar import AbstractProgressBar, truncate_operator_name
+from ray.util.debug import log_once
+
+try:
+    import rich
+    from rich.console import Console
+    from rich.live import Live
+    from rich.progress import (
+        BarColumn,
+        Progress,
+        SpinnerColumn,
+        TextColumn,
+        TimeElapsedColumn,
+    )
+    from rich.table import Column, Table
+    from rich.text import Text
+
+    needs_rich_warning = False
+except ImportError:
+    rich = None
+    needs_rich_warning = True
+
+logger = logging.getLogger(__name__)
+
+_TREE_BRANCH = "  ├─"
+_TREE_VERTICAL = "│"
+_TREE_VERTICAL_SUB_PROGRESS = "  │     -"
+_TREE_VERTICAL_INDENT = f"  {_TREE_VERTICAL}    "
+_TOTAL_PROGRESS_TOTAL = 1.0
+_RESOURCE_REPORT_HEADER = f"  {_TREE_VERTICAL} Active/total resources: "
+
+
+class _ManagerMode(str, Enum):
+    NONE = "NONE"  # no-op
+    GLOBAL_ONLY = "GLOBAL_ONLY"  # global progress
+    ALL = "ALL"  # show everything
+
+    def show_op(self) -> bool:
+        return self == self.ALL
+
+    def is_enabled(self) -> bool:
+        return self != self.NONE
+
+    @classmethod
+    def get_mode(cls) -> "_ManagerMode":
+        from ray.data.context import DataContext
+
+        ctx = DataContext.get_current()
+        if not ctx.enable_progress_bars:
+            if log_once("ray_data_progress_manager_disabled"):
+                logger.warning(
+                    "Progress bars disabled. To enable, set "
+                    "`ray.data.DataContext.get_current()."
+                    "enable_progress_bars = True`."
+                )
+            return cls.NONE
+        elif rich is None:
+            global needs_rich_warning
+            if needs_rich_warning:
+                print(
+                    "[dataset]: Run `pip install rich` to enable "
+                    "execution progress reporting."
+                )
+                needs_rich_warning = False
+            return cls.NONE
+        elif not ctx.enable_operator_progress_bars:
+            if log_once("ray_data_progress_manager_global"):
+                logger.warning(
+                    "Progress bars for operators disabled. To enable, "
+                    "set `ray.data.DataContext.get_current()."
+                    "enable_operator_progress_bars = True`."
+                )
+            return cls.GLOBAL_ONLY
+        else:
+            return cls.ALL
+
+
+class SubProgressBar(AbstractProgressBar):
+    """Thin wrapper to provide identical interface to the ProgressBar.
+
+    Updates RichExecutionProgressManager internally.
+    """
+
+    # If the name/description of the progress bar exceeds this length,
+    # it will be truncated.
+    MAX_NAME_LENGTH = 100
+
+    def __init__(
+        self,
+        name: str,
+        total: Optional[int] = None,
+        enabled: bool = True,
+        progress: Optional[Any] = None,
+        tid: Optional[Any] = None,
+    ):
+        """
+        Initialize sub-progress bar
+
+        Args:
+            name: name of sub-progress bar
+            total: total number of output rows. None for unknown.
+            enabled: whether progress bar is enabled.
+            progress: rich.Progress instance for the corresponding
+                sub-progress bar.
+            tid: rich.TaskId for the corresponding sub-progress bar task.
+        """
+        # progress, tid type Optional[Any] due to conditional rich import.
+        if enabled:
+            assert progress is not None and tid is not None
+        else:
+            progress = None
+            tid = None
+        self._total = total
+        self._completed = 0
+        self._start_time = None
+        self._enabled = enabled
+        self._progress = progress
+        self._tid = tid
+        self._desc = truncate_operator_name(name, self.MAX_NAME_LENGTH)
+
+    def set_description(self, name: str) -> None:
+        self._desc = truncate_operator_name(name, self.MAX_NAME_LENGTH)
+        if self._enabled:
+            self._progress.update(self._tid, description=self._desc)
+
+    def get_description(self) -> str:
+        return self._desc
+
+    def _update(self, completed: int, total: Optional[int] = None) -> None:
+        assert self._enabled
+        if self._start_time is None:
+            self._start_time = time.time()
+        metrics = _get_progress_metrics(self._start_time, completed, total)
+        self._progress.update(
+            self._tid,
+            completed=metrics.completed,
+            total=metrics.total,
+            rate_str=metrics.rate_str,
+            count_str=metrics.count_str,
+        )
+
+    def update(self, increment: int = 0, total: Optional[int] = None) -> None:
+        if self._enabled and increment != 0:
+            if total is not None:
+                self._total = total
+            self._completed += increment
+            self._update(self._completed, self._total)
+
+    def complete(self) -> None:
+        if self._enabled:
+            self._update(self._completed, self._completed)
+
+    def __getstate__(self):
+        return {}
+
+    def __setstate__(self, state):
+        self.enabled = False  # Progress bar is disabled on remote nodes.
+
+
+class RichExecutionProgressManager:
+    """Execution progress display using rich."""
+
+    # If the name/description of the progress bar exceeds this length,
+    # it will be truncated.
+    MAX_NAME_LENGTH = 100
+
+    def __init__(self, dataset_id: str, topology: Topology):
+        self._mode = _ManagerMode.get_mode()
+        self._dataset_id = dataset_id
+        self._sub_progress_bars: List[SubProgressBar] = []
+
+        if not self._mode.is_enabled():
+            self._live = None
+            # TODO (kyuds): for sub-progress, initialize no-op
+            for state in topology.values():
+                if _has_sub_progress_bars(state.op):
+                    self._setup_operator_sub_progress(state)
+            return
+
+        self._start_time: Optional[float] = None
+
+        # rich
+        self._console = Console(file=sys.stderr)
+        self._total = self._make_progress_bar(" ", "•", 15)
+        self._current_rows = 0
+        self._total_resources = Text(
+            f"{_RESOURCE_REPORT_HEADER}Initializing...", no_wrap=True
+        )
+
+        self._op_display = {}
+
+        self._layout_table = Table.grid(padding=(0, 1, 0, 0), expand=True)
+        self._layout_table.add_row(self._total)
+        self._layout_table.add_row(self._total_resources)
+
+        self._setup_progress_grid(topology)
+
+        # empty new line to prevent "packed" feeling
+        self._layout_table.add_row(Text())
+
+        # rich.Live is the auto-refreshing rich component display.
+        # refreshing/closing is all done through rich.Live
+        self._live = Live(
+            self._layout_table,
+            console=self._console,
+            refresh_per_second=2,
+            vertical_overflow="visible",
+        )
+
+        self._total_task_id = self._total.add_task(
+            f"Dataset {self._dataset_id} running:",
+            total=_TOTAL_PROGRESS_TOTAL,
+            rate_str="? rows/s",
+            count_str="0/?",
+        )
+
+    def _setup_progress_grid(self, topology: Topology):
+        if self._mode.show_op():
+            self._layout_table.add_row(Text(f"  {_TREE_VERTICAL}", no_wrap=True))
+        for state in topology.values():
+            if isinstance(state.op, InputDataBuffer):
+                continue
+            if self._mode.show_op():
+                uid = uuid.uuid4()
+                progress = self._make_progress_bar(_TREE_BRANCH, " ", 10)
+                stats = Text(f"{_TREE_VERTICAL_INDENT}Initializing...", no_wrap=True)
+                total = state.op.num_output_rows_total()
+                name = truncate_operator_name(state.op.name, self.MAX_NAME_LENGTH)
+                tid = progress.add_task(
+                    name,
+                    total=total if total is not None else 1,
+                    start=True,
+                    rate_str="? rows/s",
+                    count_str="0/?",
+                )
+                self._layout_table.add_row(progress)
+                self._layout_table.add_row(stats)
+                state.progress_manager_uuid = uid
+                self._op_display[uid] = (tid, progress, stats)
+
+            if _has_sub_progress_bars(state.op):
+                self._setup_operator_sub_progress(state)
+
+    def _setup_operator_sub_progress(self, state: OpState):
+        assert _has_sub_progress_bars(
+            state.op
+        ), f"Operator {state.op.name} doesn't support sub-progress bars."
+        enabled = self._mode.show_op()
+
+        sub_progress_bar_names = state.op.get_sub_progress_bar_names()
+        if sub_progress_bar_names is not None:
+            for name in sub_progress_bar_names:
+                name = truncate_operator_name(name, SubProgressBar.MAX_NAME_LENGTH)
+                progress = None
+                tid = None
+                total = None
+
+                if enabled:
+                    progress = self._make_progress_bar(
+                        _TREE_VERTICAL_SUB_PROGRESS, "", 10
+                    )
+                    total = state.op.num_output_rows_total()
+                    tid = progress.add_task(
+                        name,
+                        total=total if total is not None else 1,
+                        start=True,
+                        rate_str="? rows/s",
+                        count_str="0/?",
+                    )
+                    self._layout_table.add_row(progress)
+
+                pg = SubProgressBar(
+                    name=name,
+                    total=total,
+                    enabled=enabled,
+                    progress=progress,
+                    tid=tid,
+                )
+                state.op.set_sub_progress_bar(name, pg)
+                self._sub_progress_bars.append(pg)
+
+    def _make_progress_bar(self, indent_str, spinner_finish, bar_width):
+        # no type hints because rich import is conditional.
+        assert self._mode.is_enabled()
+        return Progress(
+            TextColumn(indent_str, table_column=Column(no_wrap=True)),
+            SpinnerColumn(finished_text=spinner_finish),
+            TextColumn(
+                "{task.description} {task.percentage:>3.0f}%",
+                table_column=Column(no_wrap=True),
+            ),
+            BarColumn(bar_width=bar_width),
+            TextColumn("{task.fields[count_str]}", table_column=Column(no_wrap=True)),
+            TextColumn("["),
+            TimeElapsedColumn(),
+            TextColumn(","),
+            TextColumn("{task.fields[rate_str]}", table_column=Column(no_wrap=True)),
+            TextColumn("]"),
+            console=self._console,
+            transient=False,
+            expand=False,
+        )
+
+    # Management
+    def start(self):
+        if self._mode.is_enabled():
+            if not self._live.is_started:
+                self._live.start()
+
+    def refresh(self):
+        if self._mode.is_enabled():
+            if self._live.is_started:
+                self._live.refresh()
+
+    def close_with_finishing_description(self, desc: str, success: bool):
+        if self._mode.is_enabled():
+            if self._live.is_started:
+                kwargs = {}
+                if success:
+                    # set everything to completed
+                    kwargs["completed"] = 1.0
+                    kwargs["total"] = 1.0
+                    for pg in self._sub_progress_bars:
+                        pg.complete()
+                    for tid, progress, _ in self._op_display.values():
+                        completed = progress.tasks[tid].completed or 0
+                        metrics = _get_progress_metrics(
+                            self._start_time, completed, completed
+                        )
+                        _update_with_conditional_rate(progress, tid, metrics)
+                self._total.update(self._total_task_id, description=desc, **kwargs)
+                self.refresh()
+                time.sleep(0.02)
+                self._live.stop()
+                logger.info(desc)
+
+    # Total Progress
+    def _can_update_total(self) -> bool:
+        return (
+            self._mode.is_enabled()
+            and self._total_task_id is not None
+            and self._total_task_id in self._total.task_ids
+        )
+
+    def update_total_progress(self, new_rows: int, total_rows: Optional[int]):
+        if not self._can_update_total():
+            return
+        if self._live.is_started:
+            if self._start_time is None:
+                self._start_time = time.time()
+            if new_rows is not None:
+                self._current_rows += new_rows
+            metrics = _get_progress_metrics(
+                self._start_time, self._current_rows, total_rows
+            )
+            _update_with_conditional_rate(self._total, self._total_task_id, metrics)
+
+    def update_resource_status(self, resource_status: str):
+        if not self._can_update_total():
+            return
+        if self._live.is_started:
+            self._total_resources.plain = _RESOURCE_REPORT_HEADER + resource_status
+
+    def _can_update_operator(self, op_state: OpState) -> bool:
+        if not self._mode.show_op():
+            return False
+        uid = op_state.progress_manager_uuid
+        if uid is None or uid not in self._op_display:
+            return False
+        tid, progress, stats = self._op_display[uid]
+        if tid is None or not progress or not stats or tid not in progress.task_ids:
+            return False
+        return True
+
+    def update_operator_progress(self, op_state: OpState):
+        if not self._can_update_operator(op_state):
+            return
+        if self._start_time is None:
+            self._start_time = time.time()
+        uid = op_state.progress_manager_uuid
+        tid, progress, stats = self._op_display[uid]
+
+        # progress
+        current_rows = op_state.output_row_count
+        total_rows = op_state.op.num_output_rows_total()
+        metrics = _get_progress_metrics(self._start_time, current_rows, total_rows)
+        _update_with_conditional_rate(progress, tid, metrics)
+        # stats
+        stats_str = op_state.op_display_metrics.display_str()
+        stats.plain = f"{_TREE_VERTICAL_INDENT}{stats_str}"
+
+
+# utilities
+def _format_k(val: int) -> str:
+    if val >= 1000:
+        fval = val / 1000.0
+        fval_str = f"{int(fval)}" if fval.is_integer() else f"{fval:.2f}"
+        return fval_str + "k"
+    return str(val)
+
+
+def _format_row_count(completed: int, total: Optional[int]) -> str:
+    """Formats row counts with k units."""
+    cstr = _format_k(completed)
+    if total is None or math.isinf(total):
+        tstr = "?k" if cstr.endswith("k") else "?"
+    else:
+        tstr = _format_k(total)
+    return f"{cstr}/{tstr}"
+
+
+@dataclass
+class _ProgressMetrics:
+    completed: int
+    total: int
+    rate_str: str
+    count_str: str
+
+
+def _get_progress_metrics(
+    start_time: float, completed_rows: int, total_rows: Optional[int]
+) -> _ProgressMetrics:
+    """
+    Args:
+        start_time: time when progress tracking started
+        completed_rows: cumulative rows outputted
+        total_rows: total rows expected (can be unknown)
+    Returns:
+        _ProgressMetrics instance containing the calculated data.
+    """
+    # Note, when total is unknown, we default the progress bar to 0.
+    # Will properly have estimates for rate and count strings though.
+    total = 1 if total_rows is None or total_rows < 1 else total_rows
+    completed = 0 if total_rows is None else completed_rows
+
+    if total_rows is None:
+        rate_str = "? row/s"
+    else:
+        elapsed = time.time() - start_time
+        rate_val = completed_rows / elapsed if elapsed > 1 else 0
+        rate_unit = "row/s"
+        if rate_val >= 1000:
+            rate_val /= 1000
+            rate_unit = "k row/s"
+        rate_str = f"{rate_val:.2f} {rate_unit}"
+    count_str = _format_row_count(completed_rows, total_rows)
+
+    return _ProgressMetrics(
+        completed=completed, total=total, rate_str=rate_str, count_str=count_str
+    )
+
+
+def _has_sub_progress_bars(op: PhysicalOperator) -> bool:
+    """Determines if operator implements sub-progress bars
+
+    Args:
+        op: Operator
+    Returns:
+        whether operator implements sub-progress bars
+    """
+    # function primarily used to avoid circular imports
+    from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
+
+    return isinstance(op, SubProgressBarMixin)
+
+
+def _update_with_conditional_rate(progress, tid, metrics):
+    # not doing type checking because rich is imported conditionally.
+    # progress: rich.Progress
+    # tid: rich.TaskId
+    # metrics: _ProgressMetrics
+    task = progress.tasks[tid]
+    kwargs = {
+        "completed": metrics.completed,
+        "total": metrics.total,
+        "count_str": metrics.count_str,
+    }
+    if task.completed != metrics.completed:
+        # update rate string only if there are new rows.
+        # this allows updates to other metric data while
+        # preserving the right rate notation.
+        kwargs["rate_str"] = metrics.rate_str
+    progress.update(tid, **kwargs)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/ranker.py b/lib/python3.12/site-packages/ray/data/_internal/execution/ranker.py
new file mode 100644
index 0000000000000000000000000000000000000000..d45af0701af62f9f2b6f84b9ef5a542f5ea47131
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/ranker.py
@@ -0,0 +1,100 @@
+"""Ranker component for operator selection in streaming executor."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, List, Protocol, Tuple, TypeVar
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.resource_manager import ResourceManager
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+# Protocol for comparable ranking values
+class Comparable(Protocol):
+    """Protocol for types that can be compared for ranking."""
+
+    def __lt__(self, other: "Comparable") -> bool:
+        ...
+
+    def __le__(self, other: "Comparable") -> bool:
+        ...
+
+    def __gt__(self, other: "Comparable") -> bool:
+        ...
+
+    def __ge__(self, other: "Comparable") -> bool:
+        ...
+
+    def __eq__(self, other: "Comparable") -> bool:
+        ...
+
+
+# Generic type for comparable ranking values
+RankingValue = TypeVar("RankingValue", bound=Comparable)
+
+
+class Ranker(ABC, Generic[RankingValue]):
+    """Abstract base class for operator ranking strategies."""
+
+    @abstractmethod
+    def rank_operator(
+        self,
+        op: PhysicalOperator,
+        topology: "Topology",
+        resource_manager: ResourceManager,
+    ) -> RankingValue:
+        """Rank operator for execution priority.
+
+        Operator to run next is selected as the one with the *smallest* value
+        of the lexicographically ordered ranks composed of (in order):
+
+        Args:
+            op: Operator to rank
+            topology: Current execution topology
+            resource_manager: Resource manager for usage information
+
+        Returns:
+            Rank (tuple) for operator
+        """
+        pass
+
+    def rank_operators(
+        self,
+        ops: List[PhysicalOperator],
+        topology: "Topology",
+        resource_manager: ResourceManager,
+    ) -> List[RankingValue]:
+
+        assert len(ops) > 0
+        return [self.rank_operator(op, topology, resource_manager) for op in ops]
+
+
+class DefaultRanker(Ranker[Tuple[int, int]]):
+    """Ranker implementation."""
+
+    def rank_operator(
+        self,
+        op: PhysicalOperator,
+        topology: "Topology",
+        resource_manager: ResourceManager,
+    ) -> Tuple[int, int]:
+        """Computes rank for op. *Lower means better rank*
+
+            1. Whether operator's could be throttled (int)
+            2. Operators' object store utilization
+
+        Args:
+            op: Operator to rank
+            topology: Current execution topology
+            resource_manager: Resource manager for usage information
+
+        Returns:
+            Rank (tuple) for operator
+        """
+
+        throttling_disabled = 0 if op.throttling_disabled() else 1
+
+        return (
+            throttling_disabled,
+            resource_manager.get_op_usage(op).object_store_memory,
+        )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/resource_manager.py b/lib/python3.12/site-packages/ray/data/_internal/execution/resource_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..58f7b6837c36b770afebe192ec268c32c38169a2
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/resource_manager.py
@@ -0,0 +1,918 @@
+import logging
+import math
+import time
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional
+
+from ray._private.ray_constants import env_bool, env_float
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    PhysicalOperator,
+    ReportsExtraResourceUsage,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+)
+from ray.data._internal.execution.operators.hash_shuffle import (
+    HashShufflingOperatorBase,
+)
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.operators.zip_operator import ZipOperator
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.util import GiB
+from ray.data.context import DataContext
+from ray.util.debug import log_once
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor_state import OpState, Topology
+
+
+logger = logging.getLogger(__name__)
+
+
+LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE: Optional[bool] = env_bool(
+    "RAY_DATA_DEBUG_RESOURCE_MANAGER", None
+)
+
+
+# These are physical operators that must receive all inputs before they start
+# processing data.
+MATERIALIZING_OPERATORS = (AllToAllOperator, ZipOperator)
+
+
+class ResourceManager:
+    """A class that manages the resource usage of a streaming executor."""
+
+    # The interval in seconds at which the global resource limits are refreshed.
+    GLOBAL_LIMITS_UPDATE_INTERVAL_S = 1
+
+    # The fraction of the object store capacity that will be used as the default object
+    # store memory limit for the streaming executor,
+    # when `ReservationOpResourceAllocator` is enabled.
+    DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION = env_float(
+        "RAY_DATA_OBJECT_STORE_MEMORY_LIMIT_FRACTION", 0.5
+    )
+
+    # The fraction of the object store capacity that will be used as the default object
+    # store memory limit for the streaming executor,
+    # when `ReservationOpResourceAllocator` is not enabled.
+    DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION = 0.25
+
+    def __init__(
+        self,
+        topology: "Topology",
+        options: ExecutionOptions,
+        get_total_resources: Callable[[], ExecutionResources],
+        data_context: DataContext,
+    ):
+        self._topology = topology
+        self._options = options
+        self._get_total_resources = get_total_resources
+        self._global_limits = ExecutionResources.zero()
+        self._global_limits_last_update_time = 0
+        self._global_usage = ExecutionResources.zero()
+        self._global_running_usage = ExecutionResources.zero()
+        self._global_pending_usage = ExecutionResources.zero()
+        self._op_usages: Dict[PhysicalOperator, ExecutionResources] = {}
+        self._op_running_usages: Dict[PhysicalOperator, ExecutionResources] = {}
+        self._op_pending_usages: Dict[PhysicalOperator, ExecutionResources] = {}
+        # Object store memory usage internal to the operator, including the
+        # pending task outputs and op's internal output buffers.
+        self._mem_op_internal: Dict[PhysicalOperator, int] = defaultdict(int)
+        # Object store memory usage of the blocks that have been taken out of
+        # the operator, including the external output buffer in OpState, and the
+        # input buffers of the downstream operators.
+        self._mem_op_outputs: Dict[PhysicalOperator, int] = defaultdict(int)
+
+        self._op_resource_allocator: Optional["OpResourceAllocator"] = None
+
+        if data_context.op_resource_reservation_enabled:
+            # We'll enable memory reservation if all operators have
+            # implemented accurate memory accounting.
+            should_enable = all(
+                op.implements_accurate_memory_accounting() for op in topology
+            )
+            if should_enable:
+                self._op_resource_allocator = ReservationOpResourceAllocator(
+                    self, data_context.op_resource_reservation_ratio
+                )
+
+        self._object_store_memory_limit_fraction = (
+            data_context.override_object_store_memory_limit_fraction
+            if data_context.override_object_store_memory_limit_fraction is not None
+            else (
+                self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION
+                if self.op_resource_allocator_enabled()
+                else self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION
+            )
+        )
+
+        self._warn_about_object_store_memory_if_needed()
+
+    def _warn_about_object_store_memory_if_needed(self):
+        """Warn if object store memory is configured below 50% of total memory."""
+        import ray
+        from ray.data.context import WARN_PREFIX
+        from ray.util.debug import log_once
+
+        if not ray.is_initialized():
+            return
+
+        cluster_resources = ray.cluster_resources()
+        total_memory = cluster_resources.get("memory", 0)
+        object_store_memory = cluster_resources.get("object_store_memory", 0)
+
+        # Check if we have actual numeric values (not mocks or None)
+        if total_memory > 0:
+            object_store_fraction = object_store_memory / total_memory
+
+            if object_store_fraction < 0.5 and log_once(
+                "ray_data_object_store_memory_warning"
+            ):
+                logger.warning(
+                    f"{WARN_PREFIX} Ray's object store is configured to use only "
+                    f"{object_store_fraction:.1%} of available memory ({object_store_memory/GiB:.1f}GiB "
+                    f"out of {total_memory/GiB:.1f}GiB total). For optimal Ray Data performance, "
+                    f"we recommend setting the object store to at least 50% of available memory. "
+                    f"You can do this by setting the 'object_store_memory' parameter when calling "
+                    f"ray.init() or by setting the RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION environment variable."
+                )
+
+    def _estimate_object_store_memory(
+        self, op: "PhysicalOperator", state: "OpState"
+    ) -> int:
+        # Don't count input refs towards dynamic memory usage, as they have been
+        # pre-created already outside this execution.
+        if isinstance(op, InputDataBuffer):
+            return 0
+
+        # Pending task outputs.
+        mem_op_internal = op.metrics.obj_store_mem_pending_task_outputs or 0
+        # Op's internal output buffers.
+        mem_op_internal += op.metrics.obj_store_mem_internal_outqueue
+
+        # Op's external output buffer.
+        mem_op_outputs = state.output_queue_bytes()
+        # Input buffers of the downstream operators.
+        for next_op in op.output_dependencies:
+            mem_op_outputs += (
+                next_op.metrics.obj_store_mem_internal_inqueue
+                + next_op.metrics.obj_store_mem_pending_task_inputs
+            )
+
+        self._mem_op_internal[op] = mem_op_internal
+        self._mem_op_outputs[op] = mem_op_outputs
+
+        return mem_op_internal + mem_op_outputs
+
+    def update_usages(self):
+        """Recalculate resource usages."""
+        # TODO(hchen): This method will be called frequently during the execution loop.
+        # And some computations are redundant. We should either remove redundant
+        # computations or remove this method entirely and compute usages on demand.
+        self._global_usage = ExecutionResources(0, 0, 0)
+        self._global_running_usage = ExecutionResources(0, 0, 0)
+        self._global_pending_usage = ExecutionResources(0, 0, 0)
+        self._op_usages.clear()
+        self._op_running_usages.clear()
+        self._op_pending_usages.clear()
+
+        # Iterate from last to first operator.
+        for op, state in reversed(self._topology.items()):
+            # Update `self._op_usages`, `self._op_running_usages`,
+            # and `self._op_pending_usages`.
+            op.update_resource_usage()
+            op_usage = op.current_processor_usage()
+            op_running_usage = op.running_processor_usage()
+            op_pending_usage = op.pending_processor_usage()
+
+            assert not op_usage.object_store_memory
+            assert not op_running_usage.object_store_memory
+            assert not op_pending_usage.object_store_memory
+
+            used_object_store = self._estimate_object_store_memory(op, state)
+
+            op_usage = op_usage.copy(object_store_memory=used_object_store)
+            op_running_usage = op_running_usage.copy(
+                object_store_memory=used_object_store
+            )
+
+            if isinstance(op, ReportsExtraResourceUsage):
+                op_usage.add(op.extra_resource_usage())
+
+            self._op_usages[op] = op_usage
+            self._op_running_usages[op] = op_running_usage
+            self._op_pending_usages[op] = op_pending_usage
+
+            # Update `self._global_usage`, `self._global_running_usage`,
+            # and `self._global_pending_usage`.
+            self._global_usage = self._global_usage.add(op_usage)
+            self._global_running_usage = self._global_running_usage.add(
+                op_running_usage
+            )
+            self._global_pending_usage = self._global_pending_usage.add(
+                op_pending_usage
+            )
+
+            # Update operator's object store usage, which is used by
+            # DatasetStats and updated on the Ray Data dashboard.
+            op._metrics.obj_store_mem_used = op_usage.object_store_memory
+
+        if self._op_resource_allocator is not None:
+            self._op_resource_allocator.update_budgets(
+                limits=self._global_limits,
+            )
+
+    def get_global_usage(self) -> ExecutionResources:
+        """Return the global resource usage at the current time."""
+        return self._global_usage
+
+    def get_global_running_usage(self) -> ExecutionResources:
+        """Return the global running resource usage at the current time."""
+        return self._global_running_usage
+
+    def get_global_pending_usage(self) -> ExecutionResources:
+        """Return the global pending resource usage at the current time."""
+        return self._global_pending_usage
+
+    def get_global_limits(self) -> ExecutionResources:
+        """Return the global resource limits at the current time.
+
+        This method autodetects any unspecified execution resource limits based on the
+        current cluster size, refreshing these values periodically to support cluster
+        autoscaling.
+        """
+        if (
+            time.time() - self._global_limits_last_update_time
+            < self.GLOBAL_LIMITS_UPDATE_INTERVAL_S
+        ):
+            return self._global_limits
+
+        self._global_limits_last_update_time = time.time()
+        default_limits = self._options.resource_limits
+        exclude = self._options.exclude_resources
+        total_resources = self._get_total_resources()
+        default_mem_fraction = self._object_store_memory_limit_fraction
+        total_resources = total_resources.copy(
+            object_store_memory=total_resources.object_store_memory
+            * default_mem_fraction
+        )
+        self._global_limits = default_limits.min(total_resources).subtract(exclude)
+        return self._global_limits
+
+    def get_op_usage(self, op: PhysicalOperator) -> ExecutionResources:
+        """Return the resource usage of the given operator at the current time."""
+        return self._op_usages[op]
+
+    def get_op_usage_str(self, op: PhysicalOperator, *, verbose: bool) -> str:
+        """Return a human-readable string representation of the resource usage of
+        the given operator."""
+        usage_str = f"{self._op_running_usages[op].cpu:.1f} CPU"
+        if self._op_running_usages[op].gpu:
+            usage_str += f", {self._op_running_usages[op].gpu:.1f} GPU"
+        usage_str += (
+            f", {self._op_running_usages[op].object_store_memory_str()} object store"
+        )
+
+        # NOTE: Config can override requested verbosity level
+        if LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE is not None:
+            verbose = LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE
+
+        if verbose:
+            usage_str += (
+                f" (in={memory_string(self._mem_op_internal[op])},"
+                f"out={memory_string(self._mem_op_outputs[op])})"
+            )
+            if self._op_resource_allocator is not None:
+                allocation = self._op_resource_allocator.get_allocation(op)
+                if allocation:
+                    usage_str += f", alloc=(cpu={allocation.cpu:.1f}"
+                    usage_str += f",gpu={allocation.gpu:.1f}"
+                    usage_str += f",obj_store={allocation.object_store_memory_str()})"
+
+                budget = self._op_resource_allocator.get_budget(op)
+                if budget:
+                    usage_str += f", budget=(cpu={budget.cpu:.1f}"
+                    usage_str += f",gpu={budget.gpu:.1f}"
+                    usage_str += f",obj_store={budget.object_store_memory_str()}"
+
+                    # Remaining memory budget for producing new task outputs.
+                    if isinstance(
+                        self._op_resource_allocator, ReservationOpResourceAllocator
+                    ):
+                        reserved_for_output = memory_string(
+                            self._op_resource_allocator._output_budgets.get(op, 0)
+                        )
+                        usage_str += f",out={reserved_for_output})"
+
+        return usage_str
+
+    def op_resource_allocator_enabled(self) -> bool:
+        """Return whether OpResourceAllocator is enabled."""
+        return self._op_resource_allocator is not None
+
+    @property
+    def op_resource_allocator(self) -> "OpResourceAllocator":
+        """Return the OpResourceAllocator."""
+        assert self._op_resource_allocator is not None
+        return self._op_resource_allocator
+
+    def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> int:
+        return self._op_resource_allocator.max_task_output_bytes_to_read(
+            op,
+            task_resource_usage=self._op_usages,
+            output_object_store_usage=self._mem_op_outputs,
+        )
+
+    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Return the budget for the given operator, or None if the operator
+        has unlimited budget."""
+        if self._op_resource_allocator is None:
+            return None
+        return self._op_resource_allocator.get_budget(op)
+
+    def is_op_eligible(self, op: PhysicalOperator) -> bool:
+        """Whether the op is eligible for memory reservation."""
+        return (
+            not op.throttling_disabled()
+            # As long as the op has finished execution, even if there are still
+            # non-taken outputs, we don't need to allocate resources for it.
+            and not op.execution_finished()
+        )
+
+    def get_eligible_ops(self) -> List[PhysicalOperator]:
+        return [op for op in self._topology if self.is_op_eligible(op)]
+
+    def get_downstream_ineligible_ops(
+        self, op: PhysicalOperator
+    ) -> Iterable[PhysicalOperator]:
+        """Get the downstream ineligible operators of the given operator.
+
+        E.g.,
+          - "cur_map->downstream_map" will return an empty list.
+          - "cur_map->limit1->limit2->downstream_map" will return [limit1, limit2].
+        """
+        for next_op in op.output_dependencies:
+            if not self.is_op_eligible(next_op):
+                yield next_op
+                yield from self.get_downstream_ineligible_ops(next_op)
+
+    def get_downstream_eligible_ops(
+        self, op: PhysicalOperator
+    ) -> Iterable[PhysicalOperator]:
+        """Get the downstream eligible operators of the given operator, ignoring
+        intermediate ineligible operators.
+
+        E.g.,
+          - "cur_map->downstream_map" will return [downstream_map].
+          - "cur_map->limit1->limit2->downstream_map" will return [downstream_map].
+        """
+        for next_op in op.output_dependencies:
+            if self.is_op_eligible(next_op):
+                yield next_op
+            else:
+                yield from self.get_downstream_eligible_ops(next_op)
+
+    def get_op_outputs_object_store_usage_with_downstream(
+        self, op: PhysicalOperator
+    ) -> int:
+        """Get the outputs memory usage of the given operator, including the downstream
+        ineligible operators.
+        """
+        # Outputs usage of the current operator.
+        op_outputs_usage = self._mem_op_outputs[op]
+        # Also account the downstream ineligible operators' memory usage.
+        op_outputs_usage += sum(
+            self.get_op_usage(next_op).object_store_memory
+            for next_op in self.get_downstream_ineligible_ops(op)
+        )
+        return op_outputs_usage
+
+    def get_op_internal_object_store_usage(self, op: PhysicalOperator) -> int:
+        """Get the internal object store memory usage of the given operator"""
+        return self._mem_op_internal[op]
+
+
+def _get_first_pending_shuffle_op(topology: "Topology") -> int:
+    for idx, op in enumerate(topology):
+        if _is_shuffle_op(op) and not op.completed():
+            return idx
+
+    return -1
+
+
+def _is_shuffle_op(op: PhysicalOperator) -> bool:
+    return isinstance(op, (AllToAllOperator, HashShufflingOperatorBase))
+
+
+class OpResourceAllocator(ABC):
+    """An interface for dynamic operator resource allocation.
+
+    This interface allows dynamically allocating available resources to each operator,
+    limiting how many tasks each operator can submit, and how much data each operator
+    can read from its running tasks.
+    """
+
+    class IdleDetector:
+        """Utility class for detecting idle operators.
+
+        Note, stalling can happen when there are less resources than Data executor
+        expects. E.g., when some resources are preempted by non-Data code, see
+        `test_no_deadlock_on_resource_contention` as an example.
+
+        This class is used to detect potential stalling and allow the execution
+        to make progress.
+        """
+
+        # The interval to detect idle operators.
+        # When downstream is idle, we'll allow reading at least one task output
+        # per this interval,
+        DETECTION_INTERVAL_S = 10.0
+        # Print a warning if an operator is idle for this time.
+        WARN_ON_IDLE_TIME_S = 60.0
+        # Whether a warning has been printed.
+        _warn_printed = False
+
+        def __init__(self):
+            # per-op fields
+            self.last_num_outputs = defaultdict(int)
+            self.last_output_time = defaultdict(lambda: time.time())
+            self.last_detection_time = defaultdict(lambda: time.time())
+
+        def detect_idle(self, op: PhysicalOperator):
+            cur_time = time.time()
+            if cur_time - self.last_detection_time[op] > self.DETECTION_INTERVAL_S:
+                cur_num_outputs = op.metrics.num_task_outputs_generated
+                if cur_num_outputs > self.last_num_outputs[op]:
+                    self.last_num_outputs[op] = cur_num_outputs
+                    self.last_output_time[op] = cur_time
+                    self.last_detection_time[op] = cur_time
+                else:
+                    self.last_detection_time[op] = cur_time
+                    self.print_warning_if_idle_for_too_long(
+                        op, cur_time - self.last_output_time[op]
+                    )
+                    return True
+
+            return False
+
+        @classmethod
+        def print_warning_if_idle_for_too_long(
+            cls, op: PhysicalOperator, idle_time: float
+        ):
+            """Print a warning if an operator is idle for too long."""
+            if idle_time < cls.WARN_ON_IDLE_TIME_S or cls._warn_printed:
+                return
+            cls._warn_printed = True
+            msg = (
+                f"Operator {op} is running but has no outputs for {idle_time} seconds."
+                " Execution may be slower than expected.\n"
+                "Ignore this warning if your UDF is expected to be slow."
+                " Otherwise, this can happen when there are fewer cluster resources"
+                " available to Ray Data than expected."
+                " If you have non-Data tasks or actors running in the cluster, exclude"
+                " their resources from Ray Data with"
+                " `DataContext.get_current().execution_options.exclude_resources`."
+                " This message will only print once."
+            )
+
+            logger.warning(msg)
+
+    def __init__(self, topology: "Topology"):
+        self._topology = topology
+        self._idle_detector = self.IdleDetector()
+
+    @abstractmethod
+    def update_budgets(
+        self,
+        *,
+        limits: ExecutionResources,
+    ):
+        """Callback to update resource usages."""
+        ...
+
+    @abstractmethod
+    def can_submit_new_task(self, op: PhysicalOperator) -> bool:
+        """Return whether the given operator can submit a new task."""
+        ...
+
+    @abstractmethod
+    def max_task_output_bytes_to_read(
+        self,
+        op: PhysicalOperator,
+        *,
+        task_resource_usage: Dict[PhysicalOperator, ExecutionResources],
+        output_object_store_usage: Dict[PhysicalOperator, int],
+    ) -> Optional[int]:
+        """Return the maximum bytes of pending task outputs can be read for
+        the given operator. None means no limit."""
+        ...
+
+    @abstractmethod
+    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Returns the budget for the given operator or `None` if the operator
+        has unlimited budget. Operator's budget is defined as:
+
+            Budget = Allocation - Usage
+        """
+        ...
+
+    @abstractmethod
+    def get_output_budget(self, op: PhysicalOperator) -> Optional[int]:
+        """Returns the budget for operator's outputs (in object store bytes) or
+        `None` if there's no limit.
+        """
+        ...
+
+    @abstractmethod
+    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        """Returns allocation for the given operator or `None` if operator's
+        allocation is unlimited."""
+        ...
+
+    def _get_eligible_ops(self) -> List[PhysicalOperator]:
+        first_pending_shuffle_op_idx = _get_first_pending_shuffle_op(self._topology)
+        return [
+            op
+            for idx, op in enumerate(self._topology)
+            if self._is_op_eligible(op)
+            and (
+                first_pending_shuffle_op_idx == -1
+                or idx <= first_pending_shuffle_op_idx
+            )
+        ]
+
+    @staticmethod
+    def _is_op_eligible(op: PhysicalOperator) -> bool:
+        """Whether the op is eligible for memory reservation."""
+        return (
+            not op.throttling_disabled()
+            # As long as the op has finished execution, even if there are still
+            # non-taken outputs, we don't need to allocate resources for it.
+            and not op.execution_finished()
+        )
+
+    def _get_downstream_eligible_ops(
+        self, op: PhysicalOperator
+    ) -> Iterable[PhysicalOperator]:
+        """Get the downstream eligible operators of the given operator, ignoring
+        intermediate ineligible operators.
+
+        E.g.,
+          - "cur_map->downstream_map" will return [downstream_map].
+          - "cur_map->limit1->limit2->downstream_map" will return [downstream_map].
+        """
+        for next_op in op.output_dependencies:
+            if self._is_op_eligible(next_op):
+                yield next_op
+            else:
+                yield from self._get_downstream_eligible_ops(next_op)
+
+    def _should_unblock_streaming_output_backpressure(
+        self, op: PhysicalOperator
+    ) -> bool:
+        # NOTE: If this operator is a terminal one, extracting outputs from it
+        #       should not be throttled
+        if not op.output_dependencies:
+            return True
+
+        # In some edge cases, the downstream operators may have no enough resources to
+        # launch tasks. Then we should temporarily unblock the streaming output
+        # backpressure by allowing reading at least 1 block. So the current operator
+        # can finish at least one task and yield resources to the downstream operators.
+        for downstream_op in self._get_downstream_eligible_ops(op):
+            if not self.can_submit_new_task(downstream_op):
+                # Case 1: the downstream operator hasn't reserved the minimum resources
+                # to run at least one task.
+                return True
+
+            # Case 2: the downstream operator has reserved the minimum resources, but
+            # the resources are preempted by non-Data tasks or actors.
+            # We don't have a good way to detect this case, so we'll unblock
+            # backpressure when the downstream operator has been idle for a while.
+            if self._idle_detector.detect_idle(downstream_op):
+                return True
+
+        return False
+
+
+class ReservationOpResourceAllocator(OpResourceAllocator):
+    """An OpResourceAllocator implementation that reserves resources for each operator.
+
+    This class reserves memory and CPU resources for eligible operators, and considers
+    runtime resource usages to limit the resources that each operator can use.
+
+    It works in the following way:
+    1. An operator is eligible for resource reservation, if it has enabled throttling
+       and hasn't completed. Ineligible operators are not throttled, but
+       their usage will be accounted for their upstream eligible operators. E.g., for
+       such a dataset "map1->limit->map2->streaming_split", we'll treat "map1->limit" as
+       a group and "map2->streaming_split" as another group.
+    2. For each eligible operator, we reserve `reservation_ratio * global_resources /
+        num_eligible_ops` resources, half of which is reserved only for the operator
+        outputs, excluding pending task outputs.
+    3. Non-reserved resources are shared among all operators.
+    4. In each scheduling iteration, each eligible operator will get "remaining of their
+       own reserved resources" + "remaining of shared resources / num_eligible_ops"
+       resources.
+
+    The `reservation_ratio` is set to 50% by default. Users can tune this value to
+    adjust how aggressive or conservative the resource allocation is. A higher value
+    will make the resource allocation more even, but may lead to underutilization and
+    worse performance. And vice versa.
+    """
+
+    def __init__(self, resource_manager: ResourceManager, reservation_ratio: float):
+        super().__init__(resource_manager._topology)
+        self._resource_manager = resource_manager
+        self._reservation_ratio = reservation_ratio
+        assert 0.0 <= self._reservation_ratio <= 1.0
+        # Per-op reserved resources, excluding `_reserved_for_op_outputs`.
+        self._op_reserved: Dict[PhysicalOperator, ExecutionResources] = {}
+        # Memory reserved exclusively for the outputs of each operator.
+        # "Op outputs" refer to blocks that have been taken out of an operator,
+        # i.e., `RessourceManager._mem_op_outputs`.
+        #
+        # Note, if we don't reserve memory for op outputs, all the budget may be used by
+        # the pending task outputs, and/or op's internal output buffers (the latter can
+        # happen when `preserve_order=True`).
+        # Then we'll have no budget to pull blocks from the op.
+        self._reserved_for_op_outputs: Dict[PhysicalOperator, float] = {}
+        # Total shared resources.
+        self._total_shared = ExecutionResources.zero()
+        # Resource budgets for each operator, excluding `_reserved_for_op_outputs`.
+        self._op_budgets: Dict[PhysicalOperator, ExecutionResources] = {}
+        # Remaining memory budget for generating new task outputs, per operator.
+        self._output_budgets: Dict[PhysicalOperator, float] = {}
+        # Whether each operator has reserved the minimum resources to run
+        # at least one task.
+        # This is used to avoid edge cases where the entire resource limits are not
+        # enough to run one task of each op.
+        # See `test_no_deadlock_on_small_cluster_resources` as an example.
+        self._reserved_min_resources: Dict[PhysicalOperator, bool] = {}
+
+        self._idle_detector = self.IdleDetector()
+
+    def _get_ineligible_ops_with_usage(self) -> List[PhysicalOperator]:
+        """
+        Resource reservation is based on the number of eligible operators.
+        However, there might be completed operators that still have blocks in their output queue, which we need to exclude them from the reservation.
+        And we also need to exclude the downstream ineligible operators.
+
+        E.g., for the following pipeline:
+        ```
+        map1 (completed, but still has blocks in its output queue) -> limit1 (ineligible, not completed) -> map2 (not completed) -> limit2 -> map3
+        ```
+
+        The reservation is based on the number of eligible operators (map2 and map3), but we need to exclude map1 and limit1 from the reservation.
+        """
+        last_completed_ops = []
+        ops_to_exclude_from_reservation = []
+        # Traverse operator tree collecting all operators that have already finished
+        for op in self._topology:
+            if not op.execution_finished():
+                for dep in op.input_dependencies:
+                    if dep.execution_finished():
+                        last_completed_ops.append(dep)
+
+        # In addition to completed operators,
+        # filter out downstream ineligible operators since they are omitted from reservation calculations.
+        for op in last_completed_ops:
+            ops_to_exclude_from_reservation.extend(
+                list(self._resource_manager.get_downstream_ineligible_ops(op))
+            )
+            ops_to_exclude_from_reservation.append(op)
+        return list(set(ops_to_exclude_from_reservation))
+
+    def _update_reservation(self, limits: ExecutionResources):
+        eligible_ops = self._resource_manager.get_eligible_ops()
+
+        self._op_reserved.clear()
+        self._reserved_for_op_outputs.clear()
+        self._reserved_min_resources.clear()
+
+        if len(eligible_ops) == 0:
+            return
+
+        remaining = limits.copy()
+
+        # Reserve `reservation_ratio * global_limits / num_ops` resources for each
+        # operator.
+        default_reserved = limits.scale(self._reservation_ratio / (len(eligible_ops)))
+        for index, op in enumerate(eligible_ops):
+            # Reserve at least half of the default reserved resources for the outputs.
+            # This makes sure that we will have enough budget to pull blocks from the
+            # op.
+            reserved_for_outputs = ExecutionResources(
+                0, 0, max(default_reserved.object_store_memory / 2, 1)
+            )
+
+            reserved_for_tasks = default_reserved.subtract(reserved_for_outputs)
+
+            min_resource_usage, max_resource_usage = op.min_max_resource_requirements()
+
+            if min_resource_usage is not None:
+                reserved_for_tasks = reserved_for_tasks.max(min_resource_usage)
+            if max_resource_usage is not None:
+                reserved_for_tasks = reserved_for_tasks.min(max_resource_usage)
+
+            # Check if the remaining resources are enough for both reserved_for_tasks
+            # and reserved_for_outputs. Note, we only consider CPU and GPU, but not
+            # object_store_memory, because object_store_memory can be oversubscribed,
+            # but CPU/GPU cannot.
+            if reserved_for_tasks.add(reserved_for_outputs).satisfies_limit(
+                remaining, ignore_object_store_memory=True
+            ):
+                self._reserved_min_resources[op] = True
+            else:
+                # If the remaining resources are not enough to reserve the minimum
+                # resources for this operator, we'll only reserve the minimum object
+                # store memory, but not the CPU and GPU resources.
+                # Because Ray Core doesn't allow CPU/GPU resources to be oversubscribed.
+                # NOTE: we prioritize upstream operators for minimum resource reservation.
+                # ops. It's fine that downstream ops don't get the minimum reservation,
+                # because they can wait for upstream ops to finish and release resources.
+                self._reserved_min_resources[op] = False
+                reserved_for_tasks = ExecutionResources(
+                    0, 0, min_resource_usage.object_store_memory
+                )
+                # Add `id(self)` to the log_once key so that it will be logged once
+                # per execution.
+                if index == 0 and log_once(f"low_resource_warning_{id(self)}"):
+                    # Log a warning if even the first operator cannot reserve
+                    # the minimum resources.
+                    logger.warning(
+                        f"Cluster resources are not enough to run any task from {op}."
+                        " The job may hang forever unless the cluster scales up."
+                    )
+
+            self._op_reserved[op] = reserved_for_tasks
+            self._reserved_for_op_outputs[op] = reserved_for_outputs.object_store_memory
+
+            op_total_reserved = reserved_for_tasks.add(reserved_for_outputs)
+            remaining = remaining.subtract(op_total_reserved)
+            remaining = remaining.max(ExecutionResources.zero())
+
+        self._total_shared = remaining
+
+    def can_submit_new_task(self, op: PhysicalOperator) -> bool:
+        """Return whether the given operator can submit a new task based on budget."""
+        budget = self.get_budget(op)
+        if budget is None:
+            return True
+        return op.incremental_resource_usage().satisfies_limit(budget)
+
+    def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        return self._op_budgets.get(op)
+
+    def get_output_budget(self, op: PhysicalOperator) -> Optional[int]:
+        return self._output_budgets.get(op)
+
+    def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]:
+        # TODO fix
+        return ExecutionResources.zero()
+
+    def max_task_output_bytes_to_read(
+        self,
+        op: PhysicalOperator,
+        *,
+        task_resource_usage: Dict[PhysicalOperator, ExecutionResources],
+        output_object_store_usage: Dict[PhysicalOperator, int],
+    ) -> Optional[int]:
+        if op not in self._op_budgets:
+            return None
+        res = self._op_budgets[op].object_store_memory
+        # Add the remaining of `_reserved_for_op_outputs`.
+        op_outputs_usage = (
+            self._resource_manager.get_op_outputs_object_store_usage_with_downstream(op)
+        )
+
+        res += max(self._reserved_for_op_outputs[op] - op_outputs_usage, 0)
+        if math.isinf(res):
+            self._output_budgets[op] = res
+            return None
+
+        res = int(res)
+        assert res >= 0
+        if res == 0 and self._should_unblock_streaming_output_backpressure(op):
+            res = 1
+        self._output_budgets[op] = res
+        return res
+
+    def update_budgets(
+        self,
+        *,
+        limits: ExecutionResources,
+    ):
+        op_to_exclude_from_reservation = self._get_ineligible_ops_with_usage()
+        for completed_op in op_to_exclude_from_reservation:
+            completed_op_usage = self._resource_manager.get_op_usage(completed_op)
+
+            limits = limits.subtract(completed_op_usage)
+            limits = limits.max(ExecutionResources.zero())
+
+        # Remaining resources to be distributed across operators
+        remaining_shared = self._update_reservation(limits)
+
+        self._op_budgets.clear()
+        eligible_ops = self._resource_manager.get_eligible_ops()
+        if len(eligible_ops) == 0:
+            return
+
+        # Remaining of shared resources.
+        remaining_shared = self._total_shared
+        for op in eligible_ops:
+            # Calculate the memory usage of the operator.
+            op_mem_usage = 0
+            # Add the memory usage of the operator itself,
+            # excluding `_reserved_for_op_outputs`.
+            op_mem_usage += self._resource_manager.get_op_internal_object_store_usage(
+                op
+            )
+            # Add the portion of op outputs usage that has
+            # exceeded `_reserved_for_op_outputs`.
+            op_outputs_usage = self._resource_manager.get_op_outputs_object_store_usage_with_downstream(
+                op
+            )
+            op_mem_usage += max(op_outputs_usage - self._reserved_for_op_outputs[op], 0)
+
+            op_usage = self._resource_manager.get_op_usage(op).copy(
+                object_store_memory=op_mem_usage
+            )
+
+            op_reserved = self._op_reserved[op]
+            # How much of the reserved resources are remaining.
+            op_reserved_remaining = op_reserved.subtract(op_usage).max(
+                ExecutionResources.zero()
+            )
+
+            self._op_budgets[op] = op_reserved_remaining
+            # How much of the reserved resources are exceeded.
+            # If exceeded, we need to subtract from the remaining shared resources.
+            op_reserved_exceeded = op_usage.subtract(op_reserved).max(
+                ExecutionResources.zero()
+            )
+            remaining_shared = remaining_shared.subtract(op_reserved_exceeded)
+
+        remaining_shared = remaining_shared.max(ExecutionResources.zero())
+
+        # Allocate the remaining shared resources to each operator.
+        for i, op in enumerate(reversed(eligible_ops)):
+            # By default, divide the remaining shared resources equally.
+            op_shared = remaining_shared.scale(1.0 / (len(eligible_ops) - i))
+            # But if the op's budget is less than `incremental_resource_usage`,
+            # it will be useless. So we'll let the downstream operator
+            # borrow some resources from the upstream operator, if remaining_shared
+            # is still enough.
+            to_borrow = (
+                op.incremental_resource_usage()
+                .subtract(self._op_budgets[op].add(op_shared))
+                .max(ExecutionResources.zero())
+            )
+            if not to_borrow.is_zero() and op_shared.add(to_borrow).satisfies_limit(
+                remaining_shared
+            ):
+                op_shared = op_shared.add(to_borrow)
+            remaining_shared = remaining_shared.subtract(op_shared)
+            assert remaining_shared.is_non_negative(), (
+                remaining_shared,
+                op,
+                op_shared,
+                to_borrow,
+            )
+
+            if op.min_max_resource_requirements()[1].gpu > 0:
+                # If an operator needs GPU, we just allocate all GPUs to it.
+                # TODO(hchen): allocate resources across multiple GPU operators.
+
+                # The op_usage can be more than the global limit in the following cases:
+                # 1. The op is setting a minimum concurrency that is larger than
+                #    available num of GPUs.
+                # 2. The cluster scales down, and the global limit decreases.
+                target_num_gpu = max(
+                    limits.gpu - self._resource_manager.get_op_usage(op).gpu,
+                    0,
+                )
+            else:
+                target_num_gpu = 0
+
+            self._op_budgets[op] = (
+                self._op_budgets[op].add(op_shared).copy(gpu=target_num_gpu)
+            )
+
+        # A materializing operator like `AllToAllOperator` waits for all its input
+        # operator's outputs before processing data. This often forces the input
+        # operator to exceed its object store memory budget. To prevent deadlock, we
+        # disable object store memory backpressure for the input operator.
+        for op in eligible_ops:
+            if any(
+                isinstance(next_op, MATERIALIZING_OPERATORS)
+                for next_op in op.output_dependencies
+            ):
+                self._op_budgets[op] = self._op_budgets[op].copy(
+                    object_store_memory=float("inf")
+                )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/streaming_executor.py b/lib/python3.12/site-packages/ray/data/_internal/execution/streaming_executor.py
new file mode 100644
index 0000000000000000000000000000000000000000..38a951479cf001891834e5f4f492d1a5100961a1
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/streaming_executor.py
@@ -0,0 +1,816 @@
+import logging
+import math
+import threading
+import time
+from typing import Dict, List, Optional, Tuple
+
+from ray.data._internal.actor_autoscaler import (
+    create_actor_autoscaler,
+)
+from ray.data._internal.cluster_autoscaler import create_cluster_autoscaler
+from ray.data._internal.execution import create_ranker
+from ray.data._internal.execution.backpressure_policy import (
+    BackpressurePolicy,
+    get_backpressure_policies,
+)
+from ray.data._internal.execution.dataset_state import DatasetState
+from ray.data._internal.execution.execution_callback import get_execution_callbacks
+from ray.data._internal.execution.interfaces import (
+    ExecutionResources,
+    Executor,
+    OutputIterator,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    InternalQueueOperatorMixin,
+)
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.progress_manager import RichExecutionProgressManager
+from ray.data._internal.execution.resource_manager import (
+    ResourceManager,
+)
+from ray.data._internal.execution.streaming_executor_state import (
+    OpState,
+    Topology,
+    build_streaming_topology,
+    process_completed_tasks,
+    select_operator_to_run,
+    update_operator_states,
+)
+from ray.data._internal.logging import (
+    get_log_directory,
+    register_dataset_logger,
+    unregister_dataset_logger,
+)
+from ray.data._internal.metadata_exporter import Topology as TopologyMetadata
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.stats import DatasetStats, Timer, _StatsManager
+from ray.data.context import OK_PREFIX, WARN_PREFIX, DataContext
+from ray.util.debug import log_once
+from ray.util.metrics import Gauge
+
+logger = logging.getLogger(__name__)
+
+# Force a progress update after this many events processed. Avoids the
+# progress seeming to stall for very large scale workloads.
+PROGRESS_BAR_UPDATE_INTERVAL = 50
+PROGRESS_MANAGER_UPDATE_INTERVAL = 20
+
+# Interval for logging execution progress updates and operator metrics.
+DEBUG_LOG_INTERVAL_SECONDS = 5
+
+# Visible for testing.
+_num_shutdown = 0
+
+
+class StreamingExecutor(Executor, threading.Thread):
+    """A streaming Dataset executor.
+
+    This implementation executes Dataset DAGs in a fully streamed way. It runs
+    by setting up the operator topology, and then routing blocks through operators in
+    a way that maximizes throughput under resource constraints.
+    """
+
+    UPDATE_METRICS_INTERVAL_S: float = 5.0
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        dataset_id: str = "unknown_dataset",
+    ):
+        self._data_context = data_context
+        self._ranker = create_ranker()
+        self._start_time: Optional[float] = None
+        self._initial_stats: Optional[DatasetStats] = None
+        self._final_stats: Optional[DatasetStats] = None
+        self._global_info: Optional[ProgressBar] = None
+        self._progress_manager: Optional[RichExecutionProgressManager] = None
+
+        # The executor can be shutdown while still running.
+        self._shutdown_lock = threading.RLock()
+        self._execution_started = False
+        self._shutdown = False
+
+        # Internal execution state shared across thread boundaries. We run the control
+        # loop on a separate thread so that it doesn't become stalled between
+        # generator `yield`s.
+        self._topology: Optional[Topology] = None
+        self._output_node: Optional[Tuple[PhysicalOperator, OpState]] = None
+        self._backpressure_policies: List[BackpressurePolicy] = []
+
+        self._dataset_id = dataset_id
+        # Stores if an operator is completed,
+        # used for marking when an op has just completed.
+        self._has_op_completed: Optional[Dict[PhysicalOperator, bool]] = None
+        self._max_errored_blocks = self._data_context.max_errored_blocks
+        self._num_errored_blocks = 0
+
+        self._last_debug_log_time = 0
+
+        self._data_context.set_dataset_logger_id(
+            register_dataset_logger(self._dataset_id)
+        )
+
+        # This stores the last time we updated the metrics.
+        # This allows us to update metrics on some interval,
+        # by comparing it with the current timestamp.
+        self._metrics_last_updated: float = 0.0
+
+        self._sched_loop_duration_s = Gauge(
+            "data_sched_loop_duration_s",
+            description="Duration of the scheduling loop in seconds",
+            tag_keys=("dataset",),
+        )
+
+        self._cpu_budget_gauge: Gauge = Gauge(
+            "data_cpu_budget",
+            "Budget (CPU) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._gpu_budget_gauge: Gauge = Gauge(
+            "data_gpu_budget",
+            "Budget (GPU) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._memory_budget_gauge: Gauge = Gauge(
+            "data_memory_budget",
+            "Budget (Memory) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._osm_budget_gauge: Gauge = Gauge(
+            "data_object_store_memory_budget",
+            "Budget (Object Store Memory) per operator",
+            tag_keys=("dataset", "operator"),
+        )
+        self._max_bytes_to_read_gauge: Gauge = Gauge(
+            "data_max_bytes_to_read",
+            description="Maximum bytes to read from streaming generator buffer.",
+            tag_keys=("dataset", "operator"),
+        )
+
+        Executor.__init__(self, self._data_context.execution_options)
+        thread_name = f"StreamingExecutor-{self._dataset_id}"
+        threading.Thread.__init__(self, daemon=True, name=thread_name)
+
+    def execute(
+        self, dag: PhysicalOperator, initial_stats: Optional[DatasetStats] = None
+    ) -> OutputIterator:
+        """Executes the DAG using a streaming execution strategy.
+
+        We take an event-loop approach to scheduling. We block on the next scheduling
+        event using `ray.wait`, updating operator state and dispatching new tasks.
+        """
+
+        self._initial_stats = initial_stats
+        self._start_time = time.perf_counter()
+
+        if not isinstance(dag, InputDataBuffer):
+            if self._data_context.print_on_execution_start:
+                message = f"Starting execution of Dataset {self._dataset_id}."
+                log_path = get_log_directory()
+                if log_path is not None:
+                    message += f" Full logs are in {log_path}"
+                logger.info(message)
+                logger.info(
+                    f"Execution plan of Dataset {self._dataset_id}: {dag.dag_str}"
+                )
+
+            logger.debug("Execution config: %s", self._options)
+
+        # Setup the streaming DAG topology and start the runner thread.
+        self._topology = build_streaming_topology(dag, self._options)
+
+        # Setup progress bars
+        if self._use_rich_progress():
+            self._progress_manager = RichExecutionProgressManager(
+                self._dataset_id, self._topology
+            )
+            self._progress_manager.start()
+        else:
+            self._global_info = ProgressBar(
+                "Running", dag.num_output_rows_total(), unit="row"
+            )
+            num_progress_bars = 1
+            for op_state in list(self._topology.values()):
+                if not isinstance(op_state.op, InputDataBuffer):
+                    num_progress_bars += op_state.initialize_progress_bars(
+                        num_progress_bars, self._options.verbose_progress
+                    )
+
+        self._resource_manager = ResourceManager(
+            self._topology,
+            self._options,
+            lambda: self._cluster_autoscaler.get_total_resources(),
+            self._data_context,
+        )
+        self._backpressure_policies = get_backpressure_policies(
+            self._data_context, self._topology, self._resource_manager
+        )
+        self._cluster_autoscaler = create_cluster_autoscaler(
+            self._topology,
+            self._resource_manager,
+            execution_id=self._dataset_id,
+        )
+        self._actor_autoscaler = create_actor_autoscaler(
+            self._topology,
+            self._resource_manager,
+            config=self._data_context.autoscaling_config,
+        )
+
+        self._has_op_completed = dict.fromkeys(self._topology, False)
+
+        self._output_node = dag, self._topology[dag]
+
+        op_to_id = {
+            op: self._get_operator_id(op, i) for i, op in enumerate(self._topology)
+        }
+        _StatsManager.register_dataset_to_stats_actor(
+            self._dataset_id,
+            self._get_operator_tags(),
+            TopologyMetadata.create_topology_metadata(dag, op_to_id),
+            self._data_context,
+        )
+        for callback in get_execution_callbacks(self._data_context):
+            callback.before_execution_starts(self)
+
+        self.start()
+        self._execution_started = True
+
+        return _ClosingIterator(self)
+
+    def __del__(self):
+        # NOTE: Upon garbage-collection we're allowing running tasks
+        #       to be terminated asynchronously (ie avoid unnecessary
+        #       synchronization on their completion)
+        self.shutdown(force=False)
+
+    def shutdown(self, force: bool, exception: Optional[Exception] = None):
+        global _num_shutdown
+
+        with self._shutdown_lock:
+            if not self._execution_started or self._shutdown:
+                return
+
+            start = time.perf_counter()
+
+            status_detail = (
+                f"failed with {exception}" if exception else "completed successfully"
+            )
+
+            logger.debug(
+                f"Shutting down executor for dataset {self._dataset_id} "
+                f"({status_detail})"
+            )
+
+            _num_shutdown += 1
+            self._shutdown = True
+            # Give the scheduling loop some time to finish processing.
+            self.join(timeout=2.0)
+            self._update_stats_metrics(
+                state=DatasetState.FINISHED.name
+                if exception is None
+                else DatasetState.FAILED.name,
+                force_update=True,
+            )
+            # Freeze the stats and save it.
+            self._final_stats = self._generate_stats()
+            stats_summary_string = self._final_stats.to_summary().to_string(
+                include_parent=False
+            )
+            # Reset the scheduling loop duration gauge + resource manager budgets/usages.
+            self._resource_manager.update_usages()
+            self.update_metrics(0)
+            if self._data_context.enable_auto_log_stats:
+                logger.info(stats_summary_string)
+            # Close the progress manager with a finishing message.
+            if exception is None:
+                desc = (
+                    f"{OK_PREFIX} Dataset {self._dataset_id} execution finished in "
+                    f"{self._final_stats.time_total_s:.2f} seconds"
+                )
+            else:
+                desc = f"{WARN_PREFIX} Dataset {self._dataset_id} execution failed"
+
+            if self._use_rich_progress() and self._progress_manager:
+                self._progress_manager.close_with_finishing_description(
+                    desc, exception is None
+                )
+            elif not self._use_rich_progress() and self._global_info:
+                logger.info(desc)
+                self._global_info.set_description(desc)
+                self._global_info.close()
+
+            timer = Timer()
+
+            for op, state in self._topology.items():
+                op.shutdown(timer, force=force)
+                if not self._use_rich_progress():
+                    # we only close sub-progress bars for the tqdm
+                    # implementation because closing is centrally
+                    # managed in the rich implementation.
+                    state.close_progress_bars()
+
+            min_ = round(timer.min(), 3)
+            max_ = round(timer.max(), 3)
+            total = round(timer.get(), 3)
+            logger.debug(
+                f"Shut down operator hierarchy for dataset {self._dataset_id}"
+                f" (min/max/total={min_}/{max_}/{total}s)"
+            )
+
+            if exception is None:
+                for callback in get_execution_callbacks(self._data_context):
+                    callback.after_execution_succeeds(self)
+            else:
+                for callback in get_execution_callbacks(self._data_context):
+                    callback.after_execution_fails(self, exception)
+
+            self._cluster_autoscaler.on_executor_shutdown()
+
+            dur = time.perf_counter() - start
+
+            logger.debug(
+                f"Shut down executor for dataset {self._dataset_id} "
+                f"(took {round(dur, 3)}s)"
+            )
+
+            # Unregister should be called after all operators are shut down to
+            # capture as many logs as possible.
+            self._data_context.set_dataset_logger_id(
+                unregister_dataset_logger(self._dataset_id)
+            )
+
+    def run(self):
+        """Run the control loop in a helper thread.
+
+        Results are returned via the output node's outqueue.
+        """
+        exc: Optional[Exception] = None
+        try:
+            # Run scheduling loop until complete.
+            while True:
+                # Use `perf_counter` rather than `process_time` to ensure we include
+                # time spent on IO, like RPCs to Ray Core.
+                t_start = time.perf_counter()
+                continue_sched = self._scheduling_loop_step(self._topology)
+
+                sched_loop_duration = time.perf_counter() - t_start
+
+                self.update_metrics(sched_loop_duration)
+                if self._initial_stats:
+                    self._initial_stats.streaming_exec_schedule_s.add(
+                        sched_loop_duration
+                    )
+
+                for callback in get_execution_callbacks(self._data_context):
+                    callback.on_execution_step(self)
+                if not continue_sched or self._shutdown:
+                    break
+        except Exception as e:
+            # Propagate it to the result iterator.
+            exc = e
+        finally:
+            # Mark state of outputting operator as finished
+            _, state = self._output_node
+            state.mark_finished(exc)
+
+    def update_metrics(self, sched_loop_duration: int):
+        self._sched_loop_duration_s.set(
+            sched_loop_duration, tags={"dataset": self._dataset_id}
+        )
+        for i, op in enumerate(self._topology):
+            tags = {
+                "dataset": self._dataset_id,
+                "operator": self._get_operator_id(op, i),
+            }
+            self._update_budget_metrics(op, tags)
+            self._update_max_bytes_to_read_metric(op, tags)
+
+    def _update_budget_metrics(self, op: PhysicalOperator, tags: Dict[str, str]):
+        budget = self._resource_manager.get_budget(op)
+        if budget is None:
+            cpu_budget = 0
+            gpu_budget = 0
+            memory_budget = 0
+            object_store_memory_budget = 0
+        else:
+            # Convert inf to -1 to represent unlimited budget in metrics
+            cpu_budget = -1 if math.isinf(budget.cpu) else budget.cpu
+            gpu_budget = -1 if math.isinf(budget.gpu) else budget.gpu
+            memory_budget = -1 if math.isinf(budget.memory) else budget.memory
+            object_store_memory_budget = (
+                -1
+                if math.isinf(budget.object_store_memory)
+                else budget.object_store_memory
+            )
+
+        self._cpu_budget_gauge.set(cpu_budget, tags=tags)
+        self._gpu_budget_gauge.set(gpu_budget, tags=tags)
+        self._memory_budget_gauge.set(memory_budget, tags=tags)
+        self._osm_budget_gauge.set(object_store_memory_budget, tags=tags)
+
+    def _update_max_bytes_to_read_metric(
+        self, op: PhysicalOperator, tags: Dict[str, str]
+    ):
+        if self._resource_manager.op_resource_allocator_enabled():
+            resource_allocator = self._resource_manager.op_resource_allocator
+            output_budget_bytes = resource_allocator.get_output_budget(op)
+            if output_budget_bytes is not None:
+                if math.isinf(output_budget_bytes):
+                    # Convert inf to -1 to represent unlimited bytes to read
+                    output_budget_bytes = -1
+                self._max_bytes_to_read_gauge.set(output_budget_bytes, tags)
+
+    def get_stats(self):
+        """Return the stats object for the streaming execution.
+
+        The stats object will be updated as streaming execution progresses.
+        """
+        if self._final_stats:
+            return self._final_stats
+        else:
+            return self._generate_stats()
+
+    def _generate_stats(self) -> DatasetStats:
+        """Create a new stats object reflecting execution status so far."""
+        stats = self._initial_stats or DatasetStats(metadata={}, parent=None)
+        for op in self._topology:
+            if isinstance(op, InputDataBuffer):
+                continue
+            builder = stats.child_builder(op.name, override_start_time=self._start_time)
+            stats = builder.build_multioperator(op.get_stats())
+            stats.extra_metrics = op.metrics.as_dict(skip_internal_metrics=True)
+        stats.streaming_exec_schedule_s = (
+            self._initial_stats.streaming_exec_schedule_s
+            if self._initial_stats
+            else None
+        )
+        return stats
+
+    def _scheduling_loop_step(self, topology: Topology) -> bool:
+        """Run one step of the scheduling loop.
+
+        This runs a few general phases:
+            1. Waiting for the next task completion using `ray.wait()`.
+            2. Pulling completed refs into operator outqueues.
+            3. Selecting and dispatching new inputs to operators.
+
+        Returns:
+            True if we should continue running the scheduling loop.
+        """
+        self._resource_manager.update_usages()
+        # Note: calling process_completed_tasks() is expensive since it incurs
+        # ray.wait() overhead, so make sure to allow multiple dispatch per call for
+        # greater parallelism.
+        num_errored_blocks = process_completed_tasks(
+            topology,
+            self._backpressure_policies,
+            self._max_errored_blocks,
+        )
+        if self._max_errored_blocks > 0:
+            self._max_errored_blocks -= num_errored_blocks
+        self._num_errored_blocks += num_errored_blocks
+
+        self._resource_manager.update_usages()
+        # Dispatch as many operators as we can for completed tasks.
+        self._report_current_usage()
+
+        i = 0
+        while True:
+            op = select_operator_to_run(
+                topology,
+                self._resource_manager,
+                self._backpressure_policies,
+                # If consumer is idling (there's nothing for it to consume)
+                # enforce liveness, ie that at least a single task gets scheduled
+                ensure_liveness=self._consumer_idling(),
+                ranker=self._ranker,
+            )
+
+            if op is None:
+                break
+
+            topology[op].dispatch_next_task()
+
+            self._resource_manager.update_usages()
+
+            i += 1
+            if not self._use_rich_progress() and i % PROGRESS_BAR_UPDATE_INTERVAL == 0:
+                self._refresh_progress_bars(topology)
+            if self._use_rich_progress() and i % PROGRESS_MANAGER_UPDATE_INTERVAL == 0:
+                self._refresh_progress_manager(topology)
+
+        # Trigger autoscaling
+        self._cluster_autoscaler.try_trigger_scaling()
+        self._actor_autoscaler.try_trigger_scaling()
+
+        update_operator_states(topology)
+        if self._use_rich_progress():
+            self._refresh_progress_manager(topology)
+        else:
+            self._refresh_progress_bars(topology)
+
+        self._update_stats_metrics(state=DatasetState.RUNNING.name)
+        if time.time() - self._last_debug_log_time >= DEBUG_LOG_INTERVAL_SECONDS:
+            _log_op_metrics(topology)
+            _debug_dump_topology(topology, self._resource_manager)
+            self._last_debug_log_time = time.time()
+
+        # Log metrics of newly completed operators.
+        for op, state in topology.items():
+            if op.completed() and not self._has_op_completed[op]:
+                log_str = (
+                    f"Operator {op} completed. "
+                    f"Operator Metrics:\n{op._metrics.as_dict(skip_internal_metrics=True)}"
+                )
+                logger.debug(log_str)
+                self._has_op_completed[op] = True
+                self._validate_operator_queues_empty(op, state)
+
+        # Keep going until all operators run to completion.
+        return not all(op.completed() for op in topology)
+
+    def _refresh_progress_bars(self, topology: Topology):
+        # Update the progress bar to reflect scheduling decisions.
+        assert not self._use_rich_progress()
+        for op_state in topology.values():
+            op_state.refresh_progress_bar(self._resource_manager)
+        # Refresh the global progress bar to update elapsed time progress.
+        if self._global_info:
+            self._global_info.refresh()
+
+    def _refresh_progress_manager(self, topology: Topology):
+        # Update the progress manager to reflect scheduling decisions.
+        assert self._use_rich_progress()
+        if self._progress_manager:
+            for op_state in topology.values():
+                if not isinstance(op_state.op, InputDataBuffer):
+                    op_state.update_display_metrics(self._resource_manager)
+                    self._progress_manager.update_operator_progress(op_state)
+            self._progress_manager.refresh()
+
+    def _consumer_idling(self) -> bool:
+        """Returns whether the user thread is blocked on topology execution."""
+        _, state = self._output_node
+        return len(state.output_queue) == 0
+
+    def _validate_operator_queues_empty(
+        self, op: PhysicalOperator, state: OpState
+    ) -> None:
+        """Validate that all queues are empty when an operator completes.
+
+        Args:
+            op: The completed operator to validate.
+            state: The operator's execution state.
+        """
+        error_msg = "Expected {} Queue for {} to be empty, but found {} bundles"
+
+        if isinstance(op, InternalQueueOperatorMixin):
+            # 1) Check Internal Input Queue is empty
+            assert op.internal_input_queue_num_blocks() == 0, error_msg.format(
+                "Internal Input", op.name, op.internal_input_queue_num_blocks()
+            )
+
+            # 2) Check Internal Output Queue is empty
+            assert op.internal_output_queue_num_blocks() == 0, error_msg.format(
+                "Internal Output",
+                op.name,
+                op.internal_output_queue_num_blocks(),
+            )
+
+        # 3) Check that External Input Queue is empty
+        for input_q in state.input_queues:
+            assert len(input_q) == 0, error_msg.format(
+                "External Input", op.name, len(input_q)
+            )
+
+    def _report_current_usage(self) -> None:
+        # running_usage is the amount of resources that have been requested but
+        # not necessarily available
+        # TODO(sofian) https://github.com/ray-project/ray/issues/47520
+        # We need to split the reported resources into running, pending-scheduling,
+        # pending-node-assignment.
+        running_usage = self._resource_manager.get_global_running_usage()
+        pending_usage = self._resource_manager.get_global_pending_usage()
+        limits = self._resource_manager.get_global_limits()
+        resources_status = (
+            f"Active & requested resources: "
+            f"{running_usage.cpu:.4g}/{limits.cpu:.4g} CPU, "
+        )
+        if running_usage.gpu > 0:
+            resources_status += f"{running_usage.gpu:.4g}/{limits.gpu:.4g} GPU, "
+        resources_status += (
+            f"{running_usage.object_store_memory_str()}/"
+            f"{limits.object_store_memory_str()} object store"
+        )
+
+        # Only include pending section when there are pending resources.
+        if pending_usage.cpu or pending_usage.gpu:
+            if pending_usage.cpu and pending_usage.gpu:
+                pending_str = (
+                    f"{pending_usage.cpu:.4g} CPU, {pending_usage.gpu:.4g} GPU"
+                )
+            elif pending_usage.cpu:
+                pending_str = f"{pending_usage.cpu:.4g} CPU"
+            else:
+                pending_str = f"{pending_usage.gpu:.4g} GPU"
+            resources_status += f" (pending: {pending_str})"
+
+        if self._use_rich_progress() and self._progress_manager:
+            self._progress_manager.update_resource_status(resources_status)
+        elif not self._use_rich_progress() and self._global_info:
+            resources_status = (
+                f"Running Dataset: {self._dataset_id}. {resources_status}"
+            )
+            self._global_info.set_description(resources_status)
+
+    def _get_operator_id(self, op: PhysicalOperator, topology_index: int) -> str:
+        return f"{op.name}_{topology_index}"
+
+    def _get_operator_tags(self):
+        """Returns a list of operator tags."""
+        return [
+            f"{self._get_operator_id(op, i)}" for i, op in enumerate(self._topology)
+        ]
+
+    def _get_state_dict(self, state):
+        last_op, last_state = list(self._topology.items())[-1]
+        return {
+            "state": state,
+            "progress": last_state.num_completed_tasks,
+            "total": last_op.num_outputs_total(),
+            "total_rows": last_op.num_output_rows_total(),
+            "end_time": time.time()
+            if state in (DatasetState.FINISHED.name, DatasetState.FAILED.name)
+            else None,
+            "operators": {
+                f"{self._get_operator_id(op, i)}": {
+                    "name": op.name,
+                    "progress": op_state.num_completed_tasks,
+                    "total": op.num_outputs_total(),
+                    "total_rows": op.num_output_rows_total(),
+                    "queued_blocks": op_state.total_enqueued_input_blocks(),
+                    "state": DatasetState.FINISHED.name
+                    if op.execution_finished()
+                    else state,
+                }
+                for i, (op, op_state) in enumerate(self._topology.items())
+            },
+        }
+
+    def _update_stats_metrics(self, state: str, force_update: bool = False):
+        now = time.time()
+        if (
+            force_update
+            or (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S
+        ):
+            _StatsManager.update_execution_metrics(
+                self._dataset_id,
+                [op.metrics for op in self._topology],
+                self._get_operator_tags(),
+                self._get_state_dict(state=state),
+            )
+            self._metrics_last_updated = now
+
+    def _use_rich_progress(self):
+        rich_enabled = self._data_context.enable_rich_progress_bars
+        use_ray_tqdm = self._data_context.use_ray_tqdm
+
+        if not rich_enabled or use_ray_tqdm:
+            if log_once("ray_data_rich_progress_disabled"):
+                logger.info(
+                    "[dataset]: A new progress UI is available. To enable, "
+                    "set `ray.data.DataContext.get_current()."
+                    "enable_rich_progress_bars = True` and `ray.data."
+                    "DataContext.get_current().use_ray_tqdm = False`."
+                )
+            return False
+        return True
+
+
+def _validate_dag(dag: PhysicalOperator, limits: ExecutionResources) -> None:
+    """Raises an exception on invalid DAGs.
+
+    It checks if the sum of min actor pool sizes are larger than the resource
+    limit, as well as other unsupported resource configurations.
+
+    This should be called prior to creating the topology from the DAG.
+
+    Args:
+        dag: The DAG to validate.
+        limits: The limits to validate against.
+    """
+
+    seen = set()
+
+    def walk(op):
+        seen.add(op)
+        for parent in op.input_dependencies:
+            if parent not in seen:
+                yield from walk(parent)
+        yield op
+
+    base_usage = ExecutionResources(cpu=1)
+    for op in walk(dag):
+        min_resource_usage, _ = op.min_max_resource_requirements()
+        base_usage = base_usage.add(min_resource_usage)
+
+    if not base_usage.satisfies_limit(limits):
+        error_message = (
+            "The current cluster doesn't have the required resources to execute your "
+            "Dataset pipeline:\n"
+        )
+        if base_usage.cpu > limits.cpu:
+            error_message += (
+                f"- Your application needs {base_usage.cpu} CPU(s), but your cluster "
+                f"only has {limits.cpu}.\n"
+            )
+        if base_usage.gpu > limits.gpu:
+            error_message += (
+                f"- Your application needs {base_usage.gpu} GPU(s), but your cluster "
+                f"only has {limits.gpu}.\n"
+            )
+        if base_usage.object_store_memory > limits.object_store_memory:
+            error_message += (
+                f"- Your application needs {base_usage.object_store_memory}B object "
+                f"store memory, but your cluster only has "
+                f"{limits.object_store_memory}B.\n"
+            )
+        raise ValueError(error_message.strip())
+
+
+def _debug_dump_topology(topology: Topology, resource_manager: ResourceManager) -> None:
+    """Log current execution state for the topology for debugging.
+
+    Args:
+        topology: The topology to debug.
+        resource_manager: The resource manager for this topology.
+    """
+    logger.debug("Execution Progress:")
+    for i, (op, state) in enumerate(topology.items()):
+        state.update_display_metrics(resource_manager)
+        logger.debug(
+            f"{i}: {state.summary_str(resource_manager, verbose=True)}, "
+            f"Blocks Outputted: {state.num_completed_tasks}/{op.num_outputs_total()}, "
+            f"Metrics: {state.op_display_metrics.display_str()}"
+        )
+
+
+def _log_op_metrics(topology: Topology) -> None:
+    """Logs the metrics of each operator.
+
+    Args:
+        topology: The topology to debug.
+    """
+    log_str = "Operator Metrics:\n"
+    for op in topology:
+        log_str += f"{op.name}: {op.metrics.as_dict(skip_internal_metrics=True)}\n"
+    logger.debug(log_str)
+
+
+class _ClosingIterator(OutputIterator):
+    """Iterator automatically shutting down executor upon exhausting the
+    iterable sequence.
+
+    NOTE: If this iterator isn't fully exhausted, executor still have to
+          be closed manually by the caller!
+    """
+
+    def __init__(self, executor: StreamingExecutor):
+        self._executor = executor
+
+    def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
+        try:
+            op, state = self._executor._output_node
+            bundle = state.get_output_blocking(output_split_idx)
+
+            # Update progress-bars
+            using_rich = self._executor._use_rich_progress()
+            if not using_rich and self._executor._global_info:
+                self._executor._global_info.update(
+                    bundle.num_rows(), op.num_output_rows_total()
+                )
+            elif using_rich and self._executor._progress_manager:
+                self._executor._progress_manager.update_total_progress(
+                    bundle.num_rows() or 0, op.num_output_rows_total()
+                )
+
+            return bundle
+
+        # Have to be BaseException to catch ``KeyboardInterrupt``
+        #
+        # NOTE: This also handles ``StopIteration``
+        except BaseException as e:
+            # Asynchronously shutdown the executor (ie avoid unnecessary
+            # synchronization on tasks termination)
+            self._executor.shutdown(
+                force=False, exception=e if not isinstance(e, StopIteration) else None
+            )
+            raise
+
+    def __del__(self):
+        # NOTE: Upon garbage-collection we're allowing running tasks
+        #       to be terminated asynchronously (ie avoid unnecessary
+        #       synchronization on their completion)
+        self._executor.shutdown(force=False)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/streaming_executor_state.py b/lib/python3.12/site-packages/ray/data/_internal/execution/streaming_executor_state.py
new file mode 100644
index 0000000000000000000000000000000000000000..e8db6b64f4f220cf155c5cfdad902b9db70b2402
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/streaming_executor_state.py
@@ -0,0 +1,846 @@
+"""Contains classes that encapsulate streaming executor state.
+
+This is split out from streaming_executor.py to facilitate better unit testing.
+"""
+
+import logging
+import threading
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+from uuid import UUID
+
+import ray
+from ray.data._internal.execution.backpressure_policy import BackpressurePolicy
+from ray.data._internal.execution.bundle_queue import create_bundle_queue
+from ray.data._internal.execution.interfaces import (
+    ExecutionOptions,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    DataOpTask,
+    MetadataOpTask,
+    OpTask,
+    Waitable,
+    _ActorPoolInfo,
+)
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+    InternalQueueOperatorMixin,
+)
+from ray.data._internal.execution.operators.hash_shuffle import (
+    HashShuffleProgressBarMixin,
+)
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.ranker import Ranker
+from ray.data._internal.execution.resource_manager import (
+    ResourceManager,
+)
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.util import (
+    unify_schemas_with_validation,
+)
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data.block import Schema
+
+logger = logging.getLogger(__name__)
+
+# Holds the full execution state of the streaming topology. It's a dict mapping each
+# operator to tracked streaming exec state.
+Topology = Dict[PhysicalOperator, "OpState"]
+
+
+class OpBufferQueue:
+    """A FIFO queue to buffer RefBundles between upstream and downstream operators.
+    This class is thread-safe.
+    """
+
+    def __init__(self):
+        self._num_blocks = 0
+        self._queue = create_bundle_queue()
+        self._num_per_split = defaultdict(int)
+        self._lock = threading.Lock()
+        # Used to buffer output RefBundles indexed by output splits.
+        self._outputs_by_split = defaultdict(create_bundle_queue)
+        super().__init__()
+
+    @property
+    def memory_usage(self) -> int:
+        """The total memory usage of the queue in bytes."""
+        with self._lock:
+            # The split queues contain bundles popped from the main queue. So, a bundle
+            # will either be in the main queue or in one of the split queues, and we
+            # don't need to worry about double counting.
+            return self._queue.estimate_size_bytes() + sum(
+                split_queue.estimate_size_bytes()
+                for split_queue in self._outputs_by_split.values()
+            )
+
+    @property
+    def num_blocks(self) -> int:
+        """The total number of blocks in the queue."""
+        with self._lock:
+            return self._num_blocks
+
+    def __len__(self):
+        with self._lock:
+            return len(self._queue)
+
+    def has_next(self, output_split_idx: Optional[int] = None) -> bool:
+        """Whether next RefBundle is available.
+
+        Args:
+            output_split_idx: If specified, only check ref bundles with the
+                given output split.
+        """
+        if output_split_idx is None:
+            with self._lock:
+                return len(self._queue) > 0
+        else:
+            with self._lock:
+                return self._num_per_split[output_split_idx] > 0
+
+    def append(self, ref: RefBundle):
+        """Append a RefBundle to the queue."""
+        with self._lock:
+            self._queue.add(ref)
+            self._num_blocks += len(ref.blocks)
+            if ref.output_split_idx is not None:
+                self._num_per_split[ref.output_split_idx] += 1
+
+    def pop(self, output_split_idx: Optional[int] = None) -> Optional[RefBundle]:
+        """Pop a RefBundle from the queue.
+        Args:
+            output_split_idx: If specified, only pop a RefBundle
+                with the given output split.
+        Returns:
+            A RefBundle if available, otherwise None.
+        """
+        ret = None
+        if output_split_idx is None:
+            try:
+                with self._lock:
+                    ret = self._queue.get_next()
+            except IndexError:
+                pass
+        else:
+            with self._lock:
+                split_queue = self._outputs_by_split[output_split_idx]
+            if len(split_queue) == 0:
+                # Move all ref bundles to their indexed queues
+                # Note, the reason why we do indexing here instead of in the append
+                # is because only the last `OpBufferQueue` in the DAG, which will call
+                # pop with output_split_idx, needs indexing.
+                # If we also index the `OpBufferQueue`s in the middle, we cannot
+                # preserve the order of ref bundles with different output splits.
+                with self._lock:
+                    while len(self._queue) > 0:
+                        ref = self._queue.get_next()
+                        self._outputs_by_split[ref.output_split_idx].add(ref)
+            try:
+                ret = split_queue.get_next()
+            except IndexError:
+                pass
+        if ret is None:
+            return None
+        with self._lock:
+            self._num_blocks -= len(ret.blocks)
+            if ret.output_split_idx is not None:
+                self._num_per_split[ret.output_split_idx] -= 1
+        return ret
+
+    def clear(self):
+        with self._lock:
+            self._queue.clear()
+            self._num_blocks = 0
+            self._num_per_split.clear()
+
+
+@dataclass
+class OpSchedulingStatus:
+    """The scheduling status of an operator.
+
+    This will be updated each time when StreamingExecutor makes
+    a scheduling decision, i.e., in each `select_operator_to_run`
+    call.
+    """
+
+    # Whether the op was considered runnable in the last scheduling
+    # decision.
+    runnable: bool = False
+    # Whether the resources were sufficient for the operator to run
+    # in the last scheduling decision.
+    under_resource_limits: bool = False
+
+
+@dataclass
+class OpDisplayMetrics:
+    """Metrics of an operator. Used for display purposes."""
+
+    cpu: float = 0.0
+    gpu: float = 0.0
+    object_store_memory: float = 0.0
+    tasks: int = 0
+    actors: int = 0
+    queued: int = 0
+    task_backpressured: bool = False
+    output_backpressured: bool = False
+    extra_info: str = ""
+
+    def display_str(self) -> str:
+        """Format metrics object to a displayable string."""
+        metrics = []
+        # resource metrics
+        gpu_str = f" {self.gpu:.1f} GPU," if self.gpu else ""
+        mem_str = memory_string(self.object_store_memory)
+        metrics.append(f"{self.cpu:.1f} CPU,{gpu_str} {mem_str} object store")
+        # task
+        task_str = f"Tasks: {self.tasks}"
+        if self.task_backpressured or self.output_backpressured:
+            backpressured = []
+            if self.task_backpressured:
+                backpressured.append("tasks")
+            if self.output_backpressured:
+                backpressured.append("outputs")
+            task_str += f" [backpressured: {','.join(backpressured)}]"
+        if self.extra_info:
+            task_str += f": {self.extra_info}"
+        metrics.append(task_str)
+        # actors
+        if self.actors:
+            metrics.append(f"Actors: {self.actors}")
+        # queue
+        metrics.append(f"Queued blocks: {self.queued}")
+        return "; ".join(metrics) + ";"
+
+
+class OpState:
+    """The execution state tracked for each PhysicalOperator.
+
+    This tracks state to manage input and output buffering for StreamingExecutor and
+    progress bars, which is separate from execution state internal to the operators.
+
+    Note: we use the `deque` data structure here because it is thread-safe, enabling
+    operator queues to be shared across threads.
+    """
+
+    def __init__(self, op: PhysicalOperator, inqueues: List[OpBufferQueue]):
+        # Each input queue is connected to another operator's output queue.
+        assert len(inqueues) == len(op.input_dependencies), (op, inqueues)
+        self.input_queues: List[OpBufferQueue] = inqueues
+        # The output queue is connected to another operator's input queue (same object).
+        #
+        # Note: this queue is also accessed concurrently from the consumer thread.
+        # (in addition to the streaming executor thread). Hence, it must be a
+        # thread-safe type such as `deque`.
+        self.output_queue: OpBufferQueue = OpBufferQueue()
+        self.op = op
+        self.progress_bar = None
+        self.num_completed_tasks = 0
+        self.inputs_done_called = False
+        # Tracks whether `input_done` is called for each input op.
+        self.input_done_called = [False] * len(op.input_dependencies)
+        # Used for StreamingExecutor to signal exception or end of execution
+        self._finished: bool = False
+        self._exception: Optional[Exception] = None
+        self._scheduling_status = OpSchedulingStatus()
+        self._schema: Optional["Schema"] = None
+        self._warned_on_schema_divergence: bool = False
+        # Progress Manager
+        self.op_display_metrics = OpDisplayMetrics()
+        self.progress_manager_uuid: Optional[UUID] = None
+        self.output_row_count: int = 0
+
+    def __repr__(self):
+        return f"OpState({self.op.name})"
+
+    def initialize_progress_bars(self, index: int, verbose_progress: bool) -> int:
+        """Create progress bars at the given index (line offset in console).
+
+        For AllToAllOperator, zero or more sub progress bar would be created.
+        Return the number of enabled progress bars created for this operator.
+        """
+        contains_sub_progress_bars = isinstance(
+            self.op, AllToAllOperator
+        ) or isinstance(self.op, HashShuffleProgressBarMixin)
+        # Only show 1:1 ops when in verbose progress mode.
+
+        ctx = DataContext.get_current()
+        progress_bar_enabled = (
+            ctx.enable_progress_bars
+            and ctx.enable_operator_progress_bars
+            and (contains_sub_progress_bars or verbose_progress)
+        )
+        self.progress_bar = ProgressBar(
+            "- " + self.op.name,
+            self.op.num_output_rows_total(),
+            unit="row",
+            position=index,
+            enabled=progress_bar_enabled,
+        )
+        num_progress_bars = 1
+        if contains_sub_progress_bars:
+            # Initialize must be called for sub progress bars, even the
+            # bars are not enabled via the DataContext.
+            num_progress_bars += self.op.initialize_sub_progress_bars(index + 1)
+        return num_progress_bars if progress_bar_enabled else 0
+
+    def close_progress_bars(self):
+        """Close all progress bars for this operator."""
+        if self.progress_bar:
+            self.progress_bar.close()
+            contains_sub_progress_bars = isinstance(
+                self.op, AllToAllOperator
+            ) or isinstance(self.op, HashShuffleProgressBarMixin)
+            if contains_sub_progress_bars:
+                # Close all sub progress bars.
+                self.op.close_sub_progress_bars()
+
+    def total_enqueued_input_blocks(self) -> int:
+        """Total number of blocks currently enqueued among:
+        1. Input queue(s) pending dispatching (``OpState.input_queues``)
+        2. Operator's internal queues (like ``MapOperator``s ref-bundler, etc)
+        """
+        external_queue_size = sum(q.num_blocks for q in self.input_queues)
+        internal_queue_size = (
+            self.op.internal_input_queue_num_blocks()
+            if isinstance(self.op, InternalQueueOperatorMixin)
+            else 0
+        )
+        return external_queue_size + internal_queue_size
+
+    def has_pending_bundles(self) -> bool:
+        return any(len(q) > 0 for q in self.input_queues)
+
+    def total_enqueued_input_blocks_bytes(self) -> int:
+        """Total number of bytes occupied by input bundles currently enqueued among:
+        1. Input queue(s) pending dispatching (``OpState.input_queues``)
+        2. Operator's internal queues (like ``MapOperator``s ref-bundler, etc)
+        """
+        internal_queue_size_bytes = (
+            self.op.internal_input_queue_num_bytes()
+            if isinstance(self.op, InternalQueueOperatorMixin)
+            else 0
+        )
+        return self.input_queue_bytes() + internal_queue_size_bytes
+
+    def update_display_metrics(self, resource_manager: ResourceManager):
+        """Update display metrics with current metrics."""
+        usage = resource_manager.get_op_usage(self.op)
+        self.op_display_metrics.cpu = usage.cpu
+        self.op_display_metrics.gpu = usage.gpu
+        self.op_display_metrics.object_store_memory = usage.object_store_memory
+
+        self.op_display_metrics.tasks = self.op.num_active_tasks()
+        self.op_display_metrics.queued = self.total_enqueued_input_blocks()
+        self.op_display_metrics.actors = self.op.get_actor_info().running
+
+        self.op_display_metrics.task_backpressured = (
+            self.op._in_task_submission_backpressure
+        )
+        self.op_display_metrics.output_backpressured = (
+            self.op._in_task_output_backpressure
+        )
+
+        self.op_display_metrics.extra_info = self.op.progress_str()
+
+    def add_output(self, ref: RefBundle) -> None:
+        """Move a bundle produced by the operator to its outqueue."""
+
+        ref, diverged = dedupe_schemas_with_validation(
+            self._schema,
+            ref,
+            warn=not self._warned_on_schema_divergence,
+            enforce_schemas=self.op.data_context.enforce_schemas,
+        )
+
+        self._schema = ref.schema
+        self._warned_on_schema_divergence |= diverged
+
+        self.output_queue.append(ref)
+        self.num_completed_tasks += 1
+
+        actor_info = self.op.get_actor_info()
+        if ref.num_rows() is not None:
+            self.output_row_count += ref.num_rows()
+            if self.progress_bar:
+                self.progress_bar.update(
+                    ref.num_rows(), self.op.num_output_rows_total()
+                )
+
+        self.op.metrics.num_alive_actors = actor_info.running
+        self.op.metrics.num_restarting_actors = actor_info.restarting
+        self.op.metrics.num_pending_actors = actor_info.pending
+        for next_op in self.op.output_dependencies:
+            next_op.metrics.num_external_inqueue_blocks += len(ref.blocks)
+            next_op.metrics.num_external_inqueue_bytes += ref.size_bytes()
+        self.op.metrics.num_external_outqueue_blocks += len(ref.blocks)
+        self.op.metrics.num_external_outqueue_bytes += ref.size_bytes()
+
+    def refresh_progress_bar(self, resource_manager: ResourceManager) -> None:
+        """Update the console with the latest operator progress."""
+        if self.progress_bar:
+            self.progress_bar.set_description(self.summary_str(resource_manager))
+            self.progress_bar.refresh()
+
+    def summary_str(
+        self, resource_manager: ResourceManager, verbose: bool = False
+    ) -> str:
+        # Active tasks
+        active = self.op.num_active_tasks()
+        desc = f"- {self.op.name}: Tasks: {active}"
+        if (
+            self.op._in_task_submission_backpressure
+            or self.op._in_task_output_backpressure
+        ):
+            backpressure_types = []
+            if self.op._in_task_submission_backpressure:
+                # The op is backpressured from submitting new tasks.
+                backpressure_types.append("tasks")
+            if self.op._in_task_output_backpressure:
+                # The op is backpressured from producing new outputs.
+                backpressure_types.append("outputs")
+            desc += f" [backpressured:{','.join(backpressure_types)}]"
+
+        # Actors info
+        desc += f"; {_actor_info_summary_str(self.op.get_actor_info())}"
+
+        # Queued blocks
+        desc += f"; Queued blocks: {self.total_enqueued_input_blocks()} ({memory_string(self.total_enqueued_input_blocks_bytes())})"
+        desc += f"; Resources: {resource_manager.get_op_usage_str(self.op, verbose=verbose)}"
+
+        # Any additional operator specific information.
+        suffix = self.op.progress_str()
+        if suffix:
+            desc += f"; {suffix}"
+
+        return desc
+
+    def dispatch_next_task(self) -> None:
+        """Move a bundle from the operator inqueue to the operator itself."""
+        for i, inqueue in enumerate(self.input_queues):
+            ref = inqueue.pop()
+            if ref is not None:
+                self.op.add_input(ref, input_index=i)
+                self.op.metrics.num_external_inqueue_bytes -= ref.size_bytes()
+                self.op.metrics.num_external_inqueue_blocks -= len(ref.blocks)
+                input_op = self.op.input_dependencies[i]
+                # TODO: This needs to be cleaned up.
+                # the input_op's output queue = curr_op's input queue
+                input_op.metrics.num_external_outqueue_blocks -= len(ref.blocks)
+                input_op.metrics.num_external_outqueue_bytes -= ref.size_bytes()
+                return
+
+        assert False, "Nothing to dispatch"
+
+    def get_output_blocking(self, output_split_idx: Optional[int]) -> RefBundle:
+        """Get an item from this node's output queue, blocking as needed.
+
+        Returns:
+            The RefBundle from the output queue, or an error / end of stream indicator.
+
+        Raises:
+            StopIteration: If all outputs are already consumed.
+            Exception: If there was an exception raised during execution.
+        """
+        while True:
+            # Check if StreamingExecutor has caught an exception or is done execution.
+            if self._exception is not None:
+                raise self._exception
+            elif self._finished and not self.output_queue.has_next(output_split_idx):
+                raise StopIteration()
+            ref = self.output_queue.pop(output_split_idx)
+            if ref is not None:
+                # Update outqueue metrics when blocks are removed from this operator's outqueue
+                # TODO: Abstract queue-releated metrics to queue.
+                self.op.metrics.num_external_outqueue_blocks -= len(ref.blocks)
+                self.op.metrics.num_external_outqueue_bytes -= ref.size_bytes()
+                return ref
+            time.sleep(0.01)
+
+    def input_queue_bytes(self) -> int:
+        """Return the object store memory of this operator's inqueue."""
+        total = 0
+        for op, inq in zip(self.op.input_dependencies, self.input_queues):
+            # Exclude existing input data items from dynamic memory usage.
+            if not isinstance(op, InputDataBuffer):
+                total += inq.memory_usage
+        return total
+
+    def output_queue_bytes(self) -> int:
+        """Return the object store memory of this operator's outqueue."""
+        return self.output_queue.memory_usage
+
+    def mark_finished(self, exception: Optional[Exception] = None):
+        """Marks this operator as finished. Used for exiting get_output_blocking."""
+        if exception is None:
+            self._finished = True
+        else:
+            self._exception = exception
+
+
+def build_streaming_topology(
+    dag: PhysicalOperator, options: ExecutionOptions
+) -> Topology:
+    """Instantiate the streaming operator state topology for the given DAG.
+
+    This involves creating the operator state for each operator in the DAG,
+    registering it with this class, and wiring up the inqueues/outqueues of
+    dependent operator states.
+
+    Args:
+        dag: The operator DAG to instantiate.
+        options: The execution options to use to start operators.
+
+    Returns:
+        The topology dict holding the streaming execution state.
+    """
+
+    topology: Topology = {}
+
+    # DFS walk to wire up operator states.
+    def setup_state(op: PhysicalOperator) -> OpState:
+        if op in topology:
+            raise ValueError("An operator can only be present in a topology once.")
+
+        # Wire up the input outqueues to this op's inqueues.
+        inqueues = []
+        for parent in op.input_dependencies:
+            parent_state = setup_state(parent)
+            inqueues.append(parent_state.output_queue)
+
+        # Create state.
+        op_state = OpState(op, inqueues)
+        topology[op] = op_state
+        op.start(options)
+        return op_state
+
+    setup_state(dag)
+    return topology
+
+
+def process_completed_tasks(
+    topology: Topology,
+    backpressure_policies: List[BackpressurePolicy],
+    max_errored_blocks: int,
+) -> int:
+    """Process any newly completed tasks. To update operator
+    states, call `update_operator_states()` afterwards.
+
+    Args:
+        topology: The toplogy of operators.
+        backpressure_policies: The backpressure policies to use.
+        max_errored_blocks: Max number of errored blocks to allow,
+            unlimited if negative.
+    Returns:
+        The number of errored blocks.
+    """
+
+    # All active tasks, keyed by their waitables.
+    active_tasks: Dict[Waitable, Tuple[OpState, OpTask]] = {}
+    for op, state in topology.items():
+        for task in op.get_active_tasks():
+            active_tasks[task.get_waitable()] = (state, task)
+
+    max_bytes_to_read_per_op: Dict[OpState, int] = {}
+    for op, state in topology.items():
+        # Check all backpressure policies for max_task_output_bytes_to_read
+        # Use the minimum limit from all policies (most restrictive)
+        max_bytes_to_read = None
+        for policy in backpressure_policies:
+            policy_limit = policy.max_task_output_bytes_to_read(op)
+            if policy_limit is not None:
+                if max_bytes_to_read is None:
+                    max_bytes_to_read = policy_limit
+                else:
+                    max_bytes_to_read = min(max_bytes_to_read, policy_limit)
+
+        # If no policy provides a limit, there's no limit
+        op.notify_in_task_output_backpressure(max_bytes_to_read == 0)
+        if max_bytes_to_read is not None:
+            max_bytes_to_read_per_op[state] = max_bytes_to_read
+
+    # Process completed Ray tasks and notify operators.
+    num_errored_blocks = 0
+    if active_tasks:
+        ready, _ = ray.wait(
+            list(active_tasks.keys()),
+            num_returns=len(active_tasks),
+            fetch_local=False,
+            timeout=0.1,
+        )
+
+        # Organize tasks by the operator they belong to, and sort them by task index.
+        # So that we'll process them in a deterministic order.
+        # This is because backpressure policies may limit the number of blocks to read
+        # per operator. In this case, we want to have fewer tasks finish quickly and
+        # yield resources, instead of having all tasks output blocks together.
+        ready_tasks_by_op = defaultdict(list)
+        for ref in ready:
+            state, task = active_tasks[ref]
+            ready_tasks_by_op[state].append(task)
+
+        for state, ready_tasks in ready_tasks_by_op.items():
+            ready_tasks = sorted(ready_tasks, key=lambda t: t.task_index())
+            for task in ready_tasks:
+                if isinstance(task, DataOpTask):
+                    try:
+                        bytes_read = task.on_data_ready(
+                            max_bytes_to_read_per_op.get(state, None)
+                        )
+                        if state in max_bytes_to_read_per_op:
+                            max_bytes_to_read_per_op[state] -= bytes_read
+                    except Exception as e:
+                        num_errored_blocks += 1
+                        should_ignore = (
+                            max_errored_blocks < 0
+                            or max_errored_blocks >= num_errored_blocks
+                        )
+                        error_message = (
+                            "An exception was raised from a task of "
+                            f'operator "{state.op.name}".'
+                        )
+                        if should_ignore:
+                            remaining = (
+                                max_errored_blocks - num_errored_blocks
+                                if max_errored_blocks >= 0
+                                else "unlimited"
+                            )
+                            error_message += (
+                                " Ignoring this exception with remaining"
+                                f" max_errored_blocks={remaining}."
+                            )
+                            logger.error(error_message, exc_info=e)
+                        else:
+                            error_message += (
+                                " Dataset execution will now abort."
+                                " To ignore this exception and continue, set"
+                                " DataContext.max_errored_blocks."
+                            )
+                            logger.exception(error_message)
+                            raise e from None
+                else:
+                    assert isinstance(task, MetadataOpTask)
+                    task.on_task_finished()
+
+    # Pull any operator outputs into the streaming op state.
+    for op, op_state in topology.items():
+        while op.has_next():
+            op_state.add_output(op.get_next())
+
+    return num_errored_blocks
+
+
+def update_operator_states(topology: Topology) -> None:
+    """Update operator states accordingly for newly completed tasks.
+    Should be called after `process_completed_tasks()`."""
+
+    for op, op_state in topology.items():
+
+        # Call inputs_done() on ops where no more inputs are coming.
+        if op_state.inputs_done_called:
+            continue
+        all_inputs_done = True
+        for idx, dep in enumerate(op.input_dependencies):
+            if dep.completed() and not topology[dep].output_queue:
+                if not op_state.input_done_called[idx]:
+                    op.input_done(idx)
+                    op_state.input_done_called[idx] = True
+            else:
+                all_inputs_done = False
+
+        if all_inputs_done:
+            op.all_inputs_done()
+            op_state.inputs_done_called = True
+
+    # Traverse the topology in reverse topological order.
+    # For each op, if all of its downstream operators have completed.
+    # call mark_execution_finished() to also complete this op.
+    for op, op_state in reversed(list(topology.items())):
+
+        dependents_completed = len(op.output_dependencies) > 0 and all(
+            dep.completed() for dep in op.output_dependencies
+        )
+        if dependents_completed:
+            op.mark_execution_finished()
+
+        # Drain external input queue if current operator is execution finished.
+        # This is needed when the limit is reached, and `mark_execution_finished`
+        # is called manually.
+        if op.execution_finished():
+            for input_queue in op_state.input_queues:
+                # Drain input queue
+                input_queue.clear()
+
+
+def get_eligible_operators(
+    topology: Topology,
+    backpressure_policies: List[BackpressurePolicy],
+    *,
+    ensure_liveness: bool,
+) -> List[PhysicalOperator]:
+    """This method returns all operators that are eligible for execution in the current state
+    of the pipeline.
+
+    Operator is considered eligible for execution iff:
+
+        1. It's NOT completed
+        2. It has at least 1 input block (in the input queue)
+        3. It can accept new inputs
+        4. It's not currently throttled (for task-submission)
+
+    """
+
+    dispatchable_ops: List[PhysicalOperator] = []
+    # Filter to ops that are eligible for execution, ie ones that are
+    #   - Dispatchable
+    #   - Not throttled
+    eligible_ops: List[PhysicalOperator] = []
+
+    for op, state in topology.items():
+        # Operator is considered being in task-submission back-pressure if any
+        # back-pressure policy is violated
+        in_backpressure = any(not p.can_add_input(op) for p in backpressure_policies)
+
+        op_runnable = False
+
+        # Check whether operator could start executing immediately:
+        #   - It's not completed
+        #   - It can accept at least one input
+        #   - Its input queue has a valid bundle
+        if not op.completed() and op.should_add_input() and state.has_pending_bundles():
+            if not in_backpressure:
+                op_runnable = True
+                eligible_ops.append(op)
+            else:
+                dispatchable_ops.append(op)
+
+        # Update scheduling status
+        state._scheduling_status = OpSchedulingStatus(
+            runnable=op_runnable,
+            under_resource_limits=not in_backpressure,
+        )
+
+        # Signal whether op in backpressure for stats collections
+        # TODO(hchen): also report which policy triggers backpressure.
+        op.notify_in_task_submission_backpressure(in_backpressure)
+
+    # To ensure liveness, allow at least 1 operator to schedule tasks regardless of
+    # limits in case when topology is entirely idle (no active tasks running)
+    if (
+        not eligible_ops
+        and ensure_liveness
+        and all(op.num_active_tasks() == 0 for op in topology)
+    ):
+        return dispatchable_ops
+
+    return eligible_ops
+
+
+def select_operator_to_run(
+    topology: Topology,
+    resource_manager: ResourceManager,
+    backpressure_policies: List[BackpressurePolicy],
+    ensure_liveness: bool,
+    ranker: "Ranker",
+) -> Optional[PhysicalOperator]:
+    """Select next operator to launch new tasks.
+
+    The objective of this method is to maximize the throughput of the overall
+    pipeline, subject to defined memory, parallelism and other constraints.
+
+    To achieve that this method implements following protocol:
+
+        1. Collects all _eligible_ to run operators (check `_get_eligible_ops`
+           for more details)
+        2. Applies stack-ranking algorithm to select the best operator (check
+           `_create_eligible_ops_ranker` for more details)
+
+    """
+    eligible_ops = get_eligible_operators(
+        topology,
+        backpressure_policies,
+        ensure_liveness=ensure_liveness,
+    )
+
+    if not eligible_ops:
+        return None
+
+    ranks = ranker.rank_operators(eligible_ops, topology, resource_manager)
+
+    assert len(eligible_ops) == len(ranks), (eligible_ops, ranks)
+
+    next_op, _ = min(zip(eligible_ops, ranks), key=lambda t: t[1])
+
+    return next_op
+
+
+def _actor_info_summary_str(info: _ActorPoolInfo) -> str:
+    total = info.running + info.pending + info.restarting
+    base = f"Actors: {total}"
+
+    if total == info.running:
+        return base
+    else:
+        return f"{base} ({info})"
+
+
+def dedupe_schemas_with_validation(
+    old_schema: Optional["Schema"],
+    bundle: "RefBundle",
+    warn: bool = True,
+    enforce_schemas: bool = False,
+) -> Tuple["RefBundle", bool]:
+    """Unify/Dedupe two schemas, warning if warn=True
+
+    Args:
+        old_schema: The old schema to unify. This can be `None`, in which case
+            the new schema will be used as the old schema.
+        bundle: The new `RefBundle` to unify with the old schema.
+        warn: Raise a warning if the schemas diverge.
+        enforce_schemas: If `True`, allow the schemas to diverge and return unified schema.
+            If `False`, but keep the old schema.
+
+    Returns:
+        A ref bundle with the unified schema of the two input schemas.
+    """
+
+    # Note, often times the refbundles correspond to only one schema. We can reduce the
+    # memory footprint of multiple schemas by keeping only one copy.
+    diverged = False
+
+    from ray.data.block import _is_empty_schema
+
+    if _is_empty_schema(old_schema):
+        return bundle, diverged
+
+    # This check is fast assuming pyarrow schemas
+    if old_schema == bundle.schema:
+        return bundle, diverged
+
+    diverged = True
+    if warn and enforce_schemas:
+        logger.warning(
+            f"Operator produced a RefBundle with a different schema "
+            f"than the previous one. Previous schema: {old_schema}, "
+            f"new schema: {bundle.schema}. This may lead to unexpected behavior."
+        )
+    if enforce_schemas:
+        old_schema = unify_schemas_with_validation([old_schema, bundle.schema])
+
+    return (
+        RefBundle(
+            bundle.blocks,
+            schema=old_schema,
+            owns_blocks=bundle.owns_blocks,
+            output_split_idx=bundle.output_split_idx,
+            _cached_object_meta=bundle._cached_object_meta,
+            _cached_preferred_locations=bundle._cached_preferred_locations,
+        ),
+        diverged,
+    )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/execution/util.py b/lib/python3.12/site-packages/ray/data/_internal/execution/util.py
new file mode 100644
index 0000000000000000000000000000000000000000..be64d2a6ed649f2552ef12cd8a156160d936ba51
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/execution/util.py
@@ -0,0 +1,82 @@
+from concurrent.futures import ThreadPoolExecutor
+from typing import TYPE_CHECKING, Any, List
+
+import ray
+from ray.data.block import BlockAccessor, CallableClass
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces import RefBundle
+
+
+def make_ref_bundles(simple_data: List[List[Any]]) -> List["RefBundle"]:
+    """Create ref bundles from a list of block data.
+
+    One bundle is created for each input block.
+    """
+    import pandas as pd
+    import pyarrow as pa
+
+    from ray.data._internal.execution.interfaces import RefBundle
+
+    output = []
+    for block in simple_data:
+        block = pd.DataFrame({"id": block})
+        output.append(
+            RefBundle(
+                [
+                    (
+                        ray.put(block),
+                        BlockAccessor.for_block(block).get_metadata(),
+                    )
+                ],
+                owns_blocks=True,
+                schema=pa.lib.Schema.from_pandas(block, preserve_index=False),
+            )
+        )
+    return output
+
+
+memory_units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
+
+
+def memory_string(num_bytes: float) -> str:
+    """Return a human-readable memory string for the given amount of bytes."""
+    k = 0
+    while num_bytes >= 1024 and k < len(memory_units) - 1:
+        num_bytes /= 1024
+        k += 1
+    return f"{num_bytes:.1f}{memory_units[k]}"
+
+
+def locality_string(locality_hits: int, locality_misses) -> str:
+    """Return a human-readable string for object locality stats."""
+    if not locality_misses:
+        return "[all objects local]"
+    return f"[{locality_hits}/{locality_hits + locality_misses} objects local]"
+
+
+def make_callable_class_concurrent(callable_cls: CallableClass) -> CallableClass:
+    """Returns a thread-safe CallableClass with the same logic as the provided
+    `callable_cls`.
+
+    This function allows the usage of concurrent actors by safeguarding user logic
+    behind a separate thread.
+
+    This allows batch slicing and formatting to occur concurrently, to overlap with the
+    user provided UDF.
+    """
+
+    class _Wrapper(callable_cls):
+        def __init__(self, *args, **kwargs):
+            self.thread_pool_executor = ThreadPoolExecutor(max_workers=1)
+            super().__init__(*args, **kwargs)
+
+        def __repr__(self):
+            return super().__repr__()
+
+        def __call__(self, *args, **kwargs):
+            # ThreadPoolExecutor will reuse the same thread for every submit call.
+            future = self.thread_pool_executor.submit(super().__call__, *args, **kwargs)
+            return future.result()
+
+    return _Wrapper
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0a246f9d191d244d0258b891a4e430d3acedb06b
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__init__.py
@@ -0,0 +1,20 @@
+from ray.data._internal.issue_detection.detectors.hanging_detector import (
+    HangingExecutionIssueDetector,
+    HangingExecutionIssueDetectorConfig,
+)
+from ray.data._internal.issue_detection.issue_detector import Issue, IssueDetector
+from ray.data._internal.issue_detection.issue_detector_configuration import (
+    IssueDetectorsConfiguration,
+)
+from ray.data._internal.issue_detection.issue_detector_manager import (
+    IssueDetectorManager,
+)
+
+__all__ = [
+    "Issue",
+    "IssueDetector",
+    "IssueDetectorManager",
+    "IssueDetectorsConfiguration",
+    "HangingExecutionIssueDetector",
+    "HangingExecutionIssueDetectorConfig",
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2db82aa5ce7d5739a4e34c4ce19b83d452ae3642
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3d27b75f8438284bdbf9bbdafa428740298690d2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector_configuration.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector_configuration.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7ee89434402754c83eeb5c873fc813e6bfbf803f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector_configuration.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector_manager.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e7015c1126cf2cc5e5637fe6451f630007a7317c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/__pycache__/issue_detector_manager.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fd93c5ce9d9716db22bebc8bc5ee89c21291b25
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__init__.py
@@ -0,0 +1,19 @@
+from ray.data._internal.issue_detection.detectors.hanging_detector import (
+    HangingExecutionIssueDetector,
+    HangingExecutionIssueDetectorConfig,
+)
+from ray.data._internal.issue_detection.detectors.hash_shuffle_detector import (
+    HashShuffleAggregatorIssueDetector,
+)
+from ray.data._internal.issue_detection.detectors.high_memory_detector import (
+    HighMemoryIssueDetector,
+    HighMemoryIssueDetectorConfig,
+)
+
+__all__ = [
+    "HangingExecutionIssueDetector",
+    "HangingExecutionIssueDetectorConfig",
+    "HighMemoryIssueDetector",
+    "HighMemoryIssueDetectorConfig",
+    "HashShuffleAggregatorIssueDetector",
+]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ee3aa04dad55aae8503cff576879c76b7ce34fc4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/hanging_detector.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/hanging_detector.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3b88f32d464df00736514b711773e9078795a9fb
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/hanging_detector.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/hash_shuffle_detector.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/hash_shuffle_detector.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..37104d8d4b14fe91062be5bbdeaeda9072543639
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/hash_shuffle_detector.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/high_memory_detector.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/high_memory_detector.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fcadf4b336a704d278d617d7ab3205b08bbfbc14
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/__pycache__/high_memory_detector.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/hanging_detector.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/hanging_detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..381ba05f7bc165f2b683ef7fbea339c98ace6a3c
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/hanging_detector.py
@@ -0,0 +1,163 @@
+import time
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Set
+
+from ray.data._internal.issue_detection.issue_detector import (
+    Issue,
+    IssueDetector,
+    IssueType,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.op_runtime_metrics import (
+        TaskDurationStats,
+    )
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+    from ray.data.context import DataContext
+
+# Default minimum count of tasks before using adaptive thresholds
+DEFAULT_OP_TASK_STATS_MIN_COUNT = 10
+# Default multiple of standard deviations to use as hanging threshold
+DEFAULT_OP_TASK_STATS_STD_FACTOR = 10
+# Default detection time interval.
+DEFAULT_DETECTION_TIME_INTERVAL_S = 30.0
+
+
+@dataclass
+class HangingExecutionState:
+    operator_id: str
+    task_idx: int
+    bytes_output: int
+    start_time_hanging: float
+
+
+@dataclass
+class HangingExecutionIssueDetectorConfig:
+    op_task_stats_min_count: int = field(default=DEFAULT_OP_TASK_STATS_MIN_COUNT)
+    op_task_stats_std_factor: float = field(default=DEFAULT_OP_TASK_STATS_STD_FACTOR)
+    detection_time_interval_s: float = DEFAULT_DETECTION_TIME_INTERVAL_S
+
+
+class HangingExecutionIssueDetector(IssueDetector):
+    def __init__(self, executor: "StreamingExecutor", ctx: "DataContext"):
+        super().__init__(executor, ctx)
+        self._detector_cfg: HangingExecutionIssueDetectorConfig = (
+            ctx.issue_detectors_config.hanging_detector_config
+        )
+        self._op_task_stats_min_count = self._detector_cfg.op_task_stats_min_count
+        self._op_task_stats_std_factor_threshold = (
+            self._detector_cfg.op_task_stats_std_factor
+        )
+
+        # Map of operator id to dict of task_idx to hanging execution info (bytes read and
+        # start time for hanging time calculation)
+        self._state_map: Dict[str, Dict[int, HangingExecutionState]] = defaultdict(dict)
+        # Map of operator id to set of task_idx that are hanging
+        self._hanging_op_tasks: Dict[str, Set[int]] = defaultdict(set)
+        # Map of operator id to operator name
+        self._op_id_to_name: Dict[str, str] = {}
+
+    def _create_issues(
+        self,
+        hanging_op_tasks: List[HangingExecutionState],
+        op_task_stats_map: Dict[str, "TaskDurationStats"],
+    ) -> List[Issue]:
+        issues = []
+        for state in hanging_op_tasks:
+            if state.task_idx not in self._hanging_op_tasks[state.operator_id]:
+                op_name = self._op_id_to_name.get(state.operator_id, state.operator_id)
+                duration = time.perf_counter() - state.start_time_hanging
+                avg_duration = op_task_stats_map[state.operator_id].mean()
+                message = (
+                    f"A task of operator {op_name} with task index "
+                    f"{state.task_idx} has been running for {duration:.2f}s, which is longer"
+                    f" than the average task duration of this operator ({avg_duration:.2f}s)."
+                    f" If this message persists, please check the stack trace of the "
+                    "task for potential hanging issues."
+                )
+                issues.append(
+                    Issue(
+                        dataset_name=self._executor._dataset_id,
+                        operator_id=state.operator_id,
+                        issue_type=IssueType.HANGING,
+                        message=message,
+                    )
+                )
+                self._hanging_op_tasks[state.operator_id].add(state.task_idx)
+
+        return issues
+
+    def detect(self) -> List[Issue]:
+        op_task_stats_map: Dict[str, "TaskDurationStats"] = {}
+        for operator, op_state in self._executor._topology.items():
+            op_metrics = operator.metrics
+            op_task_stats_map[operator.id] = op_metrics._op_task_duration_stats
+            self._op_id_to_name[operator.id] = operator.name
+            if op_state._finished:
+                # Remove finished operators / tasks from the state map
+                if operator.id in self._state_map:
+                    del self._state_map[operator.id]
+                if operator.id in self._hanging_op_tasks:
+                    del self._hanging_op_tasks[operator.id]
+            else:
+                active_tasks_idx = set()
+                for task in operator.get_active_tasks():
+                    task_info = op_metrics._running_tasks.get(task.task_index(), None)
+                    if task_info is None:
+                        # if the task is not in the running tasks map, it has finished
+                        # remove it from the state map and hanging op tasks, if present
+                        self._state_map[operator.id].pop(task.task_index(), None)
+                        self._hanging_op_tasks[operator.id].discard(task.task_index())
+                        continue
+
+                    active_tasks_idx.add(task.task_index())
+                    bytes_output = task_info.bytes_outputs
+
+                    prev_state_value = self._state_map[operator.id].get(
+                        task.task_index(), None
+                    )
+
+                    if (
+                        prev_state_value is None
+                        or bytes_output != prev_state_value.bytes_output
+                    ):
+                        self._state_map[operator.id][
+                            task.task_index()
+                        ] = HangingExecutionState(
+                            operator_id=operator.id,
+                            task_idx=task.task_index(),
+                            bytes_output=bytes_output,
+                            start_time_hanging=time.perf_counter(),
+                        )
+
+                # Remove any tasks that are no longer active
+                task_idxs_to_remove = (
+                    set(self._state_map[operator.id].keys()) - active_tasks_idx
+                )
+                for task_idx in task_idxs_to_remove:
+                    del self._state_map[operator.id][task_idx]
+
+        hanging_op_tasks = []
+        for op_id, op_state_values in self._state_map.items():
+            op_task_stats = op_task_stats_map[op_id]
+            for task_idx, state_value in op_state_values.items():
+                curr_time = time.perf_counter() - state_value.start_time_hanging
+                if op_task_stats.count() > self._op_task_stats_min_count:
+                    mean = op_task_stats.mean()
+                    stddev = op_task_stats.stddev()
+                    threshold = mean + self._op_task_stats_std_factor_threshold * stddev
+
+                    if curr_time > threshold:
+                        hanging_op_tasks.append(state_value)
+
+        # create issues for newly detected hanging tasks, then update the hanging task set
+        issues = self._create_issues(
+            hanging_op_tasks=hanging_op_tasks,
+            op_task_stats_map=op_task_stats_map,
+        )
+
+        return issues
+
+    def detection_time_interval_s(self) -> float:
+        return self._detector_cfg.detection_time_interval_s
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/hash_shuffle_detector.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/hash_shuffle_detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..345b5983aa38b67218c9be2cd0b287f68c9e2182
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/hash_shuffle_detector.py
@@ -0,0 +1,110 @@
+import time
+from typing import TYPE_CHECKING, List
+
+import ray
+from ray.data._internal.execution.operators.hash_shuffle import (
+    AggregatorHealthInfo,
+    HashShuffleOperator,
+)
+from ray.data._internal.issue_detection.issue_detector import (
+    Issue,
+    IssueDetector,
+    IssueType,
+)
+from ray.data._internal.util import GiB
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+    from ray.data.context import DataContext
+
+
+class HashShuffleAggregatorIssueDetector(IssueDetector):
+    """Detector for hash shuffle aggregator health issues."""
+
+    def __init__(self, executor: "StreamingExecutor", ctx: "DataContext"):
+        super().__init__(executor, ctx)
+        self._last_warning_times = {}  # Track per-operator warning times
+
+    def detect(self) -> List[Issue]:
+        issues = []
+        current_time = time.time()
+
+        # Find all hash shuffle operators in the topology
+        for op in self._executor._topology.keys():
+            if not isinstance(op, HashShuffleOperator):
+                continue
+
+            # Skip if operator doesn't have aggregator pool yet
+            if op._aggregator_pool is None:
+                continue
+
+            pool = op._aggregator_pool
+            aggregator_info = pool.check_aggregator_health()
+
+            if aggregator_info is None:
+                continue
+
+            # Check if we should emit a warning for this operator
+            should_warn = self._should_emit_warning(
+                op.id, current_time, aggregator_info
+            )
+
+            if should_warn:
+                message = self._format_health_warning(aggregator_info)
+                issues.append(
+                    Issue(
+                        dataset_name=self._executor._dataset_id,
+                        operator_id=op.id,
+                        issue_type=IssueType.HANGING,
+                        message=message,
+                    )
+                )
+                self._last_warning_times[op.id] = current_time
+
+        return issues
+
+    def detection_time_interval_s(self) -> float:
+        return self._ctx.hash_shuffle_aggregator_health_warning_interval_s
+
+    def _should_emit_warning(
+        self, op_id: str, current_time: float, info: AggregatorHealthInfo
+    ) -> bool:
+        """Check if we should emit a warning for this operator."""
+        if not info.has_unready_aggregators:
+            # Clear warning time if all aggregators are healthy
+            self._last_warning_times.pop(op_id, None)
+            return False
+
+        # Check if enough time has passed since start
+        if (
+            current_time - info.started_at
+            < self._ctx.min_hash_shuffle_aggregator_wait_time_in_s
+        ):
+            return False
+
+        # Check if enough time has passed since last warning
+        last_warning = self._last_warning_times.get(op_id)
+        if last_warning is None:
+            return True
+
+        return current_time - last_warning >= self.detection_time_interval_s()
+
+    def _format_health_warning(self, info: AggregatorHealthInfo) -> str:
+        """Format the health warning message."""
+        available_resources = ray.available_resources()
+        available_cpus = available_resources.get("CPU", 0)
+        cluster_resources = ray.cluster_resources()
+        total_memory = cluster_resources.get("memory", 0)
+        available_memory = available_resources.get("memory", 0)
+
+        return (
+            f"Only {info.ready_aggregators} out of {info.total_aggregators} "
+            f"hash-shuffle aggregators are ready after {info.wait_time:.1f} secs. "
+            f"This might indicate resource contention for cluster resources "
+            f"(available CPUs: {available_cpus}, required CPUs: {info.required_resources.cpu}). "
+            f"Cluster only has {available_memory / GiB:.2f} GiB available memory, required memory: {info.required_resources.memory / GiB:.2f} GiB. "
+            f"{total_memory / GiB:.2f} GiB total memory. "
+            f"Consider increasing cluster size or reducing the number of aggregators "
+            f"via `DataContext.max_hash_shuffle_aggregators`. "
+            f"Will continue checking every {self.detection_time_interval_s()}s."
+        )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/high_memory_detector.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/high_memory_detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b9a4b5d6ef3fc9bf5b2ad0fda0d703dcb568006
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/detectors/high_memory_detector.py
@@ -0,0 +1,101 @@
+import textwrap
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List
+
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.issue_detection.issue_detector import (
+    Issue,
+    IssueDetector,
+    IssueType,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+    from ray.data.context import DataContext
+
+HIGH_MEMORY_PERIODIC_WARNING = """
+Operator '{op_name}' uses {memory_per_task} of memory per task on average, but Ray
+only requests {initial_memory_request} per task at the start of the pipeline.
+
+To avoid out-of-memory errors, consider setting `memory={memory_per_task}` in the
+appropriate function or method call. (This might be unnecessary if the number of
+concurrent tasks is low.)
+
+To change the frequency of this warning, set
+`DataContext.get_current().issue_detectors_config.high_memory_detector_config.detection_time_interval_s`,
+or disable the warning by setting value to -1. (current value:
+{detection_time_interval_s})
+"""  # noqa: E501
+
+
+@dataclass
+class HighMemoryIssueDetectorConfig:
+    detection_time_interval_s: float = 30
+
+
+class HighMemoryIssueDetector(IssueDetector):
+
+    # Many nodes have a 4 GiB : 1 core ratio, but this isn't always the case (e.g., for
+    # high memory nodes).
+    _MEMORY_PER_CORE_ESTIMATE = 4 * 1024**3
+
+    def __init__(self, executor: "StreamingExecutor", ctx: "DataContext"):
+        super().__init__(executor, ctx)
+        self._detector_cfg = ctx.issue_detectors_config.high_memory_detector_config
+
+        self._initial_memory_requests: Dict[MapOperator, int] = {}
+        for op in self._executor._topology.keys():
+            if isinstance(op, MapOperator):
+                self._initial_memory_requests[op] = (
+                    op._get_dynamic_ray_remote_args().get("memory") or 0
+                )
+
+    def detect(self) -> List[Issue]:
+        issues = []
+        for op in self._executor._topology.keys():
+            if not isinstance(op, MapOperator):
+                continue
+
+            if op.metrics.average_max_uss_per_task is None:
+                continue
+
+            remote_args = op._get_dynamic_ray_remote_args()
+            num_cpus_per_task = remote_args.get("num_cpus", 1)
+            max_memory_per_task = self._MEMORY_PER_CORE_ESTIMATE * num_cpus_per_task
+
+            if (
+                op.metrics.average_max_uss_per_task > self._initial_memory_requests[op]
+                and op.metrics.average_max_uss_per_task >= max_memory_per_task
+            ):
+                message = HIGH_MEMORY_PERIODIC_WARNING.format(
+                    op_name=op.name,
+                    memory_per_task=memory_string(op.metrics.average_max_uss_per_task),
+                    initial_memory_request=memory_string(
+                        self._initial_memory_requests[op]
+                    ),
+                    detection_time_interval_s=self.detection_time_interval_s(),
+                )
+                issues.append(
+                    Issue(
+                        dataset_name=self._executor._dataset_id,
+                        operator_id=op.id,
+                        issue_type=IssueType.HIGH_MEMORY,
+                        message=_format_message(message),
+                    )
+                )
+
+        return issues
+
+    def detection_time_interval_s(self) -> float:
+        return self._detector_cfg.detection_time_interval_s
+
+
+def _format_message(message: str) -> str:
+    # Apply some formatting to make the message look nicer when printed.
+    formatted_paragraphs = []
+    for paragraph in message.split("\n\n"):
+        formatted_paragraph = textwrap.fill(paragraph, break_long_words=False).strip()
+        formatted_paragraphs.append(formatted_paragraph)
+    formatted_message = "\n\n".join(formatted_paragraphs)
+    return "\n\n" + formatted_message + "\n"
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector.py
new file mode 100644
index 0000000000000000000000000000000000000000..dbfab89f01c9c872ed8c3858087e8cb037c34af0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector.py
@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+    from ray.data.context import DataContext
+
+
+class IssueType(str, Enum):
+    HANGING = "hanging"
+    HIGH_MEMORY = "high memory"
+
+
+@dataclass
+class Issue:
+    dataset_name: str
+    operator_id: str
+    message: str
+    issue_type: IssueType
+
+
+class IssueDetector(ABC):
+    def __init__(self, executor: "StreamingExecutor", ctx: "DataContext"):
+        self._executor = executor
+        self._ctx = ctx
+
+    @abstractmethod
+    def detect(self) -> List[Issue]:
+        pass
+
+    @abstractmethod
+    def detection_time_interval_s(self) -> float:
+        """Time interval between detections, or -1 if not enabled."""
+        pass
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector_configuration.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector_configuration.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b59a7318d2324bc89df58fd86ae449692c8e470
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector_configuration.py
@@ -0,0 +1,23 @@
+from dataclasses import dataclass, field
+from typing import List, Type
+
+from ray.data._internal.issue_detection.detectors import (
+    HangingExecutionIssueDetector,
+    HangingExecutionIssueDetectorConfig,
+    HighMemoryIssueDetector,
+    HighMemoryIssueDetectorConfig,
+)
+from ray.data._internal.issue_detection.issue_detector import IssueDetector
+
+
+@dataclass
+class IssueDetectorsConfiguration:
+    hanging_detector_config: HangingExecutionIssueDetectorConfig = field(
+        default_factory=HangingExecutionIssueDetectorConfig
+    )
+    high_memory_detector_config: HighMemoryIssueDetectorConfig = field(
+        default_factory=HighMemoryIssueDetectorConfig
+    )
+    detectors: List[Type[IssueDetector]] = field(
+        default_factory=lambda: [HangingExecutionIssueDetector, HighMemoryIssueDetector]
+    )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector_manager.py b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector_manager.py
new file mode 100644
index 0000000000000000000000000000000000000000..33ebbc69dafe7364dcbea700f7c085883042a982
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/issue_detection/issue_detector_manager.py
@@ -0,0 +1,98 @@
+import logging
+import time
+from typing import TYPE_CHECKING, Dict, List
+
+from ray.core.generated.export_dataset_operator_event_pb2 import (
+    ExportDatasetOperatorEventData as ProtoOperatorEventData,
+)
+from ray.data._internal.issue_detection.issue_detector import (
+    Issue,
+    IssueDetector,
+    IssueType,
+)
+from ray.data._internal.operator_event_exporter import (
+    OperatorEvent,
+    format_export_issue_event_name,
+    get_operator_event_exporter,
+)
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.streaming_executor import StreamingExecutor
+
+logger = logging.getLogger(__name__)
+
+
+class IssueDetectorManager:
+    def __init__(self, executor: "StreamingExecutor"):
+        ctx = executor._data_context
+        self._issue_detectors: List[IssueDetector] = [
+            cls(executor, ctx) for cls in ctx.issue_detectors_config.detectors
+        ]
+        self._last_detection_times: Dict[IssueDetector, float] = {
+            detector: time.perf_counter() for detector in self._issue_detectors
+        }
+        self.executor = executor
+        self._operator_event_exporter = get_operator_event_exporter()
+
+    def invoke_detectors(self) -> None:
+        curr_time = time.perf_counter()
+        issues = []
+        for detector in self._issue_detectors:
+            if detector.detection_time_interval_s() == -1:
+                continue
+
+            if (
+                curr_time - self._last_detection_times[detector]
+                > detector.detection_time_interval_s()
+            ):
+                issues.extend(detector.detect())
+
+                self._last_detection_times[detector] = time.perf_counter()
+
+        self._report_issues(issues)
+
+    def _report_issues(self, issues: List[Issue]) -> None:
+        operators: Dict[str, "PhysicalOperator"] = {}
+        op_to_id: Dict["PhysicalOperator", str] = {}
+        for i, operator in enumerate(self.executor._topology.keys()):
+            operators[operator.id] = operator
+            op_to_id[operator] = self.executor._get_operator_id(operator, i)
+            # Reset issue detector metrics for each operator so that previous issues
+            # don't affect the current ones.
+            operator.metrics._issue_detector_hanging = 0
+            operator.metrics._issue_detector_high_memory = 0
+
+        for issue in issues:
+            logger.warning(issue.message)
+            operator = operators.get(issue.operator_id)
+            if not operator:
+                continue
+
+            issue_event_type = format_export_issue_event_name(issue.issue_type)
+            if (
+                self._operator_event_exporter is not None
+                and issue_event_type
+                in ProtoOperatorEventData.DatasetOperatorEventType.keys()
+            ):
+                event_time = time.time()
+                operator_event = OperatorEvent(
+                    dataset_id=issue.dataset_name,
+                    operator_id=op_to_id[operator],
+                    operator_name=operator.name,
+                    event_time=event_time,
+                    event_type=issue_event_type,
+                    message=issue.message,
+                )
+                self._operator_event_exporter.export_operator_event(operator_event)
+
+            if issue.issue_type == IssueType.HANGING:
+                operator.metrics._issue_detector_hanging += 1
+            if issue.issue_type == IssueType.HIGH_MEMORY:
+                operator.metrics._issue_detector_high_memory += 1
+        if len(issues) > 0:
+            logger.warning(
+                "To disable issue detection, run DataContext.get_current().issue_detectors_config.detectors = []."
+            )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/iterator/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/iterator/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e61509553bec3dd1b233e56b67880d0365908c1c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/iterator_impl.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/iterator_impl.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cc8f5487d43131257954a36668042440b2ab061e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/iterator_impl.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/stream_split_iterator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/stream_split_iterator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ca02a6aac246f7111cea8bbc0ef14ffd4d4f75b4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/iterator/__pycache__/stream_split_iterator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/iterator/iterator_impl.py b/lib/python3.12/site-packages/ray/data/_internal/iterator/iterator_impl.py
new file mode 100644
index 0000000000000000000000000000000000000000..919939a2f4af30a1cb2d1b8962f17b4789ca59ce
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/iterator/iterator_impl.py
@@ -0,0 +1,40 @@
+from typing import TYPE_CHECKING, Iterator, Optional, Tuple, Union
+
+from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
+from ray.data._internal.stats import DatasetStats
+from ray.data.context import DataContext
+from ray.data.iterator import DataIterator
+
+if TYPE_CHECKING:
+    import pyarrow
+
+    from ray.data.dataset import Dataset
+
+
+class DataIteratorImpl(DataIterator):
+    def __init__(
+        self,
+        base_dataset: "Dataset",
+    ):
+        self._base_dataset = base_dataset
+
+    def __repr__(self) -> str:
+        return f"DataIterator({self._base_dataset})"
+
+    def _to_ref_bundle_iterator(
+        self,
+    ) -> Tuple[Iterator[RefBundle], Optional[DatasetStats], bool]:
+        ref_bundles_iterator, stats = self._base_dataset._execute_to_iterator()
+        return ref_bundles_iterator, stats, False
+
+    def stats(self) -> str:
+        return self._base_dataset.stats()
+
+    def schema(self) -> Union[type, "pyarrow.lib.Schema"]:
+        return self._base_dataset.schema()
+
+    def get_context(self) -> DataContext:
+        return self._base_dataset.context
+
+    def _get_dataset_tag(self):
+        return self._base_dataset.get_dataset_id()
diff --git a/lib/python3.12/site-packages/ray/data/_internal/iterator/stream_split_iterator.py b/lib/python3.12/site-packages/ray/data/_internal/iterator/stream_split_iterator.py
new file mode 100644
index 0000000000000000000000000000000000000000..f94efe2e71b9ff78fec5a59c75b1e1146f35e140
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/iterator/stream_split_iterator.py
@@ -0,0 +1,289 @@
+import logging
+import threading
+import time
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Union
+
+import ray
+from ray.data._internal.execution.interfaces import NodeIdStr, RefBundle
+from ray.data._internal.execution.legacy_compat import execute_to_legacy_bundle_iterator
+from ray.data._internal.stats import DatasetStats
+from ray.data.block import Block
+from ray.data.context import DataContext
+from ray.data.iterator import DataIterator
+from ray.types import ObjectRef
+from ray.util.debug import log_once
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+if TYPE_CHECKING:
+    import pyarrow
+
+    from ray.data.dataset import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+BLOCKED_CLIENT_WARN_TIMEOUT = 30
+
+
+class _DatasetWrapper:
+    # A temporary workaround for https://github.com/ray-project/ray/issues/52549
+
+    def __init__(self, dataset: "Dataset") -> None:
+        self._dataset = dataset
+
+
+class StreamSplitDataIterator(DataIterator):
+    """Implements a collection of iterators over a shared data stream."""
+
+    @staticmethod
+    def create(
+        base_dataset: "Dataset",
+        n: int,
+        locality_hints: Optional[List[NodeIdStr]],
+    ) -> List["StreamSplitDataIterator"]:
+        """Create a split iterator from the given base Dataset and options.
+
+        See also: `Dataset.streaming_split`.
+        """
+        # To avoid deadlock, the concurrency on this actor must be set to at least `n`.
+        coord_actor = SplitCoordinator.options(
+            max_concurrency=n,
+            scheduling_strategy=NodeAffinitySchedulingStrategy(
+                ray.get_runtime_context().get_node_id(), soft=False
+            ),
+        ).remote(_DatasetWrapper(base_dataset), n, locality_hints)
+
+        return [
+            StreamSplitDataIterator(base_dataset, coord_actor, i, n) for i in range(n)
+        ]
+
+    def __init__(
+        self,
+        base_dataset: "Dataset",
+        coord_actor: ray.actor.ActorHandle,
+        output_split_idx: int,
+        world_size: int,
+    ):
+        self._base_dataset = base_dataset
+        self._coord_actor = coord_actor
+        self._output_split_idx = output_split_idx
+        self._world_size = world_size
+        self._iter_stats = DatasetStats(metadata={}, parent=None)
+
+    def _to_ref_bundle_iterator(
+        self,
+    ) -> Tuple[Iterator[RefBundle], Optional[DatasetStats], bool]:
+        def gen_blocks() -> Iterator[RefBundle]:
+            cur_epoch = ray.get(
+                self._coord_actor.start_epoch.remote(self._output_split_idx)
+            )
+            future: ObjectRef[
+                Optional[ObjectRef[Block]]
+            ] = self._coord_actor.get.remote(cur_epoch, self._output_split_idx)
+            while True:
+                block_ref_and_md: Optional[RefBundle] = ray.get(future)
+                if not block_ref_and_md:
+                    break
+                else:
+                    future = self._coord_actor.get.remote(
+                        cur_epoch, self._output_split_idx
+                    )
+                    yield RefBundle(
+                        blocks=block_ref_and_md.blocks,
+                        owns_blocks=False,
+                        schema=block_ref_and_md.schema,
+                    )
+
+        return gen_blocks(), self._iter_stats, False
+
+    def stats(self) -> str:
+        """Implements DataIterator."""
+        # Merge the locally recorded iter stats and the remotely recorded
+        # stream execution stats.
+        stats = ray.get(self._coord_actor.stats.remote())
+        summary = stats.to_summary()
+        summary.iter_stats = self._iter_stats.to_summary().iter_stats
+        summary.iter_stats.streaming_split_coord_time.add(
+            stats.streaming_split_coordinator_s.get()
+        )
+        return summary.to_string()
+
+    def schema(self) -> Union[type, "pyarrow.lib.Schema"]:
+        """Implements DataIterator."""
+        return self._base_dataset.schema()
+
+    def get_context(self) -> DataContext:
+        return self._base_dataset.context
+
+    def world_size(self) -> int:
+        """Returns the number of splits total."""
+        return self._world_size
+
+    def _get_dataset_tag(self):
+        return f"{self._base_dataset.get_dataset_id()}_split_{self._output_split_idx}"
+
+
+@ray.remote(num_cpus=0)
+class SplitCoordinator:
+    """Coordinator actor for routing blocks to output splits.
+
+    This actor runs a streaming executor locally on its main thread. Clients can
+    retrieve results via actor calls running on other threads.
+    """
+
+    def __init__(
+        self,
+        dataset_wrapper: _DatasetWrapper,
+        n: int,
+        locality_hints: Optional[List[NodeIdStr]],
+    ):
+        dataset = dataset_wrapper._dataset
+
+        # Set current DataContext.
+        # This needs to be a deep copy so that updates to the base dataset's
+        # context does not affect this process's global DataContext.
+        self._data_context = dataset.context.copy()
+        ray.data.DataContext._set_current(self._data_context)
+
+        if self._data_context.execution_options.locality_with_output is True:
+            self._data_context.execution_options.locality_with_output = locality_hints
+            logger.info(f"Auto configuring locality_with_output={locality_hints}")
+        self._base_dataset = dataset
+        self._n = n
+        self._locality_hints = locality_hints
+        self._lock = threading.RLock()
+        self._executor = None
+
+        # Guarded by self._lock.
+        self._next_bundle: Dict[int, RefBundle] = {}
+        self._unfinished_clients_in_epoch = n
+        self._cur_epoch = -1
+
+        # Add a new stats field to track coordinator overhead
+        self._coordinator_overhead_s = 0.0
+
+        def gen_epochs():
+            while True:
+                self._executor = self._base_dataset._plan.create_executor()
+                output_iterator = execute_to_legacy_bundle_iterator(
+                    self._executor, dataset._plan
+                )
+                yield output_iterator
+
+        self._next_epoch = gen_epochs()
+        self._output_iterator = None
+        # Store the error raised from the `gen_epoch` call.
+        self._gen_epoch_error: Optional[Exception] = None
+
+    def stats(self) -> DatasetStats:
+        """Returns stats from the base dataset."""
+        if self._executor:
+            stats = self._executor.get_stats()
+        else:
+            stats = self._base_dataset._plan.stats()
+
+        # Set the tracked overhead time
+        stats.streaming_split_coordinator_s.add(self._coordinator_overhead_s)
+
+        return stats
+
+    def start_epoch(self, split_idx: int) -> str:
+        """Called to start an epoch.
+
+        Returns:
+            UUID for the epoch, which must be used when accessing results via get().
+        """
+
+        # Wait for all clients to arrive at the barrier before starting a new epoch.
+        epoch_id = self._barrier(split_idx)
+        return epoch_id
+
+    def get(self, epoch_id: int, output_split_idx: int) -> Optional[RefBundle]:
+        """Blocking get operation.
+
+        This is intended to be called concurrently from multiple clients.
+        """
+        start_time = time.perf_counter()
+        if epoch_id != self._cur_epoch:
+            raise ValueError(
+                "Invalid iterator: the dataset has moved on to another epoch."
+            )
+
+        try:
+            # Ensure there is at least one bundle.
+            with self._lock:
+                if output_split_idx in self._next_bundle:
+                    next_bundle = self._next_bundle[output_split_idx]
+                else:
+                    next_bundle = None
+
+            # Fetch next bundle if needed.
+            while next_bundle is None or not next_bundle.blocks:
+                # This is a BLOCKING call, so do it outside the lock.
+                next_bundle = self._output_iterator.get_next(output_split_idx)
+
+            schema = next_bundle.schema
+            block = next_bundle.blocks[-1]
+            next_bundle = RefBundle(
+                blocks=next_bundle.blocks[:-1],
+                schema=next_bundle.schema,
+                owns_blocks=next_bundle.owns_blocks,
+                output_split_idx=next_bundle.output_split_idx,
+            )
+
+            # Accumulate any remaining blocks in next_bundle map as needed.
+            with self._lock:
+                self._next_bundle[output_split_idx] = next_bundle
+                if not next_bundle.blocks:
+                    del self._next_bundle[output_split_idx]
+
+            return RefBundle(
+                [block], schema=schema, owns_blocks=next_bundle.owns_blocks
+            )
+        except StopIteration:
+            return None
+        finally:
+            # Track overhead time in the instance variable
+            self._coordinator_overhead_s += time.perf_counter() - start_time
+
+    def _barrier(self, split_idx: int) -> int:
+        """Arrive and block until the start of the given epoch."""
+
+        # Decrement and await all clients to arrive here.
+        with self._lock:
+            starting_epoch = self._cur_epoch
+            self._unfinished_clients_in_epoch -= 1
+
+        start_time = time.time()
+        while (
+            self._cur_epoch == starting_epoch and self._unfinished_clients_in_epoch != 0
+        ):
+            if time.time() - start_time > BLOCKED_CLIENT_WARN_TIMEOUT:
+                if log_once(f"stream_split_blocked_{split_idx}_{starting_epoch}"):
+                    logger.warning(
+                        f"StreamSplitDataIterator(epoch={starting_epoch}, "
+                        f"split={split_idx}) blocked waiting on other clients "
+                        f"for more than {BLOCKED_CLIENT_WARN_TIMEOUT}s. All "
+                        "clients must read from the DataIterator splits at "
+                        "the same time. This warning will not be printed again "
+                        "for this epoch."
+                    )
+            time.sleep(0.1)
+
+        # Advance to the next epoch.
+        with self._lock:
+            if self._cur_epoch == starting_epoch:
+                self._cur_epoch += 1
+                self._unfinished_clients_in_epoch = self._n
+                try:
+                    self._output_iterator = next(self._next_epoch)
+                except Exception as e:
+                    self._gen_epoch_error = e
+
+        if self._gen_epoch_error is not None:
+            # If there was an error when advancing to the next epoch,
+            # re-raise it for all threads.
+            raise self._gen_epoch_error
+
+        assert self._output_iterator is not None
+        return starting_epoch + 1
diff --git a/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..050b39b4cab9773eccad0b92e21b0689c6c6e871
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/optimizers.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/optimizers.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..16d136780b11e9af709153476826fa7bd9c3bb6a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/optimizers.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/ruleset.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/ruleset.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..47b30979139fb588fa401a7874a306cfe4f9d36b
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/ruleset.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/util.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4e96c997026659642f5b76b0b09edee75e92291c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/logical/__pycache__/util.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/logical/operators/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/logical/operators/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..09db0c689f3861bb2064c2fd9e68877c94a916a3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/logical/operators/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/planner/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f25267b9cf82bf4bc50ad351cefefba098c13022
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/__init__.py
@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ray.data._internal.planner.planner import Planner
+
+
+def create_planner() -> "Planner":
+    # Import here to avoid circular import.
+    from ray.data._internal.planner.planner import Planner
+
+    return Planner()
+
+
+__all__ = ["create_planner"]
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5fe112e2c955bbb0bcf99a78dbeb9510b5173b97
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/aggregate.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/aggregate.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3a6d1cd0ee7ce56efb3eb1361e77291be90f82ba
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/aggregate.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_all_to_all_op.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_all_to_all_op.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f79fb530faa628c07dd2d6dccb33e5f474c8af5e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_all_to_all_op.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_download_op.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_download_op.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ceb3d3cd88ce24ce3696e1909f9cb7f4ff295d83
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_download_op.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_read_op.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_read_op.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a7e6d29aae85b1aec8a9da054f7885ed8ca89c76
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_read_op.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_udf_map_op.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_udf_map_op.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6b98242ce33501843bf246114b02e4f8fe425826
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_udf_map_op.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_write_op.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_write_op.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4487104e03659467a0e97a1566b80d81a6fd2e49
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/plan_write_op.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/planner.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/planner.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b7e541e138b03676de9ebe3e817bec3d6f237ef8
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/planner.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/random_shuffle.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/random_shuffle.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9a2e3b4731f3061323eddc17f3d9aa08a5931f2d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/random_shuffle.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/randomize_blocks.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/randomize_blocks.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dc1c77ac78e33884029a20c2c4f9eea959b9eb75
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/randomize_blocks.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/repartition.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/repartition.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..25f08bff279847d1379f8db52590475b3b17c894
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/repartition.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/sort.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/sort.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..385c9ec93c02b338d0776649c179215e65bf927d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/__pycache__/sort.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/aggregate.py b/lib/python3.12/site-packages/ray/data/_internal/planner/aggregate.py
new file mode 100644
index 0000000000000000000000000000000000000000..199382e226f09f37380c970e073de44a0fb7c57f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/aggregate.py
@@ -0,0 +1,103 @@
+from typing import List, Optional, Union
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.planner.exchange.aggregate_task_spec import (
+    SortAggregateTaskSpec,
+)
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data._internal.util import unify_ref_bundles_schema
+from ray.data.aggregate import AggregateFn
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_aggregate_fn(
+    key: Optional[Union[str, List[str]]],
+    aggs: List[AggregateFn],
+    batch_format: str,
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to aggregate blocks by the specified key column or key
+    function.
+    """
+    assert data_context.shuffle_strategy in [
+        ShuffleStrategy.SORT_SHUFFLE_PULL_BASED,
+        ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED,
+    ]
+
+    if len(aggs) == 0:
+        raise ValueError("Aggregate requires at least one aggregation")
+
+    def fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        blocks = []
+        metadata = []
+        for ref_bundle in refs:
+            blocks.extend(ref_bundle.block_refs)
+            metadata.extend(ref_bundle.metadata)
+        if len(blocks) == 0:
+            return (blocks, {})
+
+        unified_schema = unify_ref_bundles_schema(refs)
+        for agg_fn in aggs:
+            agg_fn._validate(unified_schema)
+
+        num_mappers = len(blocks)
+
+        sort_key = SortKey(key)
+
+        if key is None:
+            num_outputs = 1
+            boundaries = []
+        else:
+            # Use same number of output partitions.
+            num_outputs = num_mappers
+            sample_bar = ctx.sub_progress_bar_dict[
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME
+            ]
+            # Sample boundaries for aggregate key.
+            boundaries = SortTaskSpec.sample_boundaries(
+                blocks, sort_key, num_outputs, sample_bar
+            )
+
+        agg_spec = SortAggregateTaskSpec(
+            boundaries=boundaries,
+            key=sort_key,
+            aggs=aggs,
+            batch_format=batch_format,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(agg_spec)
+        elif data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PULL_BASED:
+            scheduler = PullBasedShuffleTaskScheduler(agg_spec)
+        else:
+            raise ValueError(
+                f"Invalid shuffle strategy '{data_context.shuffle_strategy}'"
+            )
+
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    return fn
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f3ccbb9898ce381955b97b765c6dc2d1ae3df646
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/aggregate_task_spec.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/aggregate_task_spec.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ed8f9f18bb80e62463982494e52d100acc823d49
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/aggregate_task_spec.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/interfaces.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/interfaces.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cd432110b4c03ef6db0160f82c77ce3c13484e88
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/interfaces.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/pull_based_shuffle_task_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/pull_based_shuffle_task_scheduler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..70f50c894cb864395bba9a6d366a9b0b196d74db
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/pull_based_shuffle_task_scheduler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/push_based_shuffle_task_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/push_based_shuffle_task_scheduler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0347d95a57452132869ef7f47a6a5e2c2d736fa4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/push_based_shuffle_task_scheduler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/shuffle_task_spec.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/shuffle_task_spec.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1244de0860cf30232817c99e97cc382c781a637b
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/shuffle_task_spec.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/sort_task_spec.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/sort_task_spec.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5c4bd85e7c51d950aad3cff482bc40ccb458e5f0
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/sort_task_spec.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/split_repartition_task_scheduler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/split_repartition_task_scheduler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1962ab0a7cc6ac55e467c77561702d53c9f3ac80
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/__pycache__/split_repartition_task_scheduler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/aggregate_task_spec.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/aggregate_task_spec.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6d8f80dca8083c17eaee2cd80a4a74f035b555f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/aggregate_task_spec.py
@@ -0,0 +1,125 @@
+from typing import List, Tuple, Union
+
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+from ray.data._internal.table_block import TableBlockAccessor
+from ray.data.aggregate import AggregateFn, AggregateFnV2, Count
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadataWithSchema,
+    KeyType,
+)
+
+
+class SortAggregateTaskSpec(ExchangeTaskSpec):
+    """
+    The implementation for sort-based aggregate tasks.
+
+    Aggregate is done in 2 steps: partial aggregate of individual blocks, and
+    final aggregate of sorted blocks.
+
+    Partial aggregate (`map`): each block is sorted locally, then partitioned into
+    smaller blocks according to the boundaries. Each partitioned block is aggregated
+    separately, then passed to a final aggregate task.
+
+    Final aggregate (`reduce`): each task would receive a block from every worker that
+    consists of items in a certain range. It then merges the sorted blocks and
+    aggregates on-the-fly.
+    """
+
+    def __init__(
+        self,
+        boundaries: List[KeyType],
+        key: SortKey,
+        aggs: List[AggregateFn],
+        batch_format: str,
+    ):
+        super().__init__(
+            map_args=[boundaries, key, aggs],
+            reduce_args=[key, aggs, batch_format],
+        )
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        boundaries: List[KeyType],
+        sort_key: SortKey,
+        aggs: List[AggregateFn],
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        stats = BlockExecStats.builder()
+
+        block = SortAggregateTaskSpec._prune_unused_columns(block, sort_key, aggs)
+        if sort_key.get_columns():
+            partitions = BlockAccessor.for_block(block).sort_and_partition(
+                boundaries,
+                sort_key,
+            )
+        else:
+            partitions = [block]
+        parts = [
+            BlockAccessor.for_block(p)._aggregate(sort_key, aggs) for p in partitions
+        ]
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta_with_schema = BlockMetadataWithSchema.from_block(
+            block, stats=stats.build()
+        )
+        return parts + [meta_with_schema]
+
+    @staticmethod
+    def reduce(
+        key: SortKey,
+        aggs: List[AggregateFn],
+        batch_format: str,
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        normalized_blocks = TableBlockAccessor.normalize_block_types(
+            mapper_outputs,
+            target_block_type=ExchangeTaskSpec._derive_target_block_type(batch_format),
+        )
+        blocks, meta_with_schema = BlockAccessor.for_block(
+            normalized_blocks[0]
+        )._combine_aggregated_blocks(
+            list(normalized_blocks), key, aggs, finalize=not partial_reduce
+        )
+        return blocks, meta_with_schema
+
+    @staticmethod
+    def _prune_unused_columns(
+        block: Block,
+        sort_key: SortKey,
+        aggs: Tuple[AggregateFn],
+    ) -> Block:
+        """Prune unused columns from block before aggregate."""
+        prune_columns = True
+        columns = set()
+        key = sort_key.get_columns()
+
+        if isinstance(key, str):
+            columns.add(key)
+        elif isinstance(key, list):
+            columns.update(key)
+        elif callable(key):
+            prune_columns = False
+
+        for agg in aggs:
+            if isinstance(agg, AggregateFnV2) and agg.get_target_column():
+                columns.add(agg.get_target_column())
+            elif not isinstance(agg, Count):
+                # Don't prune columns if any aggregate key is not string.
+                prune_columns = False
+
+        block_accessor = BlockAccessor.for_block(block)
+        if (
+            prune_columns
+            and isinstance(block_accessor, TableBlockAccessor)
+            and block_accessor.num_rows() > 0
+        ):
+            return block_accessor.select(list(columns))
+        else:
+            return block
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/interfaces.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/interfaces.py
new file mode 100644
index 0000000000000000000000000000000000000000..84388543509fc1b5236992114933dfb7ebf1d9ba
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/interfaces.py
@@ -0,0 +1,146 @@
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import ray._private.worker
+from ray.air.util.data_batch_conversion import BatchFormat
+from ray.data._internal.execution.interfaces import RefBundle
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import convert_bytes_to_human_readable_str
+from ray.data.block import Block, BlockType
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+
+    from ray.data.block import BlockMetadataWithSchema
+
+logger = logging.getLogger(__name__)
+
+
+class ExchangeTaskSpec:
+    """
+    An interface to specify the exchange map and reduce tasks.
+
+    Subclasses should implement the `map` and `reduce` static methods.
+    `map` method is to transform one input block into multiple output blocks.
+    `reduce` is to combine multiple map output blocks. Both methods are
+    single-task operations. See `ExchangeScheduler` for how to distribute
+    the operations across multiple tasks.
+
+    Any custom arguments for `map` and `reduce` methods should be specified by
+    setting `map_args` and `reduce_args`.
+
+    The concept here is similar to the exchange operator described in
+    "Volcano - An Extensible and Parallel Query Evaluation System"
+    (https://dl.acm.org/doi/10.1109/69.273032).
+    """
+
+    MAP_SUB_PROGRESS_BAR_NAME = "Shuffle Map"
+    REDUCE_SUB_PROGRESS_BAR_NAME = "Shuffle Reduce"
+
+    def __init__(self, map_args: List[Any] = None, reduce_args: List[Any] = None):
+        self._map_args = map_args or []
+        self._reduce_args = reduce_args or []
+        assert isinstance(self._map_args, list)
+        assert isinstance(self._reduce_args, list)
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        """
+        Map function to be run on each input block.
+
+        Returns list of [BlockMetadata, Block1, Block2, ..., BlockN].
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def reduce(
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        """
+        Reduce function to be run for each output block.
+
+        Args:
+            mapper_outputs: List of map output blocks to reduce.
+            partial_reduce: Whether should partially or fully reduce.
+
+        Returns:
+            The reduced block and its metadata.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def _derive_target_block_type(batch_format: str) -> Optional[BlockType]:
+        if batch_format == BatchFormat.ARROW:
+            return BlockType.ARROW
+        elif batch_format == BatchFormat.PANDAS:
+            return BlockType.PANDAS
+        else:
+            # NOTE: Unless desired batch-format is specified, avoid
+            #       overriding existing one
+            return None
+
+
+class ExchangeTaskScheduler:
+    """
+    An interface to schedule exchange tasks (`exchange_spec`) for multi-nodes
+    execution.
+    """
+
+    def __init__(self, exchange_spec: ExchangeTaskSpec):
+        """
+        Args:
+            exchange_spec: The implementation of exchange tasks to execute.
+        """
+        self._exchange_spec = exchange_spec
+        # If driver memory exceeds this threshold, warn the user. For now, this
+        # only applies to shuffle ops because most other ops are unlikely to use as
+        # much driver memory.
+        self.warn_on_driver_memory_usage_bytes: Optional[
+            int
+        ] = DataContext.get_current().warn_on_driver_memory_usage_bytes
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        warn_on_driver_memory_usage: Optional[int] = None,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        """
+        Execute the exchange tasks on input `refs`.
+        """
+        raise NotImplementedError
+
+    def warn_on_high_local_memory_store_usage(self) -> None:
+        ray_core_worker = ray._private.worker.global_worker.core_worker
+        local_memory_store_bytes_used = (
+            ray_core_worker.get_local_memory_store_bytes_used()
+        )
+        self.warn_on_driver_memory_usage(
+            local_memory_store_bytes_used,
+            "More than "
+            f"{convert_bytes_to_human_readable_str(local_memory_store_bytes_used)} "
+            "of driver memory used to store Ray Data block data and metadata. "
+            "This job may exit if driver memory is insufficient.\n\n"
+            "This can happen when many tiny blocks are created. "
+            "Check the block size using Dataset.stats() and see "
+            "https://docs.ray.io/en/latest/data/performance-tips.html"
+            " for mitigation.",
+        )
+
+    def warn_on_driver_memory_usage(
+        self, memory_usage_bytes: int, log_str: str
+    ) -> None:
+        if self.warn_on_driver_memory_usage_bytes is None:
+            return
+
+        if memory_usage_bytes > self.warn_on_driver_memory_usage_bytes:
+            logger.warning(log_str)
+            # Double the threshold to avoid verbose warnings.
+            self.warn_on_driver_memory_usage_bytes = memory_usage_bytes * 2
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/pull_based_shuffle_task_scheduler.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/pull_based_shuffle_task_scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..262382b2a08ed728bb639b8416e016f21753a0bc
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/pull_based_shuffle_task_scheduler.py
@@ -0,0 +1,160 @@
+import logging
+from typing import Any, Dict, List, Optional, Tuple
+
+from ray._private.ray_constants import CALLER_MEMORY_USAGE_PER_OBJECT_REF
+from ray.data._internal.execution.interfaces import RefBundle, TaskContext
+from ray.data._internal.planner.exchange.interfaces import (
+    ExchangeTaskScheduler,
+    ExchangeTaskSpec,
+)
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import (
+    convert_bytes_to_human_readable_str,
+    unzip,
+)
+from ray.data.block import BlockMetadataWithSchema, to_stats
+
+logger = logging.getLogger(__name__)
+
+
+class PullBasedShuffleTaskScheduler(ExchangeTaskScheduler):
+    """
+    The pull-based map-reduce shuffle scheduler.
+
+    Map tasks are first scheduled to generate map output blocks. After all map output
+    are generated, then reduce tasks are scheduled to combine map output blocks
+    together.
+
+    The concept here is similar to
+    "MapReduce: Simplified Data Processing on Large Clusters"
+    (https://dl.acm.org/doi/10.1145/1327452.1327492).
+    """
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        task_ctx: TaskContext,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        _debug_limit_execution_to_num_blocks: Optional[int] = None,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+
+        # TODO: eagerly delete the input and map output block references in order to
+        # eagerly release the blocks' memory.
+        input_blocks_list = []
+        for ref_bundle in refs:
+            input_blocks_list.extend(ref_bundle.block_refs)
+        input_num_blocks = len(input_blocks_list)
+        input_owned = all(b.owns_blocks for b in refs)
+
+        caller_memory_usage = (
+            input_num_blocks * output_num_blocks * CALLER_MEMORY_USAGE_PER_OBJECT_REF
+        )
+        self.warn_on_driver_memory_usage(
+            caller_memory_usage,
+            "Execution is estimated to use at least "
+            f"{convert_bytes_to_human_readable_str(caller_memory_usage)} "
+            "of driver memory. Ensure that the driver machine has at least "
+            "this much memory to ensure job completion.\n\n"
+            "To reduce the "
+            "amount of driver memory needed, enable push-based shuffle using "
+            "RAY_DATA_PUSH_BASED_SHUFFLE=1 "
+            "(https://docs.ray.io/en/latest/data/performance-tips.html"
+            ").",
+        )
+
+        if map_ray_remote_args is None:
+            map_ray_remote_args = {}
+        if reduce_ray_remote_args is None:
+            reduce_ray_remote_args = {}
+        if "scheduling_strategy" not in reduce_ray_remote_args:
+            reduce_ray_remote_args = reduce_ray_remote_args.copy()
+            reduce_ray_remote_args["scheduling_strategy"] = "SPREAD"
+
+        shuffle_map = cached_remote_fn(self._exchange_spec.map)
+        shuffle_reduce = cached_remote_fn(self._exchange_spec.reduce)
+
+        sub_progress_bar_dict = task_ctx.sub_progress_bar_dict
+        bar_name = ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        map_bar = sub_progress_bar_dict[bar_name]
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            input_blocks_list = input_blocks_list[:_debug_limit_execution_to_num_blocks]
+            logger.debug(f"Limiting execution to {len(input_blocks_list)} map tasks")
+
+        shuffle_map_out = [
+            shuffle_map.options(
+                **map_ray_remote_args,
+                num_returns=1 + output_num_blocks,
+            ).remote(i, block, output_num_blocks, *self._exchange_spec._map_args)
+            for i, block in enumerate(input_blocks_list)
+        ]
+
+        # The first item returned is the BlockMetadata.
+        shuffle_map_metadata_schema = []
+        for i, refs in enumerate(shuffle_map_out):
+            shuffle_map_metadata_schema.append(refs[-1])
+            shuffle_map_out[i] = refs[:-1]
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            while len(shuffle_map_out) < output_num_blocks:
+                # Repeat the first map task's results.
+                shuffle_map_out.append(shuffle_map_out[0][:])
+
+        shuffle_map_metadata_schema = map_bar.fetch_until_complete(
+            shuffle_map_metadata_schema
+        )
+
+        self.warn_on_high_local_memory_store_usage()
+
+        bar_name = ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        reduce_bar = sub_progress_bar_dict[bar_name]
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            output_num_blocks = _debug_limit_execution_to_num_blocks
+            logger.debug(f"Limiting execution to {output_num_blocks} reduce tasks")
+        shuffle_reduce_out = [
+            shuffle_reduce.options(**reduce_ray_remote_args, num_returns=2).remote(
+                *self._exchange_spec._reduce_args,
+                *[shuffle_map_out[i][j] for i in range(input_num_blocks)],
+            )
+            for j in range(output_num_blocks)
+        ]
+
+        # Release map task outputs from the Ray object store.
+        del shuffle_map_out
+
+        new_blocks, new_metadata_schema = [], []
+        if shuffle_reduce_out:
+            new_blocks, new_metadata_schema = unzip(shuffle_reduce_out)
+        new_metadata_schema: List[
+            "BlockMetadataWithSchema"
+        ] = reduce_bar.fetch_until_complete(list(new_metadata_schema))
+
+        self.warn_on_high_local_memory_store_usage()
+
+        output = []
+        for block, meta_with_schema in zip(new_blocks, new_metadata_schema):
+            output.append(
+                RefBundle(
+                    [
+                        (
+                            block,
+                            meta_with_schema.metadata,
+                        )
+                    ],
+                    owns_blocks=input_owned,
+                    schema=meta_with_schema.schema,
+                )
+            )
+
+        stats = {
+            "map": to_stats(shuffle_map_metadata_schema),
+            "reduce": to_stats(new_metadata_schema),
+        }
+
+        return (output, stats)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/push_based_shuffle_task_scheduler.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/push_based_shuffle_task_scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..901eb69bd969f9317a8dd8323b14d1d6cd171d65
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/push_based_shuffle_task_scheduler.py
@@ -0,0 +1,853 @@
+import logging
+import math
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, TypeVar, Union
+
+import ray
+from ray._private.ray_constants import CALLER_MEMORY_USAGE_PER_OBJECT_REF
+from ray.data._internal.execution.interfaces import RefBundle, TaskContext
+from ray.data._internal.planner.exchange.interfaces import (
+    ExchangeTaskScheduler,
+    ExchangeTaskSpec,
+)
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import (
+    convert_bytes_to_human_readable_str,
+    unzip,
+)
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+    _take_first_non_empty_schema,
+    to_stats,
+)
+from ray.data.context import DataContext
+from ray.types import ObjectRef
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+if TYPE_CHECKING:
+
+    from ray.data.block import BlockMetadataWithSchema
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+class _MergeTaskSchedule:
+    def __init__(self, output_num_blocks: int, num_merge_tasks_per_round: int):
+        self.output_num_blocks = output_num_blocks
+        self.num_merge_tasks_per_round = num_merge_tasks_per_round
+        self.num_reducers_per_merger = output_num_blocks // num_merge_tasks_per_round
+        self._num_mergers_with_extra_reducer = (
+            output_num_blocks % num_merge_tasks_per_round
+        )
+
+        if self.num_reducers_per_merger == 0:
+            self.num_merge_tasks_per_round = self._num_mergers_with_extra_reducer
+            self.num_reducers_per_merger = 1
+            self._num_mergers_with_extra_reducer = 0
+
+    def __repr__(self):
+        return (
+            f"    num merge tasks per round: {self.num_merge_tasks_per_round}\n"
+            f"    num reduce tasks per merge task: {self.num_reducers_per_merger}\n"
+            "    num merge tasks with extra reduce task: "
+            f"{self._num_mergers_with_extra_reducer}"
+        )
+
+    def get_num_reducers_per_merge_idx(self, merge_idx: int) -> int:
+        """
+        Each intermediate merge task will produce outputs for a partition of P
+        final reduce tasks. This helper function returns P based on the merge
+        task index.
+        """
+        assert merge_idx < self.num_merge_tasks_per_round
+        num_reducers_for_cur_merger = self.num_reducers_per_merger
+        if merge_idx < self._num_mergers_with_extra_reducer:
+            num_reducers_for_cur_merger += 1
+        return num_reducers_for_cur_merger
+
+    def get_merge_idx_for_reducer_idx(self, reducer_idx: int) -> int:
+        if (
+            reducer_idx
+            < (self.num_reducers_per_merger + 1) * self._num_mergers_with_extra_reducer
+        ):
+            merge_idx = reducer_idx // (self.num_reducers_per_merger + 1)
+        else:
+            reducer_idx -= (
+                self.num_reducers_per_merger + 1
+            ) * self._num_mergers_with_extra_reducer
+            merge_idx = (
+                self._num_mergers_with_extra_reducer
+                + reducer_idx // self.num_reducers_per_merger
+            )
+        assert merge_idx < self.num_merge_tasks_per_round
+        return merge_idx
+
+    def round_robin_reduce_idx_iterator(self):
+        """
+        When there are multiple nodes, merge tasks are spread throughout the
+        cluster to improve load-balancing. Each merge task produces outputs for
+        a contiguous partition of reduce tasks. This method creates an iterator
+        that returns reduce task indices round-robin across the merge tasks.
+        This can be used to submit reduce tasks in a way that spreads the load
+        evenly across the cluster.
+        """
+        idx = 0
+        round_idx = 0
+        while idx < self.output_num_blocks:
+            for merge_idx in range(self.num_merge_tasks_per_round):
+                if merge_idx < self._num_mergers_with_extra_reducer:
+                    reduce_idx = merge_idx * (self.num_reducers_per_merger + 1)
+                    num_reducers_for_cur_merger = self.num_reducers_per_merger + 1
+                else:
+                    reduce_idx = self._num_mergers_with_extra_reducer * (
+                        self.num_reducers_per_merger + 1
+                    )
+                    merge_idx -= self._num_mergers_with_extra_reducer
+                    reduce_idx += merge_idx * self.num_reducers_per_merger
+                    num_reducers_for_cur_merger = self.num_reducers_per_merger
+
+                if round_idx >= num_reducers_for_cur_merger:
+                    continue
+
+                reduce_idx += round_idx
+                yield reduce_idx
+                idx += 1
+            round_idx += 1
+
+
+class _PushBasedShuffleStage:
+    def __init__(
+        self,
+        output_num_blocks: int,
+        num_rounds: int,
+        num_map_tasks_per_round: int,
+        merge_task_placement: List[str],
+    ):
+        # The number of rounds of map-merge tasks. Reducer tasks are given the
+        # outputs of the merge tasks as inputs. Reducer tasks receive one input
+        # per round.
+        self.num_rounds = num_rounds
+        # The number of map tasks per round of map-merge tasks. The map task
+        # produces one output per merge task in the same round.
+        self.num_map_tasks_per_round = num_map_tasks_per_round
+
+        node_strategies = {
+            node_id: {
+                "scheduling_strategy": NodeAffinitySchedulingStrategy(
+                    node_id, soft=True
+                )
+            }
+            for node_id in set(merge_task_placement)
+        }
+        self._merge_task_options = [
+            node_strategies[node_id] for node_id in merge_task_placement
+        ]
+
+        self.merge_schedule = _MergeTaskSchedule(
+            output_num_blocks, len(merge_task_placement)
+        )
+
+    def get_estimated_num_refs(self) -> int:
+        # Number of intermediate blocks = Number of rounds x (map tasks per
+        # round * merge tasks per round).
+        num_intermediate_refs = self.num_rounds * (
+            self.num_map_tasks_per_round * self.merge_schedule.num_merge_tasks_per_round
+        )
+        # Number of input blocks + intermediate blocks + output blocks.
+        num_refs_total = (
+            (self.num_rounds * self.num_map_tasks_per_round)
+            + num_intermediate_refs
+            + self.merge_schedule.output_num_blocks
+        )
+        return num_refs_total
+
+    def get_merge_task_options(self, merge_idx):
+        return self._merge_task_options[merge_idx]
+
+    def __repr__(self):
+        return (
+            "num map tasks per round (num args per merge task): "
+            f"{self.num_map_tasks_per_round}\n"
+            f"num rounds (num args per reduce task): {self.num_rounds}\n"
+            "merge task placement: \n"
+            f"{self.merge_schedule}"
+        )
+
+
+class _PipelinedStageExecutor:
+    def __init__(
+        self,
+        stage_iter,
+        num_tasks_per_round: int,
+        max_concurrent_rounds: int = 1,
+        progress_bar: Optional[ProgressBar] = None,
+    ):
+        self._stage_iter = stage_iter
+        self._num_tasks_per_round = num_tasks_per_round
+        self._max_concurrent_rounds = max_concurrent_rounds
+        self._progress_bar = progress_bar
+
+        self._rounds: List[List[ObjectRef]] = []
+        self._task_idx = 0
+
+        self._submit_round()
+
+        self._num_block_bytes_stored_at_driver = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> List["BlockMetadataWithSchema"]:
+        """
+        Submit one round of tasks. If we already have the max concurrent rounds
+        in flight, first wait for the oldest round of tasks to finish.
+        """
+        prev_metadata_and_schema = []
+        if all(len(r) == 0 for r in self._rounds):
+            raise StopIteration
+
+        if len(self._rounds) >= self._max_concurrent_rounds:
+            prev_metadata_schema_refs = self._rounds.pop(0)
+            if prev_metadata_schema_refs:
+                if self._progress_bar is not None:
+                    prev_metadata_and_schema = self._progress_bar.fetch_until_complete(
+                        prev_metadata_schema_refs
+                    )
+                    # TODO(swang): Eagerly free the previous round's args.
+                    # See https://github.com/ray-project/ray/issues/42145.
+                else:
+                    prev_metadata_and_schema = ray.get(prev_metadata_schema_refs)
+
+        self._submit_round()
+
+        return prev_metadata_and_schema
+
+    def _submit_round(self):
+        assert len(self._rounds) < self._max_concurrent_rounds
+        task_round = []
+        for _ in range(self._num_tasks_per_round):
+            try:
+                task_round.append(next(self._stage_iter))
+            except StopIteration:
+                break
+        self._rounds.append(task_round)
+
+
+class _MapStageIterator:
+    def __init__(self, input_blocks_list, shuffle_map, map_args):
+        self._input_blocks_list = input_blocks_list
+        self._shuffle_map = shuffle_map
+        self._map_args = map_args
+
+        self._mapper_idx = 0
+        self._map_results = []
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self._input_blocks_list:
+            raise StopIteration
+
+        block = self._input_blocks_list.pop(0)
+        # NOTE(swang): Results are shuffled between map and merge tasks, so
+        # there is no advantage to colocating specific map and merge tasks.
+        # Therefore, we do not specify a node affinity policy for map tasks
+        # in case the caller or Ray has a better scheduling strategy, e.g.,
+        # based on data locality.
+        map_result = self._shuffle_map.remote(
+            self._mapper_idx,
+            block,
+            *self._map_args,
+        )
+        metadata_schema_ref = map_result.pop(-1)
+        self._map_results.append(map_result)
+        self._mapper_idx += 1
+        return metadata_schema_ref
+
+    def pop_map_results(self) -> List[List[ObjectRef]]:
+        map_results = self._map_results
+        self._map_results = []
+        return map_results
+
+
+class _MergeStageIterator:
+    def __init__(
+        self,
+        map_stage_iter: _MapStageIterator,
+        shuffle_merge,
+        stage: _PushBasedShuffleStage,
+        reduce_args,
+    ):
+        self._map_stage_iter = map_stage_iter
+        self._shuffle_merge = shuffle_merge
+        self._stage = stage
+        self._reduce_args = reduce_args
+
+        self._merge_idx = 0
+        self._map_result_buffer = None
+        # Final outputs from the map-merge stage.
+        # This is a map from merge task index to a nested list of merge results
+        # (ObjectRefs). Each merge task index corresponds to a partition of P
+        # final reduce tasks.
+        self._all_merge_results = [
+            [] for _ in range(self._stage.merge_schedule.num_merge_tasks_per_round)
+        ]
+
+    def __next__(self):
+        if not self._map_result_buffer or not self._map_result_buffer[0]:
+            assert self._merge_idx == 0
+            self._map_result_buffer = self._map_stage_iter.pop_map_results()
+
+        if not self._map_result_buffer:
+            raise StopIteration
+
+        # Shuffle the map results for the merge tasks.
+        merge_args = [map_result.pop(0) for map_result in self._map_result_buffer]
+        num_merge_returns = self._stage.merge_schedule.get_num_reducers_per_merge_idx(
+            self._merge_idx
+        )
+        merge_result = self._shuffle_merge.options(
+            num_returns=1 + num_merge_returns,
+            **self._stage.get_merge_task_options(self._merge_idx),
+        ).remote(
+            *merge_args,
+            reduce_args=self._reduce_args,
+        )
+        metadata_schema_ref = merge_result.pop(-1)
+        self._all_merge_results[self._merge_idx].append(merge_result)
+        del merge_result
+
+        self._merge_idx += 1
+        self._merge_idx %= self._stage.merge_schedule.num_merge_tasks_per_round
+        return metadata_schema_ref
+
+    def pop_merge_results(self) -> List[List[ObjectRef]]:
+        """Return a nested list of merge task results. The list at index i
+        stores the outputs of the i-th merge task submitted during each
+        map-merge round. Each merge task returns a list of outputs because it
+        may produce outputs for multiple downstream reduce tasks.
+        """
+        all_merge_results = self._all_merge_results
+        self._all_merge_results = []
+        return all_merge_results
+
+
+class _ReduceStageIterator:
+    def __init__(
+        self,
+        stage: _PushBasedShuffleStage,
+        shuffle_reduce,
+        all_merge_results: List[List[List[ObjectRef]]],
+        ray_remote_args,
+        reduce_args: List[Any],
+        _debug_limit_execution_to_num_blocks: Optional[int],
+    ):
+        self._shuffle_reduce = shuffle_reduce
+        self._stage = stage
+        self._reduce_arg_blocks: List[Tuple[int, List[ObjectRef]]] = []
+        self._ray_remote_args = ray_remote_args
+        self._reduce_args = reduce_args
+
+        for reduce_idx in self._stage.merge_schedule.round_robin_reduce_idx_iterator():
+            merge_idx = self._stage.merge_schedule.get_merge_idx_for_reducer_idx(
+                reduce_idx
+            )
+            reduce_arg_blocks = [
+                merge_results.pop(0) for merge_results in all_merge_results[merge_idx]
+            ]
+            self._reduce_arg_blocks.append((reduce_idx, reduce_arg_blocks))
+
+        assert len(self._reduce_arg_blocks) == stage.merge_schedule.output_num_blocks
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            self._reduce_arg_blocks = self._reduce_arg_blocks[
+                :_debug_limit_execution_to_num_blocks
+            ]
+            logger.debug(
+                f"Limiting execution to {len(self._reduce_arg_blocks)} reduce tasks"
+            )
+
+        for merge_idx, merge_results in enumerate(all_merge_results):
+            assert all(len(merge_result) == 0 for merge_result in merge_results), (
+                "Reduce stage did not process outputs from merge tasks at index: "
+                f"{merge_idx}"
+            )
+
+        self._reduce_results: List[Tuple[int, ObjectRef]] = []
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self._reduce_arg_blocks:
+            raise StopIteration
+
+        reduce_idx, reduce_arg_blocks = self._reduce_arg_blocks.pop(0)
+        merge_idx = self._stage.merge_schedule.get_merge_idx_for_reducer_idx(reduce_idx)
+        # Submit one partition of reduce tasks, one for each of the P
+        # outputs produced by the corresponding merge task.
+        # We also add the merge task arguments so that the reduce task
+        # is colocated with its inputs.
+        block, meta_with_schema = self._shuffle_reduce.options(
+            **self._ray_remote_args,
+            **self._stage.get_merge_task_options(merge_idx),
+            num_returns=2,
+        ).remote(*self._reduce_args, *reduce_arg_blocks, partial_reduce=False)
+        self._reduce_results.append((reduce_idx, block))
+        return meta_with_schema
+
+    def pop_reduce_results(self):
+        reduce_results = self._reduce_results
+        self._reduce_results = []
+        return reduce_results
+
+
+class PushBasedShuffleTaskScheduler(ExchangeTaskScheduler):
+    """
+    Push-based shuffle merges intermediate map outputs on the reducer nodes
+    while other map tasks are executing. The merged outputs are merged again
+    during a final reduce stage. This works as follows:
+
+    1. Submit rounds of concurrent map and merge tasks until all map inputs
+    have been processed. In each round, we execute:
+
+       M map tasks
+         Each produces N outputs. Each output contains P blocks.
+       N merge tasks
+         Takes 1 output from each of M map tasks.
+         Each produces P outputs.
+       Where M and N are chosen to maximize parallelism across CPUs. Note that
+       this assumes that all CPUs in the cluster will be dedicated to the
+       shuffle job.
+
+       Map and merge tasks are pipelined so that we always merge the previous
+       round of map outputs while executing the next round of map tasks.
+
+    2. In the final reduce stage:
+       R reduce tasks
+         Takes 1 output from one of the merge tasks from every round.
+
+    Notes:
+        N * P = R = total number of output blocks
+        M / N = merge factor - the ratio of map : merge tasks is to improve
+          pipelined parallelism. For example, if map takes twice as long to
+          execute as merge, then we should set this to 2. If pipeline bubbles
+          appear and the merge tasks are much longer than the map tasks, then
+          the merge factor should be decreased, and vice versa.
+        See paper at https://arxiv.org/abs/2203.05072 for more details.
+    """
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        task_ctx: TaskContext,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+        merge_factor: float = 2,
+        _debug_limit_execution_to_num_blocks: int = None,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        logger.debug("Using experimental push-based shuffle.")
+        # TODO: Preemptively clear the blocks list since we will incrementally delete
+        # the last remaining references as we submit the dependent map tasks during the
+        # map-merge stage.
+
+        # TODO(swang): For jobs whose reduce work is heavier than the map work,
+        # we should support fractional merge factors.
+        # TODO(swang): For large jobs, we should try to choose the merge factor
+        # automatically, e.g., by running one test round of map and merge tasks
+        # and comparing their run times.
+        # TODO(swang): Add option to automatically reduce write amplification
+        # during map-merge stage, by limiting how many partitions can be
+        # processed concurrently.
+        input_blocks_list = []
+        for ref_bundle in refs:
+            input_blocks_list.extend(ref_bundle.block_refs)
+        input_owned = all(b.owns_blocks for b in refs)
+
+        if map_ray_remote_args is None:
+            map_ray_remote_args = {}
+        if reduce_ray_remote_args is None:
+            reduce_ray_remote_args = {}
+        # The placement strategy for reduce tasks is overwritten to colocate
+        # them with their inputs from the merge stage, so remove any
+        # pre-specified scheduling strategy here.
+        reduce_ray_remote_args = reduce_ray_remote_args.copy()
+        reduce_ray_remote_args.pop("scheduling_strategy", None)
+
+        # Compute all constants used for task scheduling.
+        num_cpus_per_node_map = _get_num_cpus_per_node_map()
+        stage = self._compute_shuffle_schedule(
+            num_cpus_per_node_map,
+            len(input_blocks_list),
+            merge_factor,
+            output_num_blocks,
+        )
+
+        caller_memory_usage = (
+            stage.get_estimated_num_refs() * CALLER_MEMORY_USAGE_PER_OBJECT_REF
+        )
+        self.warn_on_driver_memory_usage(
+            caller_memory_usage,
+            "Execution is estimated to use at least "
+            f"{convert_bytes_to_human_readable_str(caller_memory_usage)}"
+            " of driver memory. Ensure that the driver machine has at least "
+            "this much memory to ensure job completion.",
+        )
+
+        # TODO(swang): Use INFO level. Currently there is no easy way to set
+        # the logging level to DEBUG from a driver script, so just print
+        # verbosely for now.
+        # See https://github.com/ray-project/ray/issues/42002.
+        logger.debug(f"Push-based shuffle schedule:\n{stage}")
+
+        map_fn = self._map_partition
+        merge_fn = self._merge
+
+        def map_partition(*args, **kwargs):
+            return map_fn(self._exchange_spec.map, *args, **kwargs)
+
+        def merge(*args, **kwargs):
+            return merge_fn(self._exchange_spec.reduce, *args, **kwargs)
+
+        shuffle_map = cached_remote_fn(map_partition)
+        shuffle_map = shuffle_map.options(
+            **map_ray_remote_args,
+            num_returns=1 + stage.merge_schedule.num_merge_tasks_per_round,
+        )
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            input_blocks_list = input_blocks_list[:_debug_limit_execution_to_num_blocks]
+            logger.debug(f"Limiting execution to {len(input_blocks_list)} map tasks")
+        map_stage_iter = _MapStageIterator(
+            input_blocks_list,
+            shuffle_map,
+            [output_num_blocks, stage.merge_schedule, *self._exchange_spec._map_args],
+        )
+
+        sub_progress_bar_dict = task_ctx.sub_progress_bar_dict
+        bar_name = ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        map_bar = sub_progress_bar_dict[bar_name]
+        map_stage_executor = _PipelinedStageExecutor(
+            map_stage_iter, stage.num_map_tasks_per_round, progress_bar=map_bar
+        )
+
+        shuffle_merge = cached_remote_fn(merge)
+        merge_stage_iter = _MergeStageIterator(
+            map_stage_iter, shuffle_merge, stage, self._exchange_spec._reduce_args
+        )
+        merge_stage_executor = _PipelinedStageExecutor(
+            merge_stage_iter,
+            stage.merge_schedule.num_merge_tasks_per_round,
+            max_concurrent_rounds=2,
+        )
+        # Execute the map-merge stage. This submits tasks in rounds of M map
+        # tasks and N merge tasks each. Task execution between map and merge is
+        # pipelined, so that while executing merge for one round of inputs, we
+        # also execute the map tasks for the following round.
+        map_done = False
+        merge_done = False
+        map_stage_metadata_schema = []
+        merge_stage_metadata_schema = []
+        while not (map_done and merge_done):
+            try:
+                map_stage_metadata_schema += next(map_stage_executor)
+            except StopIteration:
+                map_done = True
+                break
+
+            try:
+                merge_stage_metadata_schema += next(merge_stage_executor)
+            except StopIteration:
+                merge_done = True
+                break
+
+            self.warn_on_high_local_memory_store_usage()
+
+        all_merge_results = merge_stage_iter.pop_merge_results()
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            for merge_idx in range(len(all_merge_results)):
+                while len(all_merge_results[merge_idx]) < stage.num_rounds:
+                    # Repeat the first merge task's results.
+                    all_merge_results[merge_idx].append(
+                        all_merge_results[merge_idx][0][:]
+                    )
+
+        # Execute and wait for the reduce stage.
+        bar_name = ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        reduce_bar = sub_progress_bar_dict[bar_name]
+
+        shuffle_reduce = cached_remote_fn(self._exchange_spec.reduce)
+        reduce_stage_iter = _ReduceStageIterator(
+            stage,
+            shuffle_reduce,
+            all_merge_results,
+            reduce_ray_remote_args,
+            self._exchange_spec._reduce_args,
+            _debug_limit_execution_to_num_blocks,
+        )
+
+        max_reduce_tasks_in_flight = output_num_blocks
+        ctx = DataContext.get_current()
+        if ctx.pipeline_push_based_shuffle_reduce_tasks:
+            # If pipelining is enabled, we should still try to utilize all
+            # cores.
+            max_reduce_tasks_in_flight = min(
+                max_reduce_tasks_in_flight, sum(num_cpus_per_node_map.values())
+            )
+
+        reduce_stage_executor = _PipelinedStageExecutor(
+            reduce_stage_iter,
+            max_reduce_tasks_in_flight,
+            max_concurrent_rounds=2,
+            progress_bar=reduce_bar,
+        )
+        reduce_stage_metadata_schema = []
+        while True:
+            try:
+                reduce_stage_metadata_schema += next(reduce_stage_executor)
+            except StopIteration:
+                break
+
+            self.warn_on_high_local_memory_store_usage()
+
+        new_blocks = reduce_stage_iter.pop_reduce_results()
+        sorted_blocks = [
+            (block[0], block[1], reduce_stage_metadata_schema[i])
+            for i, block in enumerate(new_blocks)
+        ]
+        sorted_blocks.sort(key=lambda x: x[0])
+
+        new_blocks, reduce_stage_metadata_schema = [], []
+        if sorted_blocks:
+            res: Tuple[
+                List[Any], List[ObjectRef[Block]], List[BlockMetadataWithSchema]
+            ] = unzip(sorted_blocks)
+            _, new_blocks, reduce_stage_metadata_schema = res
+        del sorted_blocks
+
+        if _debug_limit_execution_to_num_blocks is not None:
+            output_num_blocks = min(
+                _debug_limit_execution_to_num_blocks, output_num_blocks
+            )
+
+        assert (
+            len(new_blocks) == output_num_blocks
+        ), f"Expected {output_num_blocks} outputs, produced {len(new_blocks)}"
+
+        output = []
+        for block, meta_with_schema in zip(new_blocks, reduce_stage_metadata_schema):
+            output.append(
+                RefBundle(
+                    [
+                        (
+                            block,
+                            meta_with_schema.metadata,
+                        )
+                    ],
+                    owns_blocks=input_owned,
+                    schema=meta_with_schema.schema,
+                )
+            )
+
+        stats = {
+            "map": to_stats(map_stage_metadata_schema),
+            "merge": to_stats(merge_stage_metadata_schema),
+            "reduce": to_stats(reduce_stage_metadata_schema),
+        }
+
+        return (output, stats)
+
+    @staticmethod
+    def _map_partition(
+        map_fn,
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        schedule: _MergeTaskSchedule,
+        *map_args: List[Any],
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        mapper_outputs = map_fn(idx, block, output_num_blocks, *map_args)
+
+        # A merge task may produce results for multiple downstream reducer
+        # tasks. Therefore, each map task should give each merge task a
+        # partition of its outputs, where the length of the partition is equal
+        # to the number of reducers downstream to the merge task.
+        partition = []
+        merge_idx = 0
+        while merge_idx < schedule.num_merge_tasks_per_round and mapper_outputs:
+            output = mapper_outputs.pop(0)
+            partition.append(output)
+
+            if len(partition) == schedule.get_num_reducers_per_merge_idx(merge_idx):
+                yield partition
+
+                partition = []
+                merge_idx += 1
+
+        assert not partition
+        assert len(mapper_outputs) == 1, (
+            mapper_outputs,
+            "The last output should be a BlockMetadataWithSchema",
+        )
+        assert isinstance(mapper_outputs[0], BlockMetadataWithSchema)
+        yield mapper_outputs[0]
+
+        assert merge_idx == schedule.num_merge_tasks_per_round, (
+            merge_idx,
+            schedule.num_merge_tasks_per_round,
+        )
+
+    @staticmethod
+    def _merge(
+        reduce_fn,
+        *all_mapper_outputs: List[List[Block]],
+        reduce_args: Optional[List[Any]] = None,
+    ) -> List[Union["BlockMetadataWithSchema", Block]]:
+        """
+        Returns list of [BlockMetadata, O1, O2, O3, ...output_num_blocks].
+        """
+        assert (
+            len({len(mapper_outputs) for mapper_outputs in all_mapper_outputs}) == 1
+        ), "Received different number of map inputs"
+        stats = BlockExecStats.builder()
+        if not reduce_args:
+            reduce_args = []
+
+        num_rows = 0
+        size_bytes = 0
+        schemas = []
+        for i, mapper_outputs in enumerate(zip(*all_mapper_outputs)):
+            block_meta_with_schema: Tuple[Block, "BlockMetadataWithSchema"] = reduce_fn(
+                *reduce_args, *mapper_outputs, partial_reduce=True
+            )
+            block, meta_with_schema = block_meta_with_schema
+            yield block
+
+            block = BlockAccessor.for_block(block)
+            num_rows += block.num_rows()
+            size_bytes += block.size_bytes()
+            del block
+            schemas.append(meta_with_schema.schema)
+
+        schema = _take_first_non_empty_schema(iter(schemas))
+
+        meta = BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=size_bytes,
+            input_files=None,
+            exec_stats=stats.build(),
+        )
+        meta_with_schema = BlockMetadataWithSchema(metadata=meta, schema=schema)
+        yield meta_with_schema
+
+    @staticmethod
+    def _compute_shuffle_schedule(
+        num_cpus_per_node_map: Dict[str, int],
+        num_input_blocks: int,
+        merge_factor: float,
+        num_output_blocks: int,
+    ) -> _PushBasedShuffleStage:
+        num_cpus_total = sum(v for v in num_cpus_per_node_map.values())
+        logger.debug(
+            f"Found {num_cpus_total} CPUs available CPUs for push-based shuffle."
+        )
+        num_tasks_per_map_merge_group = merge_factor + 1
+        num_total_merge_tasks = math.ceil(num_input_blocks / merge_factor)
+
+        num_merge_tasks_per_round = 0
+        merge_task_placement = []
+        leftover_cpus = 0
+        # Compute the total number of merge tasks and their node placement.
+        # Each merge task should be grouped with `merge_factor` map tasks for
+        # pipelining. These groups should then be spread across nodes according
+        # to CPU availability for load-balancing.
+        num_input_blocks_remaining = num_input_blocks
+        for node, num_cpus in num_cpus_per_node_map.items():
+            # First find how many merge tasks we should run on this node.
+            # We take the min of the number of CPUs on this node and the number
+            # of input blocks that we haven't scheduled yet, in case there are
+            # fewer input blocks than CPU slots on this node.
+            num_cpu_slots = min(num_cpus, num_input_blocks_remaining)
+            num_merge_tasks_on_cur_node = round(
+                num_cpu_slots / num_tasks_per_map_merge_group
+            )
+            # For small datasets, the number of tasks to run may be less than
+            # the total CPU slots available.
+            num_merge_tasks_on_cur_node = min(
+                num_merge_tasks_on_cur_node, num_total_merge_tasks
+            )
+            for i in range(num_merge_tasks_on_cur_node):
+                merge_task_placement.append(node)
+                # We schedule `merge_factor` many map tasks for every merge
+                # task. Subtract from the number of input blocks remaining to
+                # account for cases where the number of map tasks is smaller
+                # than the available CPU slots.
+                num_input_blocks_remaining -= merge_factor
+                num_cpus -= num_tasks_per_map_merge_group
+            num_merge_tasks_per_round += num_merge_tasks_on_cur_node
+
+            # Handle the case where a single node cannot fit a group of map and
+            # merge tasks, but we can spread the group across multiple distinct
+            # nodes.
+            leftover_cpus += num_cpus
+            if (
+                leftover_cpus >= num_tasks_per_map_merge_group
+                and num_merge_tasks_per_round < num_total_merge_tasks
+            ):
+                merge_task_placement.append(node)
+                num_merge_tasks_per_round += 1
+                leftover_cpus -= num_tasks_per_map_merge_group
+                num_input_blocks_remaining -= merge_factor
+
+            num_input_blocks_remaining = max(0, num_input_blocks_remaining)
+
+        if num_merge_tasks_per_round == 0:
+            # For small datasets, make sure we have at least one merge task.
+            for node, num_cpus in num_cpus_per_node_map.items():
+                if num_cpus >= 1:
+                    merge_task_placement.append(node)
+                    num_merge_tasks_per_round = 1
+                    break
+
+        assert num_merge_tasks_per_round == len(merge_task_placement)
+        assert num_merge_tasks_per_round > 0, num_merge_tasks_per_round
+        # Use the remaining CPUs to execute map tasks.
+        num_map_tasks_per_round = num_cpus_total - num_merge_tasks_per_round
+        num_map_tasks_per_round = min(num_map_tasks_per_round, num_input_blocks)
+        # Make sure there is at least one map task in each round.
+        num_map_tasks_per_round = max(num_map_tasks_per_round, 1)
+
+        num_rounds = math.ceil(num_input_blocks / num_map_tasks_per_round)
+        return _PushBasedShuffleStage(
+            num_output_blocks,
+            num_rounds,
+            num_map_tasks_per_round,
+            merge_task_placement,
+        )
+
+
+def _get_num_cpus_per_node_map() -> Dict[str, int]:
+    total_resources_by_node = ray.state.total_resources_per_node()
+    # Map from per-node resource name to number of CPUs available on that
+    # node.
+    num_cpus_per_node_map = {}
+    for node_id, resources in total_resources_by_node.items():
+        num_cpus = int(resources.get("CPU", 0))
+        if num_cpus == 0:
+            continue
+        num_cpus_per_node_map[node_id] = num_cpus
+    return num_cpus_per_node_map
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/shuffle_task_spec.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/shuffle_task_spec.py
new file mode 100644
index 0000000000000000000000000000000000000000..48c4c9f84a54b5b70964c6695dee21f6e8e51d0f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/shuffle_task_spec.py
@@ -0,0 +1,146 @@
+import logging
+import math
+from typing import Callable, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    BlockExecStats,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+)
+from ray.data.context import MAX_SAFE_BLOCK_SIZE_FACTOR
+
+logger = logging.getLogger(__name__)
+
+
+class ShuffleTaskSpec(ExchangeTaskSpec):
+    """
+    The implementation for shuffle tasks.
+
+    This is used by random_shuffle() and repartition().
+    """
+
+    SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME = "Split Repartition"
+
+    def __init__(
+        self,
+        target_shuffle_max_block_size: int,
+        random_shuffle: bool = False,
+        random_seed: Optional[int] = None,
+        upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]] = None,
+    ):
+        super().__init__(
+            map_args=[
+                target_shuffle_max_block_size,
+                upstream_map_fn,
+                random_shuffle,
+                random_seed,
+            ],
+            reduce_args=[random_shuffle, random_seed],
+        )
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        target_shuffle_max_block_size: int,
+        upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]],
+        random_shuffle: bool,
+        random_seed: Optional[int],
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        stats = BlockExecStats.builder()
+        if upstream_map_fn:
+            # TODO: Support dynamic block splitting in
+            # all-to-all ops, to avoid having to re-fuse
+            # upstream blocks together.
+            upstream_map_iter = upstream_map_fn([block])
+            mapped_block = next(upstream_map_iter)
+            builder = BlockAccessor.for_block(mapped_block).builder()
+            builder.add_block(mapped_block)
+            for mapped_block in upstream_map_iter:
+                builder.add_block(mapped_block)
+            # Drop the upstream inputs to reduce memory usage.
+            del mapped_block
+            block = builder.build()
+        block = BlockAccessor.for_block(block)
+        if (
+            block.size_bytes()
+            > MAX_SAFE_BLOCK_SIZE_FACTOR * target_shuffle_max_block_size
+        ):
+            logger.warning(
+                "Input block to map task has size "
+                f"{block.size_bytes() // (1024 * 1024)}MiB, which exceeds "
+                "DataContext.get_current().target_shuffle_max_block_size="
+                f"{target_shuffle_max_block_size // (1024 * 1024)}MiB. "
+                "This can lead to out-of-memory errors and can happen "
+                "when map tasks are fused to the shuffle operation. "
+                "To prevent fusion, call Dataset.materialize() on the "
+                "dataset before shuffling."
+            )
+
+        # Randomize the distribution of records to blocks.
+        if random_shuffle:
+            seed_i = random_seed + idx if random_seed is not None else None
+            block = block.random_shuffle(seed_i)
+            block = BlockAccessor.for_block(block)
+
+        # Build a list of slices to return. It's okay to put the results in a
+        # list instead of yielding them as a generator because slicing the
+        # ArrowBlock is zero-copy.
+        slice_sz = max(1, math.ceil(block.num_rows() / output_num_blocks))
+        slices = []
+        for i in range(output_num_blocks):
+            slices.append(block.slice(i * slice_sz, (i + 1) * slice_sz))
+
+        # Randomize the distribution order of the blocks (this prevents empty
+        # outputs when input blocks are very small).
+        if random_shuffle:
+            random = np.random.RandomState(seed_i)
+            random.shuffle(slices)
+
+        num_rows = sum(BlockAccessor.for_block(s).num_rows() for s in slices)
+        assert num_rows == block.num_rows(), (num_rows, block.num_rows())
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta = block.get_metadata(exec_stats=stats.build())
+        schema = block.schema()
+        meta_with_schema = BlockMetadataWithSchema(metadata=meta, schema=schema)
+        return slices + [meta_with_schema]
+
+    @staticmethod
+    def reduce(
+        random_shuffle: bool,
+        random_seed: Optional[int],
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        # TODO: Support fusion with other downstream operators.
+        stats = BlockExecStats.builder()
+        builder = DelegatingBlockBuilder()
+        for block in mapper_outputs:
+            builder.add_block(block)
+        new_block = builder.build()
+        accessor = BlockAccessor.for_block(new_block)
+        if random_shuffle:
+            new_block = accessor.random_shuffle(
+                random_seed if random_seed is not None else None
+            )
+            accessor = BlockAccessor.for_block(new_block)
+        new_metadata = BlockMetadata(
+            num_rows=accessor.num_rows(),
+            size_bytes=accessor.size_bytes(),
+            input_files=None,
+            exec_stats=stats.build(),
+        )
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta_with_schema = BlockMetadataWithSchema(
+            metadata=new_metadata, schema=accessor.schema()
+        )
+        return new_block, meta_with_schema
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/sort_task_spec.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/sort_task_spec.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e843f77eadb2b505e0f7f233cb9e0037bc0789b
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/sort_task_spec.py
@@ -0,0 +1,239 @@
+from typing import TYPE_CHECKING, List, Optional, Tuple, TypeVar, Union
+
+import numpy as np
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.table_block import TableBlockAccessor
+from ray.data._internal.util import NULL_SENTINEL
+from ray.data.block import Block, BlockAccessor, BlockExecStats
+from ray.types import ObjectRef
+
+T = TypeVar("T")
+
+if TYPE_CHECKING:
+    import pyarrow
+
+    from ray.data.block import BlockMetadataWithSchema
+
+
+class SortKey:
+    """SortKey class to convert between different sort args formats."""
+
+    def __init__(
+        self,
+        key: Optional[Union[str, List[str]]] = None,
+        descending: Union[bool, List[bool]] = False,
+        boundaries: Optional[List[T]] = None,
+    ):
+        if key is None:
+            key = []
+        if isinstance(key, str):
+            key = [key]
+        if not (isinstance(key, list) and all(isinstance(k, str) for k in key)):
+            raise ValueError(
+                f"Key must be a string or a list of strings, but got {key}."
+            )
+        if isinstance(descending, bool):
+            descending = [descending for _ in key]
+        elif isinstance(descending, list):
+            if len(descending) != len(key):
+                raise ValueError(
+                    "Length of `descending` does not match the length of the key."
+                )
+        self._columns = key
+        self._descending = descending
+        if boundaries:
+            for item in boundaries:
+                if not isinstance(item, (int, float)):
+                    raise ValueError(
+                        "The type of items in boundaries must be int or float."
+                    )
+            boundaries = list(set(boundaries))
+            boundaries.sort()
+        self._boundaries = boundaries
+
+    def get_columns(self) -> List[str]:
+        return self._columns
+
+    def get_descending(self) -> List[bool]:
+        return self._descending
+
+    def to_arrow_sort_args(self) -> List[Tuple[str, str]]:
+        return [
+            (key, "descending" if desc else "ascending")
+            for key, desc in zip(self._columns, self._descending)
+        ]
+
+    def to_pandas_sort_args(self) -> Tuple[List[str], List[bool]]:
+        return self._columns, [not desc for desc in self._descending]
+
+    def validate_schema(self, schema: Optional[Union[type, "pyarrow.lib.Schema"]]):
+        """Check the key function is valid on the given schema."""
+        if schema is None:
+            # Dataset is empty/cleared, validation not possible.
+            return
+
+        if self._columns and len(schema.names) > 0:
+            schema_names_set = set(schema.names)
+            for column in self._columns:
+                if column not in schema_names_set:
+                    raise ValueError(
+                        f"You specified the column '{column}', but there's no such "
+                        "column in the dataset. The dataset has columns: "
+                        f"{schema.names}"
+                    )
+
+    @property
+    def boundaries(self):
+        return self._boundaries
+
+
+class SortTaskSpec(ExchangeTaskSpec):
+    """
+    The implementation for distributed sort tasks.
+
+    The algorithm is similar to [External Merge Sort]
+    (https://en.wikipedia.org/wiki/External_sorting).
+    Sorting is done in 3 steps: sampling, sorting individual blocks, and
+    merging sorted blocks.
+
+    Sampling (`sample_boundaries`): we get a number of sample items from each block,
+    sort them, and use them to compute boundaries that would partition all items into
+    approximately equal ranges.
+
+    Sorting (`map`): each block is sorted locally, then partitioned into smaller
+    blocks according to the boundaries. Each partitioned block is passed to a merge
+    task.
+
+    Merging (`reduce`): a merge task would receive a block from every worker that
+    consists of items in a certain range. It then merges the sorted blocks into one
+    sorted block and becomes part of the new, sorted block.
+    """
+
+    SORT_SAMPLE_SUB_PROGRESS_BAR_NAME = "Sort Sample"
+
+    def __init__(
+        self,
+        boundaries: List[T],
+        sort_key: SortKey,
+        batch_format: str,
+    ):
+        super().__init__(
+            map_args=[boundaries, sort_key],
+            reduce_args=[sort_key, batch_format],
+        )
+
+    @staticmethod
+    def map(
+        idx: int,
+        block: Block,
+        output_num_blocks: int,
+        boundaries: List[T],
+        sort_key: SortKey,
+    ) -> List[Union[Block, "BlockMetadataWithSchema"]]:
+        stats = BlockExecStats.builder()
+        accessor = BlockAccessor.for_block(block)
+        out = accessor.sort_and_partition(boundaries, sort_key)
+        from ray.data.block import BlockMetadataWithSchema
+
+        meta_with_schema = BlockMetadataWithSchema.from_block(
+            block, stats=stats.build()
+        )
+        return out + [meta_with_schema]
+
+    @staticmethod
+    def reduce(
+        sort_key: SortKey,
+        batch_format: str,
+        *mapper_outputs: List[Block],
+        partial_reduce: bool = False,
+    ) -> Tuple[Block, "BlockMetadataWithSchema"]:
+        normalized_blocks = TableBlockAccessor.normalize_block_types(
+            mapper_outputs,
+            target_block_type=ExchangeTaskSpec._derive_target_block_type(batch_format),
+        )
+        blocks, meta_with_schema = BlockAccessor.for_block(
+            normalized_blocks[0]
+        ).merge_sorted_blocks(normalized_blocks, sort_key)
+        return blocks, meta_with_schema
+
+    @staticmethod
+    def sample_boundaries(
+        blocks: List[ObjectRef[Block]],
+        sort_key: SortKey,
+        num_reducers: int,
+        sample_bar: Optional[ProgressBar] = None,
+    ) -> List[T]:
+        """
+        Return (num_reducers - 1) items in ascending order from the blocks that
+        partition the domain into ranges with approximately equally many elements.
+        Each boundary item is a tuple of a form (col1_value, col2_value, ...).
+        """
+        columns = sort_key.get_columns()
+        n_samples = int(num_reducers * 10 / len(blocks))
+
+        sample_block = cached_remote_fn(_sample_block)
+
+        sample_results = [
+            sample_block.remote(block, n_samples, sort_key) for block in blocks
+        ]
+        if sample_bar is None:
+            sample_bar = ProgressBar(
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME,
+                len(blocks) * n_samples,
+                unit="rows",
+            )
+        # TODO(zhilong): Update sort sample bar before finished.
+        samples = sample_bar.fetch_until_complete(sample_results)
+        del sample_results
+        samples: List[Block] = [s for s in samples if len(s) > 0]
+        # The dataset is empty
+        if len(samples) == 0:
+            return [None] * (num_reducers - 1)
+
+        # Convert samples to a sorted list[tuple[...]] where each tuple represents a
+        # sample.
+        # TODO: Once we deprecate pandas blocks, we can avoid this conversion and
+        # directly sort the samples.
+        builder = DelegatingBlockBuilder()
+        for sample in samples:
+            builder.add_block(sample)
+        samples_table = builder.build()
+        samples_dict = BlockAccessor.for_block(samples_table).to_numpy(columns=columns)
+        # This zip does the transposition from list of column values to list of tuples.
+        samples_list = list(zip(*samples_dict.values()))
+
+        def is_na(x):
+            # Check if x is None or NaN. Type casting to np.array first to avoid
+            # isnan failing on strings and other types.
+            if x is None:
+                return True
+            x = np.asarray(x)
+            if np.issubdtype(x.dtype, np.number):
+                return np.isnan(x)
+            return False
+
+        # To allow multi-directional sort, we utilize Python's stable sort: we
+        # sort several times with different directions. We do this in reverse, so
+        # that the last key we sort by is the primary sort key passed by the user.
+        for i, desc in list(enumerate(sort_key.get_descending()))[::-1]:
+            # Sort the list, but Nones should be NULL_SENTINEL to ensure safe sorting.
+            samples_list.sort(
+                key=lambda sample: NULL_SENTINEL if is_na(sample[i]) else sample[i],
+                reverse=desc,
+            )
+
+        # Each boundary corresponds to a quantile of the data.
+        quantile_indices = [
+            int(q * (len(samples_list) - 1))
+            for q in np.linspace(0, 1, num_reducers + 1)
+        ]
+        # Exclude the first and last quantiles because they're 0 and 1.
+        return [samples_list[i] for i in quantile_indices[1:-1]]
+
+
+def _sample_block(block: Block, n_samples: int, sort_key: SortKey) -> Block:
+    return BlockAccessor.for_block(block).sample(n_samples, sort_key)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/split_repartition_task_scheduler.py b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/split_repartition_task_scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c3a563ef6948b52033d60ba415ea37adf26b507
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/exchange/split_repartition_task_scheduler.py
@@ -0,0 +1,166 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import ray
+from ray.data._internal.execution.interfaces import RefBundle, TaskContext
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.planner.exchange.interfaces import (
+    ExchangeTaskScheduler,
+)
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.split import _split_at_indices
+from ray.data._internal.util import unzip
+from ray.data.block import (
+    Block,
+    BlockMetadata,
+    BlockMetadataWithSchema,
+)
+from ray.types import ObjectRef
+
+
+class SplitRepartitionTaskScheduler(ExchangeTaskScheduler):
+    """
+    The split (non-shuffle) repartition scheduler.
+
+    First, we calculate global splits needed to produce `output_num_blocks` blocks.
+    After the split blocks are generated accordingly, reduce tasks are scheduled
+    to combine split blocks together.
+    """
+
+    def execute(
+        self,
+        refs: List[RefBundle],
+        output_num_blocks: int,
+        ctx: TaskContext,
+        map_ray_remote_args: Optional[Dict[str, Any]] = None,
+        reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
+    ) -> AllToAllTransformFnResult:
+        input_num_rows = 0
+        input_owned_by_consumer = True
+        for ref_bundle in refs:
+            block_num_rows = ref_bundle.num_rows()
+            if block_num_rows is None:
+                raise ValueError(
+                    "Cannot split partition on blocks with unknown number of rows."
+                )
+            input_num_rows += block_num_rows
+            if not ref_bundle.owns_blocks:
+                input_owned_by_consumer = False
+
+        # Compute the (output_num_blocks) indices needed for an equal split of the
+        # input blocks. When output_num_blocks=1, the total number of
+        # input rows is used as the end index during the split calculation,
+        # so that we can combine all input blocks into a single output block.
+        indices = []
+        if output_num_blocks == 1:
+            indices = [input_num_rows]
+        else:
+            cur_idx = 0
+            for _ in range(output_num_blocks - 1):
+                cur_idx += input_num_rows / output_num_blocks
+                indices.append(int(cur_idx))
+        assert len(indices) <= output_num_blocks, (indices, output_num_blocks)
+
+        if map_ray_remote_args is None:
+            map_ray_remote_args = {}
+        if reduce_ray_remote_args is None:
+            reduce_ray_remote_args = {}
+        if "scheduling_strategy" not in reduce_ray_remote_args:
+            reduce_ray_remote_args = reduce_ray_remote_args.copy()
+            reduce_ray_remote_args["scheduling_strategy"] = "SPREAD"
+
+        blocks_with_metadata: List[Tuple[ObjectRef[Block], BlockMetadata]] = []
+        for ref_bundle in refs:
+            blocks_with_metadata.extend(ref_bundle.blocks)
+        split_return = _split_at_indices(
+            blocks_with_metadata, indices, input_owned_by_consumer
+        )
+        split_block_refs, split_metadata = [], []
+        for b, m in zip(*split_return):
+            split_block_refs.append(b)
+            split_metadata.extend(m)
+
+        sub_progress_bar_dict = ctx.sub_progress_bar_dict
+        bar_name = ShuffleTaskSpec.SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME
+        assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
+        reduce_bar = sub_progress_bar_dict[bar_name]
+
+        reduce_task = cached_remote_fn(self._exchange_spec.reduce)
+        reduce_return = [
+            reduce_task.options(**reduce_ray_remote_args, num_returns=2).remote(
+                *self._exchange_spec._reduce_args,
+                *split_block_refs[j],
+            )
+            for j in range(output_num_blocks)
+            # Only process splits which contain blocks.
+            if len(split_block_refs[j]) > 0
+        ]
+
+        reduce_block_refs, reduce_metadata_schema = [], []
+        if reduce_return:
+            reduce_block_refs, reduce_metadata_schema = unzip(reduce_return)
+        reduce_metadata_schema: List[
+            "BlockMetadataWithSchema"
+        ] = reduce_bar.fetch_until_complete(list(reduce_metadata_schema))
+        reduce_block_refs = list(reduce_block_refs)
+
+        # Handle empty blocks.
+        if len(reduce_block_refs) < output_num_blocks:
+            import pyarrow as pa
+
+            from ray.data._internal.arrow_block import ArrowBlockBuilder
+            from ray.data._internal.pandas_block import (
+                PandasBlockBuilder,
+                PandasBlockSchema,
+            )
+
+            num_empty_blocks = output_num_blocks - len(reduce_block_refs)
+            if len(reduce_metadata_schema) > 0:
+                first_block_schema = reduce_metadata_schema[0].schema
+                if isinstance(first_block_schema, pa.Schema):
+                    builder = ArrowBlockBuilder()
+                elif isinstance(first_block_schema, PandasBlockSchema):
+                    builder = PandasBlockBuilder()
+                else:
+                    raise ValueError(
+                        "Cannot split partition on blocks with unknown block schema:"
+                        f" {first_block_schema}."
+                    )
+            else:
+                # If the result is empty, default to Arrow format for the empty blocks.
+                builder = ArrowBlockBuilder()
+
+            empty_block = builder.build()
+            empty_meta_with_schema = BlockMetadataWithSchema.from_block(
+                empty_block
+            )  # No stats for empty block.
+            empty_block_refs, empty_metadata = zip(
+                *[
+                    (ray.put(empty_block), empty_meta_with_schema)
+                    for _ in range(num_empty_blocks)
+                ]
+            )
+            reduce_block_refs.extend(empty_block_refs)
+            reduce_metadata_schema.extend(empty_metadata)
+
+        output = []
+        assert len(reduce_block_refs) == len(reduce_metadata_schema), (
+            len(reduce_block_refs),
+            len(reduce_metadata_schema),
+        )
+        for block, meta_with_schema in zip(reduce_block_refs, reduce_metadata_schema):
+            output.append(
+                RefBundle(
+                    [(block, meta_with_schema.metadata)],
+                    owns_blocks=input_owned_by_consumer,
+                    schema=meta_with_schema.schema,
+                )
+            )
+        stats = {
+            "split": split_metadata,
+            "reduce": reduce_metadata_schema,
+        }
+
+        return (output, stats)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_all_to_all_op.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_all_to_all_op.py
new file mode 100644
index 0000000000000000000000000000000000000000..75958bf81aca2c14c7edffcb8745ed6a1c6849f3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_all_to_all_op.py
@@ -0,0 +1,163 @@
+from typing import List
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.base_physical_operator import (
+    AllToAllOperator,
+)
+from ray.data._internal.logical.operators.all_to_all_operator import (
+    AbstractAllToAll,
+    Aggregate,
+    RandomizeBlocks,
+    RandomShuffle,
+    Repartition,
+    Sort,
+)
+from ray.data._internal.planner.aggregate import generate_aggregate_fn
+from ray.data._internal.planner.random_shuffle import generate_random_shuffle_fn
+from ray.data._internal.planner.randomize_blocks import generate_randomize_blocks_fn
+from ray.data._internal.planner.repartition import generate_repartition_fn
+from ray.data._internal.planner.sort import generate_sort_fn
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def _plan_hash_shuffle_repartition(
+    data_context: DataContext,
+    logical_op: Repartition,
+    input_physical_op: PhysicalOperator,
+) -> PhysicalOperator:
+    from ray.data._internal.execution.operators.hash_shuffle import (
+        HashShuffleOperator,
+    )
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+    normalized_key_columns = SortKey(logical_op._keys).get_columns()
+
+    return HashShuffleOperator(
+        input_physical_op,
+        data_context,
+        key_columns=tuple(normalized_key_columns),  # noqa: type
+        # NOTE: In case number of partitions is not specified, we fall back to
+        #       default min parallelism configured
+        num_partitions=logical_op._num_outputs,
+        should_sort=logical_op._sort,
+        # TODO wire in aggregator args overrides
+    )
+
+
+def _plan_hash_shuffle_aggregate(
+    data_context: DataContext,
+    logical_op: Aggregate,
+    input_physical_op: PhysicalOperator,
+) -> PhysicalOperator:
+    from ray.data._internal.execution.operators.hash_aggregate import (
+        HashAggregateOperator,
+    )
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+    normalized_key_columns = SortKey(logical_op._key).get_columns()
+
+    return HashAggregateOperator(
+        data_context,
+        input_physical_op,
+        key_columns=tuple(normalized_key_columns),  # noqa: type
+        aggregation_fns=tuple(logical_op._aggs),  # noqa: type
+        # NOTE: In case number of partitions is not specified, we fall back to
+        #       default min parallelism configured
+        num_partitions=logical_op._num_partitions,
+        # TODO wire in aggregator args overrides
+    )
+
+
+def plan_all_to_all_op(
+    op: AbstractAllToAll,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Get the corresponding physical operators DAG for AbstractAllToAll operators.
+
+    Note this method only converts the given `op`, but not its input dependencies.
+    See Planner.plan() for more details.
+    """
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    if isinstance(op, RandomizeBlocks):
+        fn = generate_randomize_blocks_fn(op)
+        # Randomize block order does not actually compute anything, so we
+        # want to inherit the upstream op's target max block size.
+
+    elif isinstance(op, RandomShuffle):
+        debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+            "debug_limit_shuffle_execution_to_num_blocks", None
+        )
+        fn = generate_random_shuffle_fn(
+            data_context,
+            op._seed,
+            op._num_outputs,
+            op._ray_remote_args,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+
+    elif isinstance(op, Repartition):
+        if op._keys:
+            if data_context.shuffle_strategy == ShuffleStrategy.HASH_SHUFFLE:
+                return _plan_hash_shuffle_repartition(
+                    data_context, op, input_physical_dag
+                )
+            else:
+                raise ValueError(
+                    "Key-based repartitioning only supported for "
+                    f"`DataContext.shuffle_strategy=HASH_SHUFFLE` "
+                    f"(got {data_context.shuffle_strategy})"
+                )
+
+        elif op._shuffle:
+            debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+                "debug_limit_shuffle_execution_to_num_blocks", None
+            )
+        else:
+            debug_limit_shuffle_execution_to_num_blocks = None
+
+        fn = generate_repartition_fn(
+            op._num_outputs,
+            op._shuffle,
+            data_context,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+
+    elif isinstance(op, Sort):
+        debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+            "debug_limit_shuffle_execution_to_num_blocks", None
+        )
+        fn = generate_sort_fn(
+            op._sort_key,
+            op._batch_format,
+            data_context,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+
+    elif isinstance(op, Aggregate):
+        if data_context.shuffle_strategy == ShuffleStrategy.HASH_SHUFFLE:
+            return _plan_hash_shuffle_aggregate(data_context, op, input_physical_dag)
+
+        debug_limit_shuffle_execution_to_num_blocks = data_context.get_config(
+            "debug_limit_shuffle_execution_to_num_blocks", None
+        )
+        fn = generate_aggregate_fn(
+            op._key,
+            op._aggs,
+            op._batch_format,
+            data_context,
+            debug_limit_shuffle_execution_to_num_blocks,
+        )
+    else:
+        raise ValueError(f"Found unknown logical operator during planning: {op}")
+
+    return AllToAllOperator(
+        fn,
+        input_physical_dag,
+        data_context,
+        num_outputs=op._num_outputs,
+        sub_progress_bar_names=op._sub_progress_bar_names,
+        name=op.name,
+    )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_download_op.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_download_op.py
new file mode 100644
index 0000000000000000000000000000000000000000..abc2b06897b2387739a387a1b84bb5af346ff40f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_download_op.py
@@ -0,0 +1,350 @@
+import logging
+import math
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Iterator, List
+from urllib.parse import urlparse
+
+import pyarrow as pa
+
+import ray
+from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.actor_pool_map_operator import (
+    ActorPoolMapOperator,
+)
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.logical.operators.one_to_one_operator import Download
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.util import RetryingPyFileSystem, make_async_gen
+from ray.data.block import BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+
+logger = logging.getLogger(__name__)
+
+URI_DOWNLOAD_MAX_WORKERS = 16
+
+
+def plan_download_op(
+    op: Download,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    """Plan the download operation with partitioning and downloading stages."""
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    upstream_op_is_download = False
+    if len(input_physical_dag._logical_operators) == 1 and isinstance(
+        input_physical_dag._logical_operators[0], Download
+    ):
+        upstream_op_is_download = True
+
+    uri_column_names = op.uri_column_names
+    uri_column_names_str = ", ".join(uri_column_names)
+    output_bytes_column_names = op.output_bytes_column_names
+    ray_remote_args = op.ray_remote_args
+
+    # Import _get_udf from the main planner file
+    from ray.data._internal.planner.plan_udf_map_op import (
+        _generate_transform_fn_for_map_batches,
+        _get_udf,
+    )
+
+    # If we have multiple download operators in a row, we should only include the partition actor
+    # at the start of the chain. This is primarily done to prevent partition actors from bottlenecking
+    # the chain becuase the interleaved operators would be a single actor. As a result, the
+    # URIDownloader physical operator is responsible for outputting appropriately sized blocks.
+    partition_map_operator = None
+    if not upstream_op_is_download:
+        # PartitionActor is a callable class, so we need ActorPoolStrategy
+        partition_compute = ActorPoolStrategy(
+            size=1
+        )  # Use single actor for partitioning
+
+        fn, init_fn = _get_udf(
+            PartitionActor, (), {}, (uri_column_names, data_context), {}
+        )
+        block_fn = _generate_transform_fn_for_map_batches(fn)
+
+        partition_transform_fns = [
+            BlockMapTransformFn(
+                block_fn,
+                # NOTE: Disable block-shaping to produce blocks as is
+                disable_block_shaping=True,
+            ),
+        ]
+        partition_map_transformer = MapTransformer(
+            partition_transform_fns,
+            init_fn=init_fn,
+        )
+
+        partition_map_operator = ActorPoolMapOperator(
+            partition_map_transformer,
+            input_physical_dag,
+            data_context,
+            name=f"Partition({uri_column_names_str})",
+            # NOTE: Partition actor doesn't use the user-provided `ray_remote_args`
+            #       since those only apply to the actual download tasks. Partitioning is
+            #       a lightweight internal operation that doesn't need custom resource
+            #       requirements.
+            ray_remote_args=None,
+            compute_strategy=partition_compute,  # Use actor-based compute for callable class
+            # NOTE: We set `_generator_backpressure_num_objects` to -1 to unblock
+            #       backpressure since partitioning is extremely fast. Without this, the
+            #       partition actor gets bottlenecked by the Ray Data scheduler, which
+            #       can prevent Ray Data from launching enough download tasks.
+            ray_actor_task_remote_args={"_generator_backpressure_num_objects": -1},
+        )
+
+    fn, init_fn = _get_udf(
+        download_bytes_threaded,
+        (uri_column_names, output_bytes_column_names, data_context),
+        {},
+        None,
+        None,
+    )
+
+    download_transform_fn = _generate_transform_fn_for_map_batches(fn)
+    transform_fns = [
+        BlockMapTransformFn(
+            download_transform_fn,
+            output_block_size_option=OutputBlockSizeOption.of(
+                target_max_block_size=data_context.target_max_block_size
+            ),
+        ),
+    ]
+
+    download_compute = TaskPoolStrategy()
+    download_map_transformer = MapTransformer(
+        transform_fns,
+        init_fn=init_fn,
+    )
+
+    download_map_operator = MapOperator.create(
+        download_map_transformer,
+        partition_map_operator if partition_map_operator else input_physical_dag,
+        data_context,
+        name=f"Download({uri_column_names_str})",
+        compute_strategy=download_compute,
+        ray_remote_args=ray_remote_args,
+    )
+
+    return download_map_operator
+
+
+def uri_to_path(uri: str) -> str:
+    """Convert a URI to a filesystem path."""
+    # TODO(mowen): urlparse might be slow. in the future we could use a faster alternative.
+    parsed = urlparse(uri)
+    if parsed.scheme == "file":
+        return parsed.path
+    return parsed.netloc + parsed.path
+
+
+def _arrow_batcher(table: pa.Table, output_batch_size: int):
+    """Batch a PyArrow table into smaller tables of size n using zero-copy slicing."""
+    num_rows = table.num_rows
+    for i in range(0, num_rows, output_batch_size):
+        end_idx = min(i + output_batch_size, num_rows)
+        # Use PyArrow's zero-copy slice operation
+        batch_table = table.slice(i, end_idx - i)
+        yield batch_table
+
+
+def download_bytes_threaded(
+    block: pa.Table,
+    uri_column_names: List[str],
+    output_bytes_column_names: List[str],
+    data_context: DataContext,
+) -> Iterator[pa.Table]:
+    """Optimized version that uses make_async_gen for concurrent downloads.
+
+    Supports downloading from multiple URI columns in a single operation.
+    """
+    if not isinstance(block, pa.Table):
+        block = BlockAccessor.for_block(block).to_arrow()
+
+    output_block = block
+
+    # Download each URI column and add it to the output block
+    for uri_column_name, output_bytes_column_name in zip(
+        uri_column_names, output_bytes_column_names
+    ):
+        # Extract URIs from PyArrow table
+        uris = output_block.column(uri_column_name).to_pylist()
+
+        if len(uris) == 0:
+            continue
+
+        paths, fs = _resolve_paths_and_filesystem(uris)
+        fs = RetryingPyFileSystem.wrap(
+            fs, retryable_errors=data_context.retried_io_errors
+        )
+
+        def load_uri_bytes(uri_path_iterator):
+            """Function that takes an iterator of URI paths and yields downloaded bytes for each."""
+            for uri_path in uri_path_iterator:
+                read_bytes = None
+                try:
+                    # Use open_input_stream to handle the rare scenario where the data source is not seekable.
+                    with fs.open_input_stream(uri_path) as f:
+                        read_bytes = f.read()
+                except OSError as e:
+                    logger.debug(
+                        f"OSError reading uri '{uri_path}' for column '{uri_column_name}': {e}"
+                    )
+                except Exception as e:
+                    # Catch unexpected errors like pyarrow.lib.ArrowInvalid caused by an invalid uri like
+                    # `foo://bar` to avoid failing because of one invalid uri.
+                    logger.warning(
+                        f"Unexpected error reading uri '{uri_path}' for column '{uri_column_name}': {e}"
+                    )
+                finally:
+                    yield read_bytes
+
+        # Use make_async_gen to download URI bytes concurrently
+        # This preserves the order of results to match the input URIs
+        uri_bytes = list(
+            make_async_gen(
+                base_iterator=iter(paths),
+                fn=load_uri_bytes,
+                preserve_ordering=True,
+                num_workers=URI_DOWNLOAD_MAX_WORKERS,
+            )
+        )
+
+        # Add the new column to the PyArrow table
+        output_block = output_block.add_column(
+            len(output_block.column_names),
+            output_bytes_column_name,
+            pa.array(uri_bytes),
+        )
+
+    output_block_size = output_block.nbytes
+    ctx = ray.data.context.DatasetContext.get_current()
+    max_bytes = ctx.target_max_block_size
+    if max_bytes is not None and output_block_size > max_bytes:
+        num_blocks = math.ceil(output_block_size / max_bytes)
+        num_rows = output_block.num_rows
+        yield from _arrow_batcher(output_block, int(math.ceil(num_rows / num_blocks)))
+    else:
+        yield output_block
+
+
+class PartitionActor:
+    """Actor that partitions download operations based on estimated file sizes.
+
+    For multiple URI columns, estimates the combined size across all columns.
+    """
+
+    INIT_SAMPLE_BATCH_SIZE = 25
+
+    def __init__(self, uri_column_names: List[str], data_context: DataContext):
+        self._uri_column_names = uri_column_names
+        self._data_context = data_context
+        self._batch_size_estimate = None
+
+    def __call__(self, block: pa.Table) -> Iterator[pa.Table]:
+        if not isinstance(block, pa.Table):
+            block = BlockAccessor.for_block(block).to_arrow()
+
+        # Validate all URI columns exist
+        for uri_column_name in self._uri_column_names:
+            if uri_column_name not in block.column_names:
+                raise ValueError(
+                    "Ray Data tried to download URIs from a column named "
+                    f"{uri_column_name!r}, but a column with that name doesn't "
+                    "exist. Is the specified download column correct?"
+                )
+
+        if self._batch_size_estimate is None:
+            self._batch_size_estimate = self._estimate_nrows_per_partition(block)
+
+        yield from _arrow_batcher(block, self._batch_size_estimate)
+
+    def _estimate_nrows_per_partition(self, block: pa.Table) -> int:
+        sampled_file_sizes_by_column = {}
+        for uri_column_name in self._uri_column_names:
+            # Extract URIs from PyArrow table for sampling
+            uris = block.column(uri_column_name).to_pylist()
+            sample_uris = uris[: self.INIT_SAMPLE_BATCH_SIZE]
+            sampled_file_sizes = self._sample_sizes(sample_uris)
+            sampled_file_sizes_by_column[uri_column_name] = sampled_file_sizes
+
+        # If we sample HTTP URIs, or if an error occurs during sampling, then the file
+        # sizes might be `None`. In these cases, we replace the `file_size` with 0.
+        sampled_file_sizes_by_column = {
+            uri_column_name: [
+                file_size if file_size is not None else 0
+                for file_size in sampled_file_sizes
+            ]
+            for uri_column_name, sampled_file_sizes in sampled_file_sizes_by_column.items()
+        }
+
+        # This is some fancy Python code to compute the file size of each row.
+        row_sizes = [
+            sum(file_sizes_in_row)
+            for file_sizes_in_row in zip(*sampled_file_sizes_by_column.values())
+        ]
+
+        target_nbytes_per_partition = self._data_context.target_max_block_size
+        avg_nbytes_per_row = sum(row_sizes) / len(row_sizes)
+        if avg_nbytes_per_row == 0:
+            logger.warning(
+                "Estimated average row size is 0. Falling back to using the number of "
+                "rows in the block as the partition size."
+            )
+            return len(block)
+
+        nrows_per_partition = math.floor(
+            target_nbytes_per_partition / avg_nbytes_per_row
+        )
+        return nrows_per_partition
+
+    def _sample_sizes(self, uris: List[str]) -> List[int]:
+        """Fetch file sizes in parallel using ThreadPoolExecutor."""
+
+        def get_file_size(uri_path, fs):
+            try:
+                return fs.get_file_info(uri_path).size
+            except Exception:
+                return None
+
+        # If no URIs, return empty list
+        if not uris:
+            return []
+
+        # Get the filesystem from the first URI
+        paths, fs = _resolve_paths_and_filesystem(uris)
+        fs = RetryingPyFileSystem.wrap(
+            fs, retryable_errors=self._data_context.retried_io_errors
+        )
+
+        # Use ThreadPoolExecutor for concurrent size fetching
+        file_sizes = [None] * len(paths)
+        with ThreadPoolExecutor(max_workers=URI_DOWNLOAD_MAX_WORKERS) as executor:
+            # Submit all size fetch tasks
+            future_to_file_index = {
+                executor.submit(get_file_size, uri_path, fs): file_index
+                for file_index, uri_path in enumerate(paths)
+            }
+
+            # Collect results as they complete (order doesn't matter)
+            for future in as_completed(future_to_file_index):
+                file_index = future_to_file_index[future]
+                try:
+                    size = future.result()
+                    file_sizes[file_index] = size if size is not None else 0
+                except Exception as e:
+                    logger.warning(f"Error fetching file size for download: {e}")
+                    file_sizes[file_index] = 0
+
+        assert all(
+            fs is not None for fs in file_sizes
+        ), "File size sampling did not complete for all paths"
+        return file_sizes
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__init__.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2365dd88bbb788db8caa207a3c67338892889e2e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/expression_evaluator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/expression_evaluator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..62629f164ca353742f31ae51ff53894eaf701e5c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/expression_evaluator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/expression_visitors.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/expression_visitors.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..347fb736065673177b002687cb133ea88a1fe2d2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/__pycache__/expression_visitors.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/expression_evaluator.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/expression_evaluator.py
new file mode 100644
index 0000000000000000000000000000000000000000..7cb630a8e5e9324cbba7d4ac406cabd410d75c12
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/expression_evaluator.py
@@ -0,0 +1,754 @@
+from __future__ import annotations
+
+import ast
+import logging
+import operator
+from typing import Any, Callable, Dict, List, TypeVar, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.dataset as ds
+
+from ray.data._internal.logical.rules.projection_pushdown import (
+    _extract_input_columns_renaming_mapping,
+)
+from ray.data.block import Block, BlockAccessor, BlockColumn, BlockType
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    Expr,
+    LiteralExpr,
+    Operation,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    _ExprVisitor,
+    col,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _pa_is_in(left: Any, right: Any) -> Any:
+    if not isinstance(right, (pa.Array, pa.ChunkedArray)):
+        right = pa.array(right.as_py() if isinstance(right, pa.Scalar) else right)
+    return pc.is_in(left, right)
+
+
+_PANDAS_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
+    Operation.ADD: operator.add,
+    Operation.SUB: operator.sub,
+    Operation.MUL: operator.mul,
+    Operation.DIV: operator.truediv,
+    Operation.FLOORDIV: operator.floordiv,
+    Operation.GT: operator.gt,
+    Operation.LT: operator.lt,
+    Operation.GE: operator.ge,
+    Operation.LE: operator.le,
+    Operation.EQ: operator.eq,
+    Operation.NE: operator.ne,
+    Operation.AND: operator.and_,
+    Operation.OR: operator.or_,
+    Operation.NOT: operator.invert,
+    Operation.IS_NULL: pd.isna,
+    Operation.IS_NOT_NULL: pd.notna,
+    Operation.IN: lambda left, right: left.isin(right),
+    Operation.NOT_IN: lambda left, right: ~left.isin(right),
+}
+
+
+def _is_pa_string_type(t: pa.DataType) -> bool:
+    return pa.types.is_string(t) or pa.types.is_large_string(t)
+
+
+def _is_pa_string_like(x: Union[pa.Array, pa.ChunkedArray]) -> bool:
+    t = x.type
+    if pa.types.is_dictionary(t):
+        t = t.value_type
+    return _is_pa_string_type(t)
+
+
+def _pa_decode_dict_string_array(x: Union[pa.Array, pa.ChunkedArray]) -> Any:
+    """Convert Arrow dictionary-encoded string arrays to regular string arrays.
+
+    Dictionary encoding stores strings as indices into a dictionary of unique values.
+    This function converts them back to regular string arrays for string operations.
+
+    Example:
+        # Input: pa.array(['a', 'b']).dictionary_encode()
+        #   -- dictionary: ["a", "b"]
+        #   -- indices: [0, 1]
+        # Output: regular string array ["a", "b"]
+    Args:
+        x: The input array to convert.
+    Returns:
+        The converted string array.
+    """
+    if pa.types.is_dictionary(x.type) and _is_pa_string_type(x.type.value_type):
+        return pc.cast(x, pa.string())
+    return x
+
+
+def _to_pa_string_input(x: Any) -> Any:
+    if isinstance(x, str):
+        return pa.scalar(x)
+    elif _is_pa_string_like(x) and isinstance(x, (pa.Array, pa.ChunkedArray)):
+        x = _pa_decode_dict_string_array(x)
+    else:
+        raise
+    return x
+
+
+def _pa_add_or_concat(left: Any, right: Any) -> Any:
+    if isinstance(left, pa.Scalar):
+        left = left.as_py()
+    if isinstance(right, pa.Scalar):
+        right = right.as_py()
+    # If either side is string-like, perform string concatenation.
+    if (
+        isinstance(left, str)
+        or isinstance(right, str)
+        or (isinstance(left, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(left))
+        or (
+            isinstance(right, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(right)
+        )
+    ):
+        left_input = _to_pa_string_input(left)
+        right_input = _to_pa_string_input(right)
+        return pc.binary_join_element_wise(left_input, right_input, "")
+    return pc.add(left, right)
+
+
+_ARROW_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
+    Operation.ADD: _pa_add_or_concat,
+    Operation.SUB: pc.subtract,
+    Operation.MUL: pc.multiply,
+    Operation.DIV: pc.divide,
+    Operation.FLOORDIV: lambda left, right: pc.floor(pc.divide(left, right)),
+    Operation.GT: pc.greater,
+    Operation.LT: pc.less,
+    Operation.GE: pc.greater_equal,
+    Operation.LE: pc.less_equal,
+    Operation.EQ: pc.equal,
+    Operation.NE: pc.not_equal,
+    Operation.AND: pc.and_kleene,
+    Operation.OR: pc.or_kleene,
+    Operation.NOT: pc.invert,
+    Operation.IS_NULL: pc.is_null,
+    Operation.IS_NOT_NULL: pc.is_valid,
+    Operation.IN: _pa_is_in,
+    Operation.NOT_IN: lambda left, right: pc.invert(_pa_is_in(left, right)),
+}
+
+
+# NOTE: (srinathk) There are 3 distinct stages of handling passed in exprs:
+# 1. Parsing it (as text)
+# 2. Resolving unbound names (to schema)
+# 3. Converting resolved expressions to PA ones
+# Need to break up the abstraction provided by ExpressionEvaluator.
+
+ScalarType = TypeVar("ScalarType")
+
+
+class ExpressionEvaluator:
+    @staticmethod
+    def get_filters(expression: str) -> ds.Expression:
+        """Parse and evaluate the expression to generate a filter condition.
+
+        Args:
+            expression: A string representing the filter expression to parse.
+
+        Returns:
+            A PyArrow compute expression for filtering data.
+
+        """
+        try:
+            tree = ast.parse(expression, mode="eval")
+            return _ConvertToArrowExpressionVisitor().visit(tree.body)
+        except SyntaxError as e:
+            raise ValueError(f"Invalid syntax in the expression: {expression}") from e
+        except Exception as e:
+            logger.exception(f"Error processing expression: {e}")
+            raise
+
+    @staticmethod
+    def parse_native_expression(expression: str) -> "Expr":
+        """Parse and evaluate the expression to generate a Ray Data expression.
+
+        Args:
+            expression: A string representing the filter expression to parse.
+
+        Returns:
+            A Ray Data Expr object for filtering data.
+
+        """
+        try:
+            tree = ast.parse(expression, mode="eval")
+            return _ConvertToNativeExpressionVisitor().visit(tree.body)
+        except SyntaxError as e:
+            raise ValueError(f"Invalid syntax in the expression: {expression}") from e
+        except Exception as e:
+            logger.exception(f"Error processing expression: {e}")
+            raise
+
+
+class _ConvertToArrowExpressionVisitor(ast.NodeVisitor):
+    # TODO: Deprecate this visitor after we remove string support in filter API.
+    def visit_Compare(self, node: ast.Compare) -> ds.Expression:
+        """Handle comparison operations (e.g., a == b, a < b, a in b).
+
+        Args:
+            node: The AST node representing a comparison operation.
+
+        Returns:
+            An expression representing the comparison.
+        """
+        # Handle left operand
+        # TODO Validate columns
+        if isinstance(node.left, ast.Attribute):
+            # Visit and handle attributes
+            left_expr = self.visit(node.left)
+        elif isinstance(node.left, ast.Name):
+            # Treat as a simple field
+            left_expr = self.visit(node.left)
+        elif isinstance(node.left, ast.Constant):
+            # Constant values are used directly
+            left_expr = node.left.value
+        else:
+            raise ValueError(f"Unsupported left operand type: {type(node.left)}")
+
+        comparators = [self.visit(comp) for comp in node.comparators]
+
+        op = node.ops[0]
+        if isinstance(op, ast.In):
+            return pc.is_in(left_expr, comparators[0])
+        elif isinstance(op, ast.NotIn):
+            return ~pc.is_in(left_expr, comparators[0])
+        elif isinstance(op, ast.Eq):
+            return left_expr == comparators[0]
+        elif isinstance(op, ast.NotEq):
+            return left_expr != comparators[0]
+        elif isinstance(op, ast.Lt):
+            return left_expr < comparators[0]
+        elif isinstance(op, ast.LtE):
+            return left_expr <= comparators[0]
+        elif isinstance(op, ast.Gt):
+            return left_expr > comparators[0]
+        elif isinstance(op, ast.GtE):
+            return left_expr >= comparators[0]
+        else:
+            raise ValueError(f"Unsupported operator type: {op}")
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> ds.Expression:
+        """Handle logical operations (e.g., a and b, a or b).
+
+        Args:
+            node: The AST node representing a boolean operation.
+
+        Returns:
+            An expression representing the logical operation.
+        """
+        conditions = [self.visit(value) for value in node.values]
+        combined_expr = conditions[0]
+
+        for condition in conditions[1:]:
+            if isinstance(node.op, ast.And):
+                # Combine conditions with logical AND
+                combined_expr &= condition
+            elif isinstance(node.op, ast.Or):
+                # Combine conditions with logical OR
+                combined_expr |= condition
+            else:
+                raise ValueError(
+                    f"Unsupported logical operator: {type(node.op).__name__}"
+                )
+
+        return combined_expr
+
+    def visit_Name(self, node: ast.Name) -> ds.Expression:
+        """Handle variable (name) nodes and return them as pa.dataset.Expression.
+
+        Even if the name contains periods, it's treated as a single string.
+
+        Args:
+            node: The AST node representing a variable.
+
+        Returns:
+            The variable wrapped as a pa.dataset.Expression.
+        """
+        # Directly use the field name as a string (even if it contains periods)
+        field_name = node.id
+        return pc.field(field_name)
+
+    def visit_Attribute(self, node: ast.Attribute) -> object:
+        """Handle attribute access (e.g., np.nan).
+
+        Args:
+            node: The AST node representing an attribute access.
+
+        Returns:
+            object: The attribute value.
+
+        Raises:
+            ValueError: If the attribute is unsupported.
+        """
+        # Recursively visit the left side (base object or previous attribute)
+        if isinstance(node.value, ast.Attribute):
+            # If the value is an attribute, recursively resolve it
+            left_expr = self.visit(node.value)
+            return pc.field(f"{left_expr}.{node.attr}")
+
+        elif isinstance(node.value, ast.Name):
+            # If the value is a name (e.g., "foo"), we can directly return the field
+            left_name = node.value.id  # The base name, e.g., "foo"
+            return pc.field(f"{left_name}.{node.attr}")
+
+        raise ValueError(f"Unsupported attribute: {node.attr}")
+
+    def visit_List(self, node: ast.List) -> ds.Expression:
+        """Handle list literals.
+
+        Args:
+            node: The AST node representing a list.
+
+        Returns:
+            The list of elements wrapped as a pa.dataset.Expression.
+        """
+        elements = [self.visit(elt) for elt in node.elts]
+        return pa.array(elements)
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> ds.Expression:
+        """Handle case where comparator is UnaryOP (e.g., a == -1).
+        AST for this expression will be Compare(left=Name(id='a'), ops=[Eq()],
+        comparators=[UnaryOp(op=USub(), operand=Constant(value=1))])
+
+        Args:
+            node: The constant value."""
+
+        op = node.op
+        if isinstance(op, ast.USub):
+            return pc.scalar(-node.operand.value)
+        else:
+            raise ValueError(f"Unsupported unary operator: {op}")
+
+    # TODO (srinathk) Note that visit_Constant does not return pa.dataset.Expression
+    # because to support function in() which takes in a List, the elements in the List
+    # needs to values instead of pa.dataset.Expression per pyarrow.dataset.Expression
+    # specification. May be down the road, we can update it as Arrow relaxes this
+    # constraint.
+    def visit_Constant(self, node: ast.Constant) -> object:
+        """Handle constant values (e.g., numbers, strings).
+
+        Args:
+            node: The AST node representing a constant value.
+
+        Returns:
+            object: The constant value itself (e.g., number, string, or boolean).
+        """
+        return node.value  # Return the constant value directly.
+
+    def visit_Call(self, node: ast.Call) -> ds.Expression:
+        """Handle function calls (e.g., is_nan(a), is_valid(b)).
+
+        Args:
+            node: The AST node representing a function call.
+
+        Returns:
+            The corresponding expression based on the function called.
+
+        Raises:
+            ValueError: If the function is unsupported or has incorrect arguments.
+        """
+        func_name = node.func.id
+        function_map = {
+            "is_nan": lambda arg: arg.is_nan(),
+            "is_null": lambda arg, nan_is_null=False: arg.is_null(
+                nan_is_null=nan_is_null
+            ),
+            "is_valid": lambda arg: arg.is_valid(),
+            "is_in": lambda arg1, arg2: pc.is_in(arg1, arg2),
+        }
+
+        if func_name in function_map:
+            # Visit all arguments of the function call
+            args = [self.visit(arg) for arg in node.args]
+            # Handle the "is_null" function with one or two arguments
+            if func_name == "is_null":
+                if len(args) == 1:
+                    return function_map[func_name](args[0])
+                elif len(args) == 2:
+                    return function_map[func_name](args[0], args[1])
+                else:
+                    raise ValueError("is_null function requires one or two arguments.")
+            # Handle the "is_in" function with exactly two arguments
+            elif func_name == "is_in" and len(args) != 2:
+                raise ValueError("is_in function requires two arguments.")
+            # Ensure the function has one argument (for functions like is_valid)
+            elif func_name != "is_in" and len(args) != 1:
+                raise ValueError(f"{func_name} function requires exactly one argument.")
+            # Call the corresponding function with the arguments
+            return function_map[func_name](*args)
+        else:
+            raise ValueError(f"Unsupported function: {func_name}")
+
+
+class _ConvertToNativeExpressionVisitor(ast.NodeVisitor):
+    """AST visitor that converts string expressions to Ray Data expressions."""
+
+    def visit_Compare(self, node: ast.Compare) -> "Expr":
+        """Handle comparison operations (e.g., a == b, a < b, a in b)."""
+        from ray.data.expressions import BinaryExpr, Operation
+
+        if len(node.ops) != 1 or len(node.comparators) != 1:
+            raise ValueError("Only simple binary comparisons are supported")
+
+        left = self.visit(node.left)
+        right = self.visit(node.comparators[0])
+        op = node.ops[0]
+
+        # Map AST comparison operators to Ray Data operations
+        op_map = {
+            ast.Eq: Operation.EQ,
+            ast.NotEq: Operation.NE,
+            ast.Lt: Operation.LT,
+            ast.LtE: Operation.LE,
+            ast.Gt: Operation.GT,
+            ast.GtE: Operation.GE,
+            ast.In: Operation.IN,
+            ast.NotIn: Operation.NOT_IN,
+        }
+
+        if type(op) not in op_map:
+            raise ValueError(f"Unsupported comparison operator: {type(op).__name__}")
+
+        return BinaryExpr(op_map[type(op)], left, right)
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> "Expr":
+        """Handle logical operations (e.g., a and b, a or b)."""
+        from ray.data.expressions import BinaryExpr, Operation
+
+        conditions = [self.visit(value) for value in node.values]
+        combined_expr = conditions[0]
+
+        for condition in conditions[1:]:
+            if isinstance(node.op, ast.And):
+                combined_expr = BinaryExpr(Operation.AND, combined_expr, condition)
+            elif isinstance(node.op, ast.Or):
+                combined_expr = BinaryExpr(Operation.OR, combined_expr, condition)
+            else:
+                raise ValueError(
+                    f"Unsupported logical operator: {type(node.op).__name__}"
+                )
+
+        return combined_expr
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> "Expr":
+        """Handle unary operations (e.g., not a, -5)."""
+        from ray.data.expressions import Operation, UnaryExpr, lit
+
+        if isinstance(node.op, ast.Not):
+            operand = self.visit(node.operand)
+            return UnaryExpr(Operation.NOT, operand)
+        elif isinstance(node.op, ast.USub):
+            operand = self.visit(node.operand)
+            return operand * lit(-1)
+        else:
+            raise ValueError(f"Unsupported unary operator: {type(node.op).__name__}")
+
+    def visit_Name(self, node: ast.Name) -> "Expr":
+        """Handle variable names (column references)."""
+        from ray.data.expressions import col
+
+        return col(node.id)
+
+    def visit_Constant(self, node: ast.Constant) -> "Expr":
+        """Handle constant values (numbers, strings, booleans)."""
+        from ray.data.expressions import lit
+
+        return lit(node.value)
+
+    def visit_List(self, node: ast.List) -> "Expr":
+        """Handle list literals."""
+        from ray.data.expressions import LiteralExpr, lit
+
+        # Visit all elements first
+        visited_elements = [self.visit(elt) for elt in node.elts]
+
+        # Try to extract constant values for literal list
+        elements = []
+        for elem in visited_elements:
+            if isinstance(elem, LiteralExpr):
+                elements.append(elem.value)
+            else:
+                # For compatibility with Arrow visitor, we need to support non-literals
+                # but Ray Data expressions may have limitations here
+                raise ValueError(
+                    "List contains non-constant expressions. Ray Data expressions "
+                    "currently only support lists of constant values."
+                )
+
+        return lit(elements)
+
+    def visit_Attribute(self, node: ast.Attribute) -> "Expr":
+        """Handle attribute access (e.g., for nested column names)."""
+        from ray.data.expressions import col
+
+        # For nested column names like "user.age", combine them with dots
+        if isinstance(node.value, ast.Name):
+            return col(f"{node.value.id}.{node.attr}")
+        elif isinstance(node.value, ast.Attribute):
+            # Recursively handle nested attributes
+            left_expr = self.visit(node.value)
+            if isinstance(left_expr, ColumnExpr):
+                return col(f"{left_expr._name}.{node.attr}")
+
+        raise ValueError(
+            f"Unsupported attribute access: {node.attr}. Node details: {ast.dump(node)}"
+        )
+
+    def visit_Call(self, node: ast.Call) -> "Expr":
+        """Handle function calls for operations like is_null, is_not_null, is_nan."""
+        from ray.data.expressions import BinaryExpr, Operation, UnaryExpr
+
+        func_name = node.func.id if isinstance(node.func, ast.Name) else str(node.func)
+
+        if func_name == "is_null":
+            if len(node.args) != 1:
+                raise ValueError("is_null() expects exactly one argument")
+            operand = self.visit(node.args[0])
+            return UnaryExpr(Operation.IS_NULL, operand)
+        # Adding this conditional to keep it consistent with the current implementation,
+        # of carrying Pyarrow's semantic of `is_valid`
+        elif func_name == "is_valid" or func_name == "is_not_null":
+            if len(node.args) != 1:
+                raise ValueError(f"{func_name}() expects exactly one argument")
+            operand = self.visit(node.args[0])
+            return UnaryExpr(Operation.IS_NOT_NULL, operand)
+        elif func_name == "is_nan":
+            if len(node.args) != 1:
+                raise ValueError("is_nan() expects exactly one argument")
+            operand = self.visit(node.args[0])
+            # Use x != x pattern for NaN detection (NaN != NaN is True)
+            return BinaryExpr(Operation.NE, operand, operand)
+        elif func_name == "is_in":
+            if len(node.args) != 2:
+                raise ValueError("is_in() expects exactly two arguments")
+            left = self.visit(node.args[0])
+            right = self.visit(node.args[1])
+            return BinaryExpr(Operation.IN, left, right)
+        else:
+            raise ValueError(f"Unsupported function: {func_name}")
+
+
+class NativeExpressionEvaluator(_ExprVisitor[Union[BlockColumn, ScalarType]]):
+    """Visitor-based expression evaluator that uses Block and BlockColumns
+
+    This evaluator implements the visitor pattern to traverse expression trees
+    and evaluate them against Block data structures. It maintains operation
+    mappings in shared state and returns consistent BlockColumn types.
+    """
+
+    def __init__(self, block: Block):
+        """Initialize the evaluator with a block and operation mappings.
+
+        Args:
+            block: The Block to evaluate expressions against.
+        """
+        self.block = block
+        self.block_accessor = BlockAccessor.for_block(block)
+
+        # Use BlockAccessor to determine operation mappings
+        block_type = self.block_accessor.block_type()
+        if block_type == BlockType.PANDAS:
+            self.ops = _PANDAS_EXPR_OPS_MAP
+        elif block_type == BlockType.ARROW:
+            self.ops = _ARROW_EXPR_OPS_MAP
+        else:
+            raise TypeError(f"Unsupported block type: {block_type}")
+
+    def visit_column(self, expr: ColumnExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a column expression and return the column data.
+
+        Args:
+            expr: The column expression.
+
+        Returns:
+            The column data as a BlockColumn.
+        """
+        return self.block[expr.name]
+
+    def visit_literal(self, expr: LiteralExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a literal expression and return the literal value.
+
+        Args:
+            expr: The literal expression.
+
+        Returns:
+            The literal value.
+        """
+        # Given that expressions support pandas blocks, we need to return the value as is.
+        # Pandas has multiple dtype_backends, so there's no guarantee on the return type.
+        return expr.value
+
+    def visit_binary(self, expr: BinaryExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a binary expression and return the result of the operation.
+
+        Args:
+            expr: The binary expression.
+
+        Returns:
+            The result of the binary operation as a BlockColumn.
+        """
+        left_result = self.visit(expr.left)
+        right_result = self.visit(expr.right)
+
+        return self.ops[expr.op](left_result, right_result)
+
+    def visit_unary(self, expr: UnaryExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a unary expression and return the result of the operation.
+
+        Args:
+            expr: The unary expression.
+
+        Returns:
+            The result of the unary operation as a BlockColumn.
+        """
+        operand_result = self.visit(expr.operand)
+        return self.ops[expr.op](operand_result)
+
+    def visit_udf(self, expr: UDFExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a UDF expression and return the result of the function call.
+
+        Args:
+            expr: The UDF expression.
+
+        Returns:
+            The result of the UDF call as a BlockColumn.
+        """
+        args = [self.visit(arg) for arg in expr.args]
+        kwargs = {k: self.visit(v) for k, v in expr.kwargs.items()}
+        result = expr.fn(*args, **kwargs)
+
+        if not isinstance(result, (pd.Series, np.ndarray, pa.Array, pa.ChunkedArray)):
+            function_name = expr.fn.__name__
+            raise TypeError(
+                f"UDF '{function_name}' returned invalid type {type(result).__name__}. "
+                f"Expected type (pandas.Series, numpy.ndarray, pyarrow.Array, or pyarrow.ChunkedArray)"
+            )
+
+        return result
+
+    def visit_alias(self, expr: AliasExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit an alias expression and return the renamed result.
+
+        Args:
+            expr: The alias expression.
+
+        Returns:
+            A Block with the data from the inner expression.
+        """
+        # Evaluate the inner expression
+        return self.visit(expr.expr)
+
+    def visit_star(self, expr: StarExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a star expression.
+
+        Args:
+            expr: The star expression.
+
+        Returns:
+            TypeError: StarExpr cannot be evaluated as a regular expression.
+        """
+        # star() should not be evaluated directly - it's handled at Project level
+        raise TypeError(
+            "StarExpr cannot be evaluated as a regular expression. "
+            "It should only be used in Project operations."
+        )
+
+    def visit_download(self, expr: DownloadExpr) -> Union[BlockColumn, ScalarType]:
+        """Visit a download expression.
+
+        Args:
+            expr: The download expression.
+
+        Returns:
+            TypeError: DownloadExpr evaluation not yet implemented.
+        """
+        raise TypeError(
+            "DownloadExpr evaluation is not yet implemented in NativeExpressionEvaluator."
+        )
+
+
+def eval_expr(expr: Expr, block: Block) -> Union[BlockColumn, ScalarType]:
+    """Evaluate an expression against a block using the visitor pattern.
+
+    Args:
+        expr: The expression to evaluate.
+        block: The Block to evaluate against.
+
+    Returns:
+        The evaluated result as a BlockColumn or a scalar value.
+    """
+    evaluator = NativeExpressionEvaluator(block)
+    return evaluator.visit(expr)
+
+
+def eval_projection(projection_exprs: List[Expr], block: Block) -> Block:
+    """
+    Evaluate a projection (list of expressions) against a block.
+
+    Handles projection semantics including:
+    - Empty projections
+    - Star() expressions for preserving existing columns
+    - Rename detection
+    - Column ordering
+
+    Args:
+        projection_exprs: List of expressions to evaluate (may include StarExpr)
+        block: The block to project
+
+    Returns:
+        A new block with the projected schema
+    """
+    block_accessor = BlockAccessor.for_block(block)
+
+    # Skip projection only for schema-less empty blocks.
+    if block_accessor.num_rows() == 0 and len(block_accessor.column_names()) == 0:
+        return block
+
+    # Handle simple cases early.
+    if len(projection_exprs) == 0:
+        return block_accessor.select([])
+
+    input_column_names = list(block_accessor.column_names())
+    # Collect input column rename map from the projection list
+    input_column_rename_map = _extract_input_columns_renaming_mapping(projection_exprs)
+
+    # Expand star expr (if any)
+    if isinstance(projection_exprs[0], StarExpr):
+        # Cherry-pick input block's columns that aren't explicitly removed via
+        # renaming
+        input_column_ref_exprs = [
+            col(c) for c in input_column_names if c not in input_column_rename_map
+        ]
+
+        projection_exprs = input_column_ref_exprs + projection_exprs[1:]
+
+    names, output_cols = zip(*[(e.name, eval_expr(e, block)) for e in projection_exprs])
+
+    # This clumsy workaround is necessary to be able to fill in Pyarrow tables
+    # that has to be "seeded" from existing table with N rows, and couldn't be
+    # started from a truly empty table.
+    #
+    # TODO fix
+    new_block = BlockAccessor.for_block(block).fill_column("__stub__", None)
+    new_block = BlockAccessor.for_block(new_block).drop(input_column_names)
+
+    for name, output_col in zip(names, output_cols):
+        new_block = BlockAccessor.for_block(new_block).fill_column(name, output_col)
+
+    return BlockAccessor.for_block(new_block).drop(["__stub__"])
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/expression_visitors.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/expression_visitors.py
new file mode 100644
index 0000000000000000000000000000000000000000..b01ed2e87710b2a5a73835daef28419af85c2cd5
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_expression/expression_visitors.py
@@ -0,0 +1,349 @@
+from dataclasses import replace
+from typing import Dict, List, TypeVar
+
+from ray.data.expressions import (
+    AliasExpr,
+    BinaryExpr,
+    ColumnExpr,
+    DownloadExpr,
+    Expr,
+    LiteralExpr,
+    StarExpr,
+    UDFExpr,
+    UnaryExpr,
+    _ExprVisitor,
+)
+
+T = TypeVar("T")
+
+
+class _ExprVisitorBase(_ExprVisitor[None]):
+    """Base visitor that provides automatic recursive traversal.
+
+    This class extends _ExprVisitor and provides default implementations
+    for composite nodes that automatically traverse child expressions.
+    """
+
+    def visit_binary(self, expr: "BinaryExpr") -> None:
+        """Default implementation: recursively visit both operands."""
+        super().visit(expr.left)
+        super().visit(expr.right)
+
+    def visit_unary(self, expr: "UnaryExpr") -> None:
+        """Default implementation: recursively visit the operand."""
+        super().visit(expr.operand)
+
+    def visit_alias(self, expr: "AliasExpr") -> None:
+        """Default implementation: recursively visit the inner expression."""
+        super().visit(expr.expr)
+
+    def visit_udf(self, expr: "UDFExpr") -> None:
+        """Default implementation: recursively visit all arguments."""
+        for arg in expr.args:
+            super().visit(arg)
+        for value in expr.kwargs.values():
+            super().visit(value)
+
+    def visit_literal(self, expr: LiteralExpr) -> None:
+        """Visit a literal expression (no columns to collect)."""
+        pass
+
+    def visit_star(self, expr: StarExpr) -> None:
+        """Visit a star expression (no columns to collect)."""
+        pass
+
+    def visit_download(self, expr: "Expr") -> None:
+        """Visit a download expression (no columns to collect)."""
+        pass
+
+
+class _ColumnReferenceCollector(_ExprVisitorBase):
+    """Visitor that collects all column references from expression trees.
+
+    This visitor traverses expression trees and accumulates column names
+    referenced in ColumnExpr nodes.
+    """
+
+    def __init__(self):
+        """Initialize with an empty set of referenced columns."""
+
+        # NOTE: We're using dict to maintain insertion ordering
+        self._col_refs: Dict[str, None] = dict()
+
+    def get_column_refs(self) -> List[str]:
+        return list(self._col_refs.keys())
+
+    def visit_column(self, expr: ColumnExpr) -> None:
+        """Visit a column expression and collect its name.
+
+        Args:
+            expr: The column expression.
+
+        Returns:
+            None (only collects columns as a side effect).
+        """
+        self._col_refs[expr.name] = None
+
+    def visit_alias(self, expr: AliasExpr) -> None:
+        """Visit an alias expression and collect from its inner expression.
+
+        Args:
+            expr: The alias expression.
+
+        Returns:
+            None (only collects columns as a side effect).
+        """
+        self.visit(expr.expr)
+
+
+class _ColumnSubstitutionVisitor(_ExprVisitor[Expr]):
+    """Visitor rebinding column references in ``Expression``s.
+
+    This visitor traverses given ``Expression`` trees and substitutes column references
+    according to a provided substitution map.
+    """
+
+    def __init__(self, column_ref_substitutions: Dict[str, Expr]):
+        """Initialize with a column substitution map.
+
+        Args:
+            column_ref_substitutions: Mapping from column names to replacement expressions.
+        """
+        self._col_ref_substitutions = column_ref_substitutions
+
+    def visit_column(self, expr: ColumnExpr) -> Expr:
+        """Visit a column expression and substitute it.
+
+        Args:
+            expr: The column expression.
+
+        Returns:
+            The substituted expression or the original if no substitution exists.
+        """
+        substitution = self._col_ref_substitutions.get(expr.name)
+
+        return substitution if substitution is not None else expr
+
+    def visit_literal(self, expr: LiteralExpr) -> Expr:
+        """Visit a literal expression (no rewriting needed).
+
+        Args:
+            expr: The literal expression.
+
+        Returns:
+            The original literal expression.
+        """
+        return expr
+
+    def visit_binary(self, expr: BinaryExpr) -> Expr:
+        """Visit a binary expression and rewrite its operands.
+
+        Args:
+            expr: The binary expression.
+
+        Returns:
+            A new binary expression with rewritten operands.
+        """
+        return BinaryExpr(
+            expr.op,
+            self.visit(expr.left),
+            self.visit(expr.right),
+        )
+
+    def visit_unary(self, expr: UnaryExpr) -> Expr:
+        """Visit a unary expression and rewrite its operand.
+
+        Args:
+            expr: The unary expression.
+
+        Returns:
+            A new unary expression with rewritten operand.
+        """
+        return UnaryExpr(expr.op, self.visit(expr.operand))
+
+    def visit_udf(self, expr: UDFExpr) -> Expr:
+        """Visit a UDF expression and rewrite its arguments.
+
+        Args:
+            expr: The UDF expression.
+
+        Returns:
+            A new UDF expression with rewritten arguments.
+        """
+        new_args = [self.visit(arg) for arg in expr.args]
+        new_kwargs = {key: self.visit(value) for key, value in expr.kwargs.items()}
+        return UDFExpr(
+            fn=expr.fn, data_type=expr.data_type, args=new_args, kwargs=new_kwargs
+        )
+
+    def visit_alias(self, expr: AliasExpr) -> Expr:
+        """Visit an alias expression and rewrite its inner expression.
+
+        Args:
+            expr: The alias expression.
+
+        Returns:
+            A new alias expression with rewritten inner expression and preserved name.
+        """
+        # We unalias returned expression to avoid nested aliasing
+        visited = self.visit(expr.expr)._unalias()
+        # NOTE: We're carrying over all of the other aspects of the alias
+        #       only replacing inner expre
+        return replace(
+            expr,
+            expr=visited,
+            # Alias expression will remain a renaming one (ie replacing source column)
+            # so long as it's referencing another column (and not otherwise)
+            #
+            # TODO replace w/ standalone rename expr
+            _is_rename=expr._is_rename and _is_col_expr(visited),
+        )
+
+    def visit_download(self, expr: "Expr") -> Expr:
+        """Visit a download expression (no rewriting needed).
+
+        Args:
+            expr: The download expression.
+
+        Returns:
+            The original download expression.
+        """
+        return expr
+
+    def visit_star(self, expr: StarExpr) -> Expr:
+        """Visit a star expression (no rewriting needed).
+
+        Args:
+            expr: The star expression.
+
+        Returns:
+            The original star expression.
+        """
+        return expr
+
+
+def _is_col_expr(expr: Expr) -> bool:
+    return isinstance(expr, ColumnExpr) or (
+        isinstance(expr, AliasExpr) and isinstance(expr.expr, ColumnExpr)
+    )
+
+
+class _TreeReprVisitor(_ExprVisitor[str]):
+    """Visitor that generates a readable tree representation of expressions. Returns in pre-order traversal."""
+
+    def __init__(self, prefix: str = "", is_last: bool = True):
+        """
+        Initialize the tree representation visitor.
+
+        Args:
+            prefix: The prefix string for indentation (accumulated from parent nodes)
+            is_last: Whether this node is the last child of its parent
+        """
+        self.prefix = prefix
+        self.is_last = is_last
+        self._max_length = 50  # Maximum length of the node label
+
+    def _make_tree_lines(
+        self,
+        node_label: str,
+        children: List[tuple[str, "Expr"]] = None,
+        expr: "Expr" = None,
+    ) -> str:
+        """
+        Format a node and its children with tree box-drawing characters.
+
+        Args:
+            node_label: The label for this node (e.g., "ADD")
+            children: List of (label, child_expr) tuples to render as children
+            expr: The expression node (used to extract datatype)
+
+        Returns:
+            Multi-line string representation of the tree
+        """
+        lines = [node_label]
+
+        if children:
+            for i, (label, child_expr) in enumerate(children):
+                is_last_child = i == len(children) - 1
+
+                # Build prefix for the child based on whether current node is last
+                child_prefix = self.prefix + ("    " if self.is_last else "│   ")
+
+                # Choose connector: └── for last child, ├── for others
+                connector = "└── " if is_last_child else "├── "
+
+                # Recursively visit the child with updated prefix
+                child_visitor = _TreeReprVisitor(child_prefix, is_last_child)
+                child_lines = child_visitor.visit(child_expr).split("\n")
+
+                # Add the first line with label and connector
+                if label:
+                    lines.append(f"{child_prefix}{connector}{label}: {child_lines[0]}")
+                else:
+                    lines.append(f"{child_prefix}{connector}{child_lines[0]}")
+
+                # Add remaining lines from child with proper indentation
+                for line in child_lines[1:]:
+                    lines.append(line)
+
+        return "\n".join(lines)
+
+    def visit_column(self, expr: "ColumnExpr") -> str:
+        return self._make_tree_lines(f"COL({expr.name!r})", expr=expr)
+
+    def visit_literal(self, expr: "LiteralExpr") -> str:
+        # Truncate long values for readability
+        value_repr = repr(expr.value)
+        if len(value_repr) > self._max_length:
+            value_repr = value_repr[: self._max_length - 3] + "..."
+        return self._make_tree_lines(f"LIT({value_repr})", expr=expr)
+
+    def visit_binary(self, expr: "BinaryExpr") -> str:
+        return self._make_tree_lines(
+            f"{expr.op.name}",
+            children=[
+                ("left", expr.left),
+                ("right", expr.right),
+            ],
+            expr=expr,
+        )
+
+    def visit_unary(self, expr: "UnaryExpr") -> str:
+        return self._make_tree_lines(
+            f"{expr.op.name}",
+            children=[("operand", expr.operand)],
+            expr=expr,
+        )
+
+    def visit_alias(self, expr: "AliasExpr") -> str:
+        rename_marker = " [rename]" if expr._is_rename else ""
+        return self._make_tree_lines(
+            f"ALIAS({expr.name!r}){rename_marker}",
+            children=[("", expr.expr)],
+            expr=expr,
+        )
+
+    def visit_udf(self, expr: "UDFExpr") -> str:
+        # Get function name for better readability
+        fn_name = getattr(expr.fn, "__name__", str(expr.fn))
+
+        children = []
+        # Add positional arguments
+        for i, arg in enumerate(expr.args):
+            children.append((f"arg[{i}]", arg))
+
+        # Add keyword arguments
+        for key, value in expr.kwargs.items():
+            children.append((f"kwarg[{key!r}]", value))
+
+        return self._make_tree_lines(
+            f"UDF({fn_name})",
+            children=children if children else None,
+            expr=expr,
+        )
+
+    def visit_download(self, expr: "DownloadExpr") -> str:
+        return self._make_tree_lines(f"DOWNLOAD({expr.uri_column_name!r})", expr=expr)
+
+    def visit_star(self, expr: "StarExpr") -> str:
+        return self._make_tree_lines("COL(*)", expr=expr)
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_read_op.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_read_op.py
new file mode 100644
index 0000000000000000000000000000000000000000..149c95aa8bb49b81d2effef4038f396e34dfd903
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_read_op.py
@@ -0,0 +1,129 @@
+import logging
+import warnings
+from typing import Iterable, List
+
+import ray
+from ray import ObjectRef
+from ray.data._internal.compute import TaskPoolStrategy
+from ray.data._internal.execution.interfaces import PhysicalOperator, RefBundle
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.execution.util import memory_string
+from ray.data._internal.logical.operators.read_operator import Read
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.util import _warn_on_high_parallelism
+from ray.data.block import Block, BlockMetadata
+from ray.data.context import DataContext
+from ray.data.datasource.datasource import ReadTask
+from ray.experimental.locations import get_local_object_locations
+from ray.util.debug import log_once
+
+TASK_SIZE_WARN_THRESHOLD_BYTES = 1024 * 1024  # 1 MiB
+
+logger = logging.getLogger(__name__)
+
+
+def _derive_metadata(read_task: ReadTask, read_task_ref: ObjectRef) -> BlockMetadata:
+    # NOTE: Use the `get_local_object_locations` API to get the size of the
+    # serialized ReadTask, instead of pickling.
+    # Because the ReadTask may capture ObjectRef objects, which cannot
+    # be serialized out-of-band.
+    locations = get_local_object_locations([read_task_ref])
+    task_size = locations[read_task_ref]["object_size"]
+    if task_size > TASK_SIZE_WARN_THRESHOLD_BYTES and log_once(
+        f"large_read_task_{read_task.read_fn.__name__}"
+    ):
+        warnings.warn(
+            "The serialized size of your read function named "
+            f"'{read_task.read_fn.__name__}' is {memory_string(task_size)}. This size "
+            "is relatively large. As a result, Ray might excessively "
+            "spill objects during execution. To fix this issue, avoid accessing "
+            f"`self` or other large objects in '{read_task.read_fn.__name__}'."
+        )
+
+    return BlockMetadata(
+        num_rows=1,
+        size_bytes=task_size,
+        exec_stats=None,
+        input_files=None,
+    )
+
+
+def plan_read_op(
+    op: Read,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Get the corresponding DAG of physical operators for Read.
+
+    Note this method only converts the given `op`, but not its input dependencies.
+    See Planner.plan() for more details.
+    """
+    assert len(physical_children) == 0
+
+    def get_input_data(target_max_block_size) -> List[RefBundle]:
+        parallelism = op.get_detected_parallelism()
+        assert (
+            parallelism is not None
+        ), "Read parallelism must be set by the optimizer before execution"
+
+        # Get the original read tasks
+        read_tasks = op._datasource_or_legacy_reader.get_read_tasks(
+            parallelism, per_task_row_limit=op._per_block_limit
+        )
+
+        _warn_on_high_parallelism(parallelism, len(read_tasks))
+
+        ret = []
+        for read_task in read_tasks:
+            read_task_ref = ray.put(read_task)
+            ref_bundle = RefBundle(
+                (
+                    (
+                        # TODO: figure out a better way to pass read
+                        # tasks other than ray.put().
+                        read_task_ref,
+                        _derive_metadata(read_task, read_task_ref),
+                    ),
+                ),
+                # `owns_blocks` is False, because these refs are the root of the
+                # DAG. We shouldn't eagerly free them. Otherwise, the DAG cannot
+                # be reconstructed.
+                owns_blocks=False,
+                schema=None,
+            )
+            ret.append(ref_bundle)
+        return ret
+
+    inputs = InputDataBuffer(data_context, input_data_factory=get_input_data)
+
+    def do_read(blocks: Iterable[ReadTask], _: TaskContext) -> Iterable[Block]:
+        for read_task in blocks:
+            yield from read_task()
+
+    # Create a MapTransformer for a read operator
+    map_transformer = MapTransformer(
+        [
+            BlockMapTransformFn(
+                do_read,
+                is_udf=False,
+                output_block_size_option=OutputBlockSizeOption.of(
+                    target_max_block_size=data_context.target_max_block_size,
+                ),
+            ),
+        ]
+    )
+
+    return MapOperator.create(
+        map_transformer,
+        inputs,
+        data_context,
+        name=op.name,
+        compute_strategy=TaskPoolStrategy(op._concurrency),
+        ray_remote_args=op._ray_remote_args,
+    )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_udf_map_op.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_udf_map_op.py
new file mode 100644
index 0000000000000000000000000000000000000000..e95580571a06e999032b8d16ab57060a548e9972
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_udf_map_op.py
@@ -0,0 +1,794 @@
+import asyncio
+import collections
+import inspect
+import logging
+import queue
+from threading import Thread
+from types import GeneratorType
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+)
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+
+import ray
+from ray._common.utils import get_or_create_event_loop
+from ray._private.ray_constants import env_integer
+from ray.data._internal.compute import get_compute
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BatchMapTransformFn,
+    BlockMapTransformFn,
+    MapTransformCallable,
+    MapTransformer,
+    Row,
+    RowMapTransformFn,
+)
+from ray.data._internal.execution.util import make_callable_class_concurrent
+from ray.data._internal.logical.operators.map_operator import (
+    AbstractUDFMap,
+    Filter,
+    FlatMap,
+    MapBatches,
+    MapRows,
+    Project,
+    StreamingRepartition,
+)
+from ray.data._internal.numpy_support import _is_valid_column_values
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.streaming_repartition import StreamingRepartitionRefBundler
+from ray.data._internal.util import _truncated_repr
+from ray.data.block import (
+    Block,
+    BlockAccessor,
+    CallableClass,
+    DataBatch,
+    UserDefinedFunction,
+)
+from ray.data.context import DataContext
+from ray.data.exceptions import UserCodeException
+from ray.util.rpdb import _is_ray_debugger_post_mortem_enabled
+
+logger = logging.getLogger(__name__)
+
+
+# Controls default max-concurrency setting for async row-based UDFs
+DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY = env_integer(
+    "RAY_DATA_DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY", 16
+)
+
+# Controls default max-concurrency setting for async batch-based UDFs
+DEFAULT_ASYNC_BATCH_UDF_MAX_CONCURRENCY = env_integer(
+    "RAY_DATA_DEFAULT_ASYNC_BATCH_UDF_MAX_CONCURRENCY", 4
+)
+
+
+class _MapActorContext:
+    def __init__(
+        self,
+        udf_map_cls: UserDefinedFunction,
+        udf_map_fn: Callable[[Any], Any],
+        is_async: bool,
+    ):
+        self.udf_map_cls = udf_map_cls
+        self.udf_map_fn = udf_map_fn
+        self.is_async = is_async
+        self.udf_map_asyncio_loop = None
+        self.udf_map_asyncio_thread = None
+
+        if is_async:
+            self._init_async()
+
+    def _init_async(self):
+        # Only used for callable class with async generator `__call__` method.
+        loop = get_or_create_event_loop()
+
+        def run_loop():
+            asyncio.set_event_loop(loop)
+            loop.run_forever()
+
+        thread = Thread(target=run_loop)
+        thread.start()
+        self.udf_map_asyncio_loop = loop
+        self.udf_map_asyncio_thread = thread
+
+
+def plan_project_op(
+    op: Project,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    # Extract op.exprs before defining the closure to prevent cloudpickle from
+    # serializing the entire op object (which may contain references to non-serializable
+    # datasources with weak references, e.g., PyIceberg tables)
+    projection_exprs = op.exprs
+
+    def _project_block(block: Block) -> Block:
+        try:
+            from ray.data._internal.planner.plan_expression.expression_evaluator import (
+                eval_projection,
+            )
+
+            return eval_projection(projection_exprs, block)
+        except Exception as e:
+            _try_wrap_udf_exception(e)
+
+    compute = get_compute(op._compute)
+    map_transformer = MapTransformer(
+        [
+            BlockMapTransformFn(
+                _generate_transform_fn_for_map_block(_project_block),
+                disable_block_shaping=(len(op.exprs) == 0),
+            )
+        ]
+    )
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        ray_remote_args=op._ray_remote_args,
+        ray_remote_args_fn=op._ray_remote_args_fn,
+    )
+
+
+def plan_streaming_repartition_op(
+    op: StreamingRepartition,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+    compute = get_compute(op._compute)
+    transform_fn = BlockMapTransformFn(
+        lambda blocks, ctx: blocks,
+        output_block_size_option=OutputBlockSizeOption.of(
+            target_num_rows_per_block=op.target_num_rows_per_block,  # To split n*target_max_block_size row into n blocks
+        ),
+    )
+    map_transformer = MapTransformer([transform_fn])
+
+    # Disable fusion for streaming repartition with the downstream op.
+    operator = MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        ref_bundler=StreamingRepartitionRefBundler(op.target_num_rows_per_block),
+        ray_remote_args=op._ray_remote_args,
+        ray_remote_args_fn=op._ray_remote_args_fn,
+        supports_fusion=False,
+    )
+
+    return operator
+
+
+def plan_filter_op(
+    op: Filter,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    output_block_size_option = OutputBlockSizeOption.of(
+        target_max_block_size=data_context.target_max_block_size,
+    )
+
+    predicate_expr = op._predicate_expr
+    compute = get_compute(op._compute)
+    if predicate_expr is not None:
+
+        def filter_block_fn(
+            blocks: Iterable[Block], ctx: TaskContext
+        ) -> Iterable[Block]:
+            for block in blocks:
+                block_accessor = BlockAccessor.for_block(block)
+                filtered_block = block_accessor.filter(predicate_expr)
+                yield filtered_block
+
+        init_fn = None
+        transform_fn = BlockMapTransformFn(
+            filter_block_fn,
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+    else:
+        udf_is_callable_class = isinstance(op._fn, CallableClass)
+        filter_fn, init_fn = _get_udf(
+            op._fn,
+            op._fn_args,
+            op._fn_kwargs,
+            op._fn_constructor_args if udf_is_callable_class else None,
+            op._fn_constructor_kwargs if udf_is_callable_class else None,
+        )
+
+        transform_fn = RowMapTransformFn(
+            _generate_transform_fn_for_filter(filter_fn),
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+
+    map_transformer = MapTransformer([transform_fn], init_fn=init_fn)
+
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        ray_remote_args=op._ray_remote_args,
+        ray_remote_args_fn=op._ray_remote_args_fn,
+    )
+
+
+def plan_udf_map_op(
+    op: AbstractUDFMap,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> MapOperator:
+    """Get the corresponding physical operators DAG for AbstractUDFMap operators.
+
+    Note this method only converts the given `op`, but not its input dependencies.
+    See Planner.plan() for more details.
+    """
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    output_block_size_option = OutputBlockSizeOption.of(
+        target_max_block_size=data_context.target_max_block_size,
+    )
+
+    compute = get_compute(op._compute)
+    udf_is_callable_class = isinstance(op._fn, CallableClass)
+    fn, init_fn = _get_udf(
+        op._fn,
+        op._fn_args,
+        op._fn_kwargs,
+        op._fn_constructor_args if udf_is_callable_class else None,
+        op._fn_constructor_kwargs if udf_is_callable_class else None,
+    )
+
+    if isinstance(op, MapBatches):
+        transform_fn = BatchMapTransformFn(
+            _generate_transform_fn_for_map_batches(fn),
+            batch_size=op._batch_size,
+            batch_format=op._batch_format,
+            zero_copy_batch=op._zero_copy_batch,
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+
+    else:
+        if isinstance(op, MapRows):
+            udf_fn = _generate_transform_fn_for_map_rows(fn)
+        elif isinstance(op, FlatMap):
+            udf_fn = _generate_transform_fn_for_flat_map(fn)
+        else:
+            raise ValueError(f"Found unknown logical operator during planning: {op}")
+
+        transform_fn = RowMapTransformFn(
+            udf_fn,
+            is_udf=True,
+            output_block_size_option=output_block_size_option,
+        )
+
+    map_transformer = MapTransformer([transform_fn], init_fn=init_fn)
+
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name=op.name,
+        compute_strategy=compute,
+        min_rows_per_bundle=op._min_rows_per_bundled_input,
+        ray_remote_args_fn=op._ray_remote_args_fn,
+        ray_remote_args=op._ray_remote_args,
+        per_block_limit=op._per_block_limit,
+    )
+
+
+def _get_udf(
+    op_fn: Callable,
+    op_fn_args: Tuple[Any, ...],
+    op_fn_kwargs: Dict[str, Any],
+    op_fn_constructor_args: Optional[Tuple[Any, ...]],
+    op_fn_constructor_kwargs: Optional[Dict[str, Any]],
+):
+    # Note, it's important to define these standalone variables.
+    # So the parsed functions won't need to capture the entire operator, which may not
+    # be serializable.
+    udf = op_fn
+    fn_args = op_fn_args or ()
+    fn_kwargs = op_fn_kwargs or {}
+
+    if isinstance(udf, CallableClass):
+        fn_constructor_args = op_fn_constructor_args or ()
+        fn_constructor_kwargs = op_fn_constructor_kwargs or {}
+
+        is_async_udf = _is_async_udf(udf.__call__)
+
+        if not is_async_udf:
+            # TODO(ak) this constrains concurrency for user UDFs to run in a single
+            #          thread irrespective of max_concurrency. Remove
+            udf = make_callable_class_concurrent(udf)
+
+        def init_fn():
+            if ray.data._map_actor_context is None:
+                ray.data._map_actor_context = _MapActorContext(
+                    udf_map_cls=udf,
+                    udf_map_fn=udf(
+                        *fn_constructor_args,
+                        **fn_constructor_kwargs,
+                    ),
+                    is_async=is_async_udf,
+                )
+
+        if inspect.iscoroutinefunction(udf.__call__):
+
+            async def _wrapped_udf_map_fn(item: Any) -> Any:
+                assert ray.data._map_actor_context is not None
+                assert ray.data._map_actor_context.is_async
+
+                try:
+                    return await ray.data._map_actor_context.udf_map_fn(
+                        item,
+                        *fn_args,
+                        **fn_kwargs,
+                    )
+                except Exception as e:
+                    _try_wrap_udf_exception(e)
+
+        elif inspect.isasyncgenfunction(udf.__call__):
+
+            async def _wrapped_udf_map_fn(item: Any) -> Any:
+                assert ray.data._map_actor_context is not None
+                assert ray.data._map_actor_context.is_async
+
+                try:
+                    gen = ray.data._map_actor_context.udf_map_fn(
+                        item,
+                        *fn_args,
+                        **fn_kwargs,
+                    )
+
+                    async for res in gen:
+                        yield res
+                except Exception as e:
+                    _try_wrap_udf_exception(e, item)
+
+        else:
+            assert isinstance(
+                udf.__call__, Callable
+            ), f"Expected Callable, got {udf.__call__} ({type(udf.__call__)})"
+
+            def _wrapped_udf_map_fn(item: Any) -> Any:
+                assert ray.data._map_actor_context is not None
+                assert not ray.data._map_actor_context.is_async
+                try:
+                    return ray.data._map_actor_context.udf_map_fn(
+                        item,
+                        *fn_args,
+                        **fn_kwargs,
+                    )
+                except Exception as e:
+                    _try_wrap_udf_exception(e)
+
+    else:
+
+        def _wrapped_udf_map_fn(item: Any) -> Any:
+            try:
+                return udf(item, *fn_args, **fn_kwargs)
+            except Exception as e:
+                _try_wrap_udf_exception(e)
+
+        def init_fn():
+            pass
+
+    return _wrapped_udf_map_fn, init_fn
+
+
+def _try_wrap_udf_exception(e: Exception, item: Any = None):
+    """If the Ray Debugger is enabled, keep the full stack trace unmodified
+    so that the debugger can stop at the initial unhandled exception.
+    Otherwise, clear the stack trace to omit noisy internal code path."""
+    ctx = ray.data.DataContext.get_current()
+    if _is_ray_debugger_post_mortem_enabled() or ctx.raise_original_map_exception:
+        raise e
+    else:
+        raise UserCodeException("UDF failed to process a data block.") from e
+
+
+# Following are util functions for converting UDFs to `MapTransformCallable`s.
+
+
+def _validate_batch_output(batch: Block) -> None:
+    if not isinstance(
+        batch,
+        (
+            list,
+            pa.Table,
+            np.ndarray,
+            collections.abc.Mapping,
+            pd.core.frame.DataFrame,
+            dict,
+        ),
+    ):
+        raise ValueError(
+            "The `fn` you passed to `map_batches` returned a value of type "
+            f"{type(batch)}. This isn't allowed -- `map_batches` expects "
+            "`fn` to return a `pandas.DataFrame`, `pyarrow.Table`, "
+            "`numpy.ndarray`, `list`, or `dict[str, numpy.ndarray]`."
+        )
+
+    if isinstance(batch, list):
+        raise ValueError(
+            f"Error validating {_truncated_repr(batch)}: "
+            "Returning a list of objects from `map_batches` is not "
+            "allowed in Ray 2.5. To return Python objects, "
+            "wrap them in a named dict field, e.g., "
+            "return `{'results': objects}` instead of just `objects`."
+        )
+
+    if isinstance(batch, collections.abc.Mapping):
+        for key, value in list(batch.items()):
+            if not _is_valid_column_values(value):
+                raise ValueError(
+                    f"Error validating {_truncated_repr(batch)}: "
+                    "The `fn` you passed to `map_batches` returned a "
+                    f"`dict`. `map_batches` expects all `dict` values "
+                    f"to be `list` or `np.ndarray` type, but the value "
+                    f"corresponding to key {key!r} is of type "
+                    f"{type(value)}. To fix this issue, convert "
+                    f"the {type(value)} to a `np.ndarray`."
+                )
+
+
+def _generate_transform_fn_for_map_batches(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[DataBatch, DataBatch]:
+
+    if _is_async_udf(fn):
+        transform_fn = _generate_transform_fn_for_async_map(
+            fn,
+            _validate_batch_output,
+            max_concurrency=DEFAULT_ASYNC_BATCH_UDF_MAX_CONCURRENCY,
+        )
+
+    else:
+
+        def transform_fn(
+            batches: Iterable[DataBatch], _: TaskContext
+        ) -> Iterable[DataBatch]:
+            for batch in batches:
+                try:
+                    if (
+                        not isinstance(batch, collections.abc.Mapping)
+                        and BlockAccessor.for_block(batch).num_rows() == 0
+                    ):
+                        # For empty input blocks, we directly output them without
+                        # calling the UDF.
+                        # TODO(hchen): This workaround is because some all-to-all
+                        # operators output empty blocks with no schema.
+                        res = [batch]
+                    else:
+                        res = fn(batch)
+                        if not isinstance(res, GeneratorType):
+                            res = [res]
+                except ValueError as e:
+                    read_only_msgs = [
+                        "assignment destination is read-only",
+                        "buffer source array is read-only",
+                    ]
+                    err_msg = str(e)
+                    if any(msg in err_msg for msg in read_only_msgs):
+                        raise ValueError(
+                            f"Batch mapper function {fn.__name__} tried to mutate a "
+                            "zero-copy read-only batch. To be able to mutate the "
+                            "batch, pass zero_copy_batch=False to map_batches(); "
+                            "this will create a writable copy of the batch before "
+                            "giving it to fn. To elide this copy, modify your mapper "
+                            "function so it doesn't try to mutate its input."
+                        ) from e
+                    else:
+                        raise e from None
+                else:
+                    for out_batch in res:
+                        _validate_batch_output(out_batch)
+                        yield out_batch
+
+    return transform_fn
+
+
+def _is_async_udf(fn: UserDefinedFunction) -> bool:
+    return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn)
+
+
+def _validate_row_output(item):
+    if not isinstance(item, collections.abc.Mapping):
+        raise ValueError(
+            f"Error validating {_truncated_repr(item)}: "
+            "Standalone Python objects are not "
+            "allowed in Ray >= 2.5. To return Python objects from map(), "
+            "wrap them in a dict, e.g., "
+            "return `{'item': item}` instead of just `item`."
+        )
+
+
+def _generate_transform_fn_for_map_rows(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Row, Row]:
+
+    if _is_async_udf(fn):
+        transform_fn = _generate_transform_fn_for_async_map(
+            fn,
+            _validate_row_output,
+            # NOTE: UDF concurrency is limited
+            max_concurrency=DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY,
+        )
+
+    else:
+
+        def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+            for row in rows:
+                out_row = fn(row)
+                _validate_row_output(out_row)
+                yield out_row
+
+    return transform_fn
+
+
+def _generate_transform_fn_for_flat_map(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Row, Iterable[Row]]:
+    if _is_async_udf(fn):
+        # UDF is a callable class with async generator `__call__` method.
+        transform_fn = _generate_transform_fn_for_async_map(
+            fn,
+            _validate_row_output,
+            max_concurrency=DEFAULT_ASYNC_ROW_UDF_MAX_CONCURRENCY,
+            is_flat_map=True,
+        )
+
+    else:
+
+        def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+            for row in rows:
+                for out_row in fn(row):
+                    _validate_row_output(out_row)
+                    yield out_row
+
+    return transform_fn
+
+
+def _generate_transform_fn_for_filter(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Row, Row]:
+    def transform_fn(rows: Iterable[Row], _: TaskContext) -> Iterable[Row]:
+        for row in rows:
+            if fn(row):
+                yield row
+
+    return transform_fn
+
+
+def _generate_transform_fn_for_map_block(
+    fn: UserDefinedFunction,
+) -> MapTransformCallable[Block, Block]:
+    def transform_fn(blocks: Iterable[Block], _: TaskContext) -> Iterable[Block]:
+        for block in blocks:
+            out_block = fn(block)
+            yield out_block
+
+    return transform_fn
+
+
+_SENTINEL = object()
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def _generate_transform_fn_for_async_map(
+    fn: UserDefinedFunction,
+    validate_fn: Callable,
+    *,
+    max_concurrency: int,
+    is_flat_map: bool = False,
+) -> MapTransformCallable:
+    assert max_concurrency > 0, "Max concurrency must be positive"
+
+    if inspect.isasyncgenfunction(fn):
+
+        async def _apply_udf(item: T) -> List[U]:
+            gen = fn(item)
+            # NOTE: Async generator is unrolled inside the task to maintain
+            #       requested concurrency level (`max_concurrent_batches`)
+            return [out async for out in gen]
+
+    elif inspect.iscoroutinefunction(fn):
+
+        async def _apply_udf(item: T) -> List[U]:
+            res = await fn(item)
+            return res if is_flat_map else [res]
+
+    else:
+        raise ValueError(f"Expected a coroutine function, got {fn}")
+
+    # Goals of the algorithm applying async UDF application to the provided iterator
+    # are following:
+    #
+    #   - No more than `max_concurrency` async tasks are running
+    #     at any given moment
+    #   - Slow consumption from the output queue should result in
+    #     the processing to get back-pressured (so that output queue
+    #     doesn't grow unbounded)
+    #   - Order of the items (rows/batches) produced by this method
+    #     *must be* deterministic (though is not guaranteed to be specified
+    #     if max_concurrency > 1)
+    #
+    # To achieve that, algorithm applying async UDF to elements of the provided sequence
+    # is structured like following:
+    #
+    #   - Task scheduling and subsequent results re-ordering are performed as
+    #     different stages (inside `_schedule` and `_report` methods respectively)
+    #
+    #   - Scheduling stage aim to schedule and run no more than `max_concurrency` tasks
+    #     at any given moment
+    #
+    #   - Once task completes it's added into task completion queue for its results to be
+    #     subsequently reported with deterministic ordering). Task completion queue is
+    #     capped at `maxsize=max_concurrency` elements to make sure scheduling stage is
+    #     throttled (and task completion queue isn't growing unbounded) in case when
+    #     reporting stage isn't able to keep up.
+    #
+    #   - Reporting stage dequeues completed tasks from completion queue, reorders
+    #     them (to *always* produce deterministic ordering) and adds its results into
+    #     output queue.
+    #
+    #   - Output queue is capped at `maxsize=max_concurrency` elements to make sure that
+    #     reporting stage is throttled (and output queue doesn't grow unbounded) in case
+    #     when consumer (Ray task itself) isn't able to keep up
+    #
+    async def _execute_transform(it: Iterator[T], output_queue: queue.Queue) -> None:
+        loop = asyncio.get_running_loop()
+
+        # NOTE: Individual tasks could complete in arbitrary order.
+        #       To make sure that the ordering produced by this transformation
+        #       is deterministic we utilize subsequent reordering stage to
+        #       to keep the output ordering the same as that one of the input
+        #       iterator.
+        completed_tasks_queue = asyncio.Queue(maxsize=max_concurrency)
+        # NOTE: This method is nested to support Python 3.9 where we only can
+        #       init `asyncio.Queue` inside the async function
+        async def _reorder() -> None:
+            completed_task_map: Dict[int, asyncio.Task] = dict()
+            next_idx = 0
+            completed_scheduling = False
+
+            try:
+                while not completed_scheduling:
+                    task, idx = await completed_tasks_queue.get()
+
+                    if isinstance(task, Exception):
+                        raise task
+                    elif task is _SENTINEL:
+                        completed_scheduling = True
+                    else:
+                        completed_task_map[idx] = task
+
+                    while next_idx in completed_task_map:
+                        next_task = completed_task_map.pop(next_idx)
+
+                        # NOTE: Once output queue fills up, this will block
+                        #       therefore serving as back-pressure for scheduling tasks
+                        #       preventing it from scheduling new tasks.
+                        # NOTE: This will block the whole event-loop not just this task
+                        output_queue.put(await next_task)
+
+                        next_idx += 1
+
+                assert (
+                    len(completed_task_map) == 0
+                ), f"{next_idx=}, {completed_task_map.keys()=}"
+                sentinel = _SENTINEL
+
+            except BaseException as e:
+                sentinel = e
+            finally:
+                output_queue.put(sentinel)
+
+        # NOTE: Reordering is an async process
+        asyncio.create_task(_reorder())
+
+        cur_task_map: Dict[asyncio.Task, int] = dict()
+        consumed = False
+
+        sentinel = _SENTINEL
+        enumerated_it = enumerate(it)
+
+        try:
+            while True:
+                while len(cur_task_map) < max_concurrency and not consumed:
+                    try:
+                        idx, item = next(enumerated_it)
+                        # Launch async task while keeping track of its
+                        # index in the enumerated sequence
+                        task = loop.create_task(_apply_udf(item))
+                        cur_task_map[task] = idx
+                    except StopIteration:
+                        consumed = True
+                        break
+
+                # Check if any running tasks remaining
+                if not cur_task_map:
+                    break
+
+                done, pending = await asyncio.wait(
+                    cur_task_map.keys(), return_when=asyncio.FIRST_COMPLETED
+                )
+
+                for task in done:
+                    # Report completed tasks along w/ its corresponding
+                    # index in the input sequence
+                    #
+                    # NOTE: Once completed tasks queue fills up, this will block
+                    #       therefore serving as back-pressure for scheduling tasks
+                    #       preventing it from scheduling new tasks
+                    await completed_tasks_queue.put((task, cur_task_map[task]))
+
+                    cur_task_map.pop(task)
+
+        except BaseException as e:
+            for cur_task in cur_task_map:
+                if not cur_task.done():
+                    cur_task.cancel()
+
+            sentinel = e
+        finally:
+            assert len(cur_task_map) == 0, f"{cur_task_map}"
+            await completed_tasks_queue.put((sentinel, None))
+
+    def _transform(batch_iter: Iterable[T], task_context: TaskContext) -> Iterable[U]:
+        output_queue = queue.Queue(maxsize=max_concurrency)
+
+        loop = ray.data._map_actor_context.udf_map_asyncio_loop
+
+        asyncio.run_coroutine_threadsafe(
+            _execute_transform(iter(batch_iter), output_queue), loop
+        )
+
+        while True:
+            items = output_queue.get()
+            if items is _SENTINEL:
+                break
+            elif isinstance(items, Exception):
+                raise items
+            else:
+                # NOTE: Sequences from individual UDFs are combined into a single
+                #       sequence here, as compared to letting individual UDFs to
+                #       add into the output queue to guarantee *deterministic* ordering
+                #       (necessary for Ray Data to be able to guarantee task retries
+                #       producing the same results)
+                for item in items:
+                    validate_fn(item)
+                    yield item
+
+    return _transform
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/plan_write_op.py b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_write_op.py
new file mode 100644
index 0000000000000000000000000000000000000000..8bd3f200cd593bc891a95a022dca9ddfff6cd6c0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/plan_write_op.py
@@ -0,0 +1,120 @@
+import itertools
+import uuid
+from typing import Callable, Iterator, List, Union
+
+from ray.data._internal.compute import TaskPoolStrategy
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.interfaces.task_context import TaskContext
+from ray.data._internal.execution.operators.map_operator import MapOperator
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformer,
+)
+from ray.data._internal.logical.operators.write_operator import Write
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.datasink import Datasink
+from ray.data.datasource.datasource import Datasource
+
+WRITE_UUID_KWARG_NAME = "write_uuid"
+
+
+def generate_write_fn(
+    datasink_or_legacy_datasource: Union[Datasink, Datasource], **write_args
+) -> Callable[[Iterator[Block], TaskContext], Iterator[Block]]:
+    def fn(blocks: Iterator[Block], ctx: TaskContext) -> Iterator[Block]:
+        """Writes the blocks to the given datasink or legacy datasource.
+
+        Outputs the original blocks to be written."""
+        # Create a copy of the iterator, so we can return the original blocks.
+        it1, it2 = itertools.tee(blocks, 2)
+        if isinstance(datasink_or_legacy_datasource, Datasink):
+            ctx.kwargs["_datasink_write_return"] = datasink_or_legacy_datasource.write(
+                it1, ctx
+            )
+        else:
+            datasink_or_legacy_datasource.write(it1, ctx, **write_args)
+
+        return it2
+
+    return fn
+
+
+def generate_collect_write_stats_fn() -> BlockMapTransformFn:
+    # If the write op succeeds, the resulting Dataset is a list of
+    # one Block which contain stats/metrics about the write.
+    # Otherwise, an error will be raised. The Datasource can handle
+    # execution outcomes with `on_write_complete()`` and `on_write_failed()``.
+    def fn(blocks: Iterator[Block], ctx: TaskContext) -> Iterator[Block]:
+        """Handles stats collection for block writes."""
+        block_accessors = [BlockAccessor.for_block(block) for block in blocks]
+        total_num_rows = sum(ba.num_rows() for ba in block_accessors)
+        total_size_bytes = sum(ba.size_bytes() for ba in block_accessors)
+
+        # NOTE: Write tasks can return anything, so we need to wrap it in a valid block
+        # type.
+        import pandas as pd
+
+        block = pd.DataFrame(
+            {
+                "num_rows": [total_num_rows],
+                "size_bytes": [total_size_bytes],
+                "write_return": [ctx.kwargs.get("_datasink_write_return", None)],
+            }
+        )
+        return iter([block])
+
+    return BlockMapTransformFn(
+        fn,
+        is_udf=False,
+        disable_block_shaping=True,
+    )
+
+
+def plan_write_op(
+    op: Write,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    collect_stats_fn = generate_collect_write_stats_fn()
+
+    return _plan_write_op_internal(
+        op, physical_children, data_context, extra_transformations=[collect_stats_fn]
+    )
+
+
+def _plan_write_op_internal(
+    op: Write,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+    extra_transformations: List[BlockMapTransformFn],
+) -> PhysicalOperator:
+    assert len(physical_children) == 1
+    input_physical_dag = physical_children[0]
+
+    write_fn = generate_write_fn(op._datasink_or_legacy_datasource, **op._write_args)
+
+    # Create a MapTransformer for a write operator
+    transform_fns = [
+        BlockMapTransformFn(
+            write_fn,
+            is_udf=False,
+            # NOTE: No need for block-shaping
+            disable_block_shaping=True,
+        ),
+    ] + extra_transformations
+
+    map_transformer = MapTransformer(transform_fns)
+
+    return MapOperator.create(
+        map_transformer,
+        input_physical_dag,
+        data_context,
+        name="Write",
+        # Add a UUID to write tasks to prevent filename collisions. This a UUID for the
+        # overall write operation, not the individual write tasks.
+        map_task_kwargs={WRITE_UUID_KWARG_NAME: uuid.uuid4().hex},
+        ray_remote_args=op._ray_remote_args,
+        min_rows_per_bundle=op._min_rows_per_bundled_input,
+        compute_strategy=TaskPoolStrategy(op._concurrency),
+    )
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/planner.py b/lib/python3.12/site-packages/ray/data/_internal/planner/planner.py
new file mode 100644
index 0000000000000000000000000000000000000000..c12fffa1cd1a7c4d683c17c0775099d39849c18a
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/planner.py
@@ -0,0 +1,248 @@
+from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar
+
+from ray.data._internal.execution.interfaces import PhysicalOperator
+from ray.data._internal.execution.operators.aggregate_num_rows import (
+    AggregateNumRows,
+)
+from ray.data._internal.execution.operators.input_data_buffer import (
+    InputDataBuffer,
+)
+from ray.data._internal.execution.operators.join import JoinOperator
+from ray.data._internal.execution.operators.limit_operator import LimitOperator
+from ray.data._internal.execution.operators.output_splitter import OutputSplitter
+from ray.data._internal.execution.operators.union_operator import UnionOperator
+from ray.data._internal.execution.operators.zip_operator import ZipOperator
+from ray.data._internal.logical.interfaces import (
+    LogicalOperator,
+    LogicalPlan,
+    PhysicalPlan,
+)
+from ray.data._internal.logical.operators.all_to_all_operator import (
+    AbstractAllToAll,
+)
+from ray.data._internal.logical.operators.count_operator import Count
+from ray.data._internal.logical.operators.from_operators import AbstractFrom
+from ray.data._internal.logical.operators.input_data_operator import InputData
+from ray.data._internal.logical.operators.join_operator import Join
+from ray.data._internal.logical.operators.map_operator import (
+    AbstractUDFMap,
+    Filter,
+    Project,
+    StreamingRepartition,
+)
+from ray.data._internal.logical.operators.n_ary_operator import Union, Zip
+from ray.data._internal.logical.operators.one_to_one_operator import Download, Limit
+from ray.data._internal.logical.operators.read_operator import Read
+from ray.data._internal.logical.operators.streaming_split_operator import StreamingSplit
+from ray.data._internal.logical.operators.write_operator import Write
+from ray.data._internal.planner.plan_all_to_all_op import plan_all_to_all_op
+from ray.data._internal.planner.plan_download_op import plan_download_op
+from ray.data._internal.planner.plan_read_op import plan_read_op
+from ray.data._internal.planner.plan_udf_map_op import (
+    plan_filter_op,
+    plan_project_op,
+    plan_streaming_repartition_op,
+    plan_udf_map_op,
+)
+from ray.data._internal.planner.plan_write_op import plan_write_op
+from ray.data.context import DataContext
+
+LogicalOperatorType = TypeVar("LogicalOperatorType", bound=LogicalOperator)
+PlanLogicalOpFn = Callable[
+    [LogicalOperatorType, List[PhysicalOperator], DataContext], PhysicalOperator
+]
+
+
+def plan_input_data_op(
+    logical_op: InputData,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    """Get the corresponding DAG of physical operators for InputData."""
+    assert len(physical_children) == 0
+
+    return InputDataBuffer(
+        data_context,
+        input_data=logical_op.input_data,
+    )
+
+
+def plan_from_op(
+    op: AbstractFrom,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    assert len(physical_children) == 0
+    return InputDataBuffer(data_context, op.input_data)
+
+
+def plan_zip_op(_, physical_children, data_context):
+    assert len(physical_children) >= 2
+    return ZipOperator(data_context, *physical_children)
+
+
+def plan_union_op(_, physical_children, data_context):
+    assert len(physical_children) >= 2
+    return UnionOperator(data_context, *physical_children)
+
+
+def plan_limit_op(logical_op, physical_children, data_context):
+    assert len(physical_children) == 1
+    return LimitOperator(logical_op._limit, physical_children[0], data_context)
+
+
+def plan_count_op(logical_op, physical_children, data_context):
+    assert len(physical_children) == 1
+    return AggregateNumRows(
+        [physical_children[0]], data_context, column_name=Count.COLUMN_NAME
+    )
+
+
+def plan_join_op(
+    logical_op: Join,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+) -> PhysicalOperator:
+    assert len(physical_children) == 2
+    return JoinOperator(
+        data_context=data_context,
+        left_input_op=physical_children[0],
+        right_input_op=physical_children[1],
+        join_type=logical_op._join_type,
+        left_key_columns=logical_op._left_key_columns,
+        right_key_columns=logical_op._right_key_columns,
+        left_columns_suffix=logical_op._left_columns_suffix,
+        right_columns_suffix=logical_op._right_columns_suffix,
+        num_partitions=logical_op._num_outputs,
+        partition_size_hint=logical_op._partition_size_hint,
+        aggregator_ray_remote_args_override=logical_op._aggregator_ray_remote_args,
+    )
+
+
+def plan_streaming_split_op(
+    logical_op: StreamingSplit,
+    physical_children: List[PhysicalOperator],
+    data_context: DataContext,
+):
+    assert len(physical_children) == 1
+    return OutputSplitter(
+        physical_children[0],
+        n=logical_op._num_splits,
+        equal=logical_op._equal,
+        data_context=data_context,
+        locality_hints=logical_op._locality_hints,
+    )
+
+
+class Planner:
+    """The planner to convert optimized logical to physical operators.
+
+    Note that planner is only doing operators conversion. Physical optimization work is
+    done by physical optimizer.
+    """
+
+    _DEFAULT_PLAN_FNS = {
+        Read: plan_read_op,
+        InputData: plan_input_data_op,
+        Write: plan_write_op,
+        AbstractFrom: plan_from_op,
+        Filter: plan_filter_op,
+        AbstractUDFMap: plan_udf_map_op,
+        AbstractAllToAll: plan_all_to_all_op,
+        Union: plan_union_op,
+        Zip: plan_zip_op,
+        Limit: plan_limit_op,
+        Count: plan_count_op,
+        Project: plan_project_op,
+        StreamingRepartition: plan_streaming_repartition_op,
+        Join: plan_join_op,
+        StreamingSplit: plan_streaming_split_op,
+        Download: plan_download_op,
+    }
+
+    def plan(self, logical_plan: LogicalPlan) -> PhysicalPlan:
+        """Convert logical to physical operators recursively in post-order."""
+        physical_dag, op_map = self._plan_recursively(
+            logical_plan.dag, logical_plan.context
+        )
+        physical_plan = PhysicalPlan(physical_dag, op_map, logical_plan.context)
+        return physical_plan
+
+    def get_plan_fn(self, logical_op: LogicalOperator) -> PlanLogicalOpFn:
+        plan_fn = find_plan_fn(logical_op, self._DEFAULT_PLAN_FNS)
+        if plan_fn is not None:
+            return plan_fn
+
+        raise ValueError(
+            f"Found unknown logical operator during planning: {logical_op}"
+        )
+
+    def _plan_recursively(
+        self, logical_op: LogicalOperator, data_context: DataContext
+    ) -> Tuple[PhysicalOperator, Dict[LogicalOperator, PhysicalOperator]]:
+        """Plan a logical operator and its input dependencies recursively.
+
+        Args:
+            logical_op: The logical operator to plan.
+            data_context: The data context.
+
+        Returns:
+            A tuple of the physical operator corresponding to the logical operator, and
+            a mapping from physical to logical operators.
+        """
+        op_map: Dict[PhysicalOperator, LogicalOperator] = {}
+
+        # Plan the input dependencies first.
+        physical_children = []
+        for child in logical_op.input_dependencies:
+            physical_child, child_op_map = self._plan_recursively(child, data_context)
+            physical_children.append(physical_child)
+            op_map.update(child_op_map)
+
+        plan_fn = self.get_plan_fn(logical_op)
+        # We will call `set_logical_operators()` in the following for-loop,
+        # no need to do it here.
+        physical_op = plan_fn(logical_op, physical_children, data_context)
+
+        # Traverse up the DAG, and set the mapping from physical to logical operators.
+        # At this point, all physical operators without logical operators set
+        # must have been created by the current logical operator.
+        queue = [physical_op]
+        while queue:
+            curr_physical_op = queue.pop()
+            # Once we find an operator with a logical operator set, we can stop.
+            if curr_physical_op._logical_operators:
+                break
+
+            curr_physical_op.set_logical_operators(logical_op)
+            # Add this operator to the op_map so optimizer can find it
+            op_map[curr_physical_op] = logical_op
+            queue.extend(curr_physical_op.input_dependencies)
+
+        # Also add the final operator (in case the loop didn't catch it)
+        op_map[physical_op] = logical_op
+        return physical_op, op_map
+
+
+def find_plan_fn(
+    logical_op: LogicalOperator, plan_fns: Dict[Type[LogicalOperator], PlanLogicalOpFn]
+) -> Optional[PlanLogicalOpFn]:
+    """Find the plan function for a logical operator.
+
+    This function goes through the plan functions in order and returns the first one
+    that is an instance of the logical operator type.
+
+    Args:
+        logical_op: The logical operator to find the plan function for.
+        plan_fns: The dictionary of plan functions.
+
+    Returns:
+        The plan function for the logical operator, or None if no plan function is
+        found.
+    """
+    # TODO: This implementation doesn't account for type hierarchies conflicts or
+    # multiple inheritance.
+    for op_type, plan_fn in plan_fns.items():
+        if isinstance(logical_op, op_type):
+            return plan_fn
+    return None
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/random_shuffle.py b/lib/python3.12/site-packages/ray/data/_internal/planner/random_shuffle.py
new file mode 100644
index 0000000000000000000000000000000000000000..d94b2e2f2e632bb5fef99377ffc7e37abdb26bdb
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/random_shuffle.py
@@ -0,0 +1,90 @@
+import time
+from typing import Any, Dict, List, Optional
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data.context import DataContext, ShuffleStrategy
+from ray.util.common import INT32_MAX
+
+
+def generate_random_shuffle_fn(
+    data_context: DataContext,
+    seed: Optional[int],
+    num_outputs: Optional[int] = None,
+    ray_remote_args: Optional[Dict[str, Any]] = None,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to randomly shuffle each records of blocks."""
+
+    # If no seed has been specified, pin timestamp based one
+    # so that task could be safely retried (w/o changing their output)
+    seed = seed if seed is not None else (time.time_ns() % INT32_MAX)
+
+    def fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        num_input_blocks = sum(len(r.blocks) for r in refs)
+
+        # If map_transformer is specified (e.g. from fusing
+        # MapOperator->AllToAllOperator), we pass a map function which
+        # is applied to each block before shuffling.
+        map_transformer: Optional[MapTransformer] = ctx.upstream_map_transformer
+        upstream_map_fn = None
+        nonlocal ray_remote_args
+        if map_transformer:
+            # NOTE: We override target max-block sizing of the previous
+            #       transformation to avoid unnecessary block shaping (if any)
+            map_transformer.override_target_max_block_size(None)
+
+            def upstream_map_fn(blocks):
+                return map_transformer.apply_transform(blocks, ctx)
+
+            # If there is a fused upstream operator,
+            # also use the ray_remote_args from the fused upstream operator.
+            ray_remote_args = ctx.upstream_map_ray_remote_args
+
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=True,
+            random_seed=seed,
+            upstream_map_fn=upstream_map_fn,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            if num_outputs is not None:
+                raise NotImplementedError(
+                    "Push-based shuffle doesn't support setting num_blocks yet."
+                )
+            scheduler = PushBasedShuffleTaskScheduler(shuffle_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(shuffle_spec)
+
+        return scheduler.execute(
+            refs,
+            num_outputs or num_input_blocks,
+            task_ctx=ctx,
+            map_ray_remote_args=ray_remote_args,
+            reduce_ray_remote_args=ray_remote_args,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    return fn
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/randomize_blocks.py b/lib/python3.12/site-packages/ray/data/_internal/planner/randomize_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..dea35bfd9720d70872a09c1a43ad6360952f41e6
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/randomize_blocks.py
@@ -0,0 +1,59 @@
+from typing import List
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.logical.operators.all_to_all_operator import RandomizeBlocks
+
+
+def generate_randomize_blocks_fn(
+    op: RandomizeBlocks,
+) -> AllToAllTransformFn:
+    """Generate function to randomize order of blocks."""
+
+    def fn(
+        refs: List[RefBundle],
+        context: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        import random
+
+        nonlocal op
+        blocks_with_metadata = []
+        index_to_schema = [None] * len(refs)
+        for i, ref_bundle in enumerate(refs):
+            index_to_schema[i] = ref_bundle.schema
+            blocks_with_metadata.extend(
+                (block, meta, i) for block, meta in ref_bundle.blocks
+            )
+
+        if len(blocks_with_metadata) == 0:
+            return refs, {op._name: []}
+        else:
+            if op._seed is not None:
+                random.seed(op._seed)
+            input_owned = all(b.owns_blocks for b in refs)
+            random.shuffle(blocks_with_metadata)
+            output = []
+            stats_list = []
+            for block, meta, i in blocks_with_metadata:
+                stats_list.append(meta.to_stats())
+                output.append(
+                    RefBundle(
+                        [
+                            (
+                                block,
+                                meta,
+                            )
+                        ],
+                        owns_blocks=input_owned,
+                        schema=index_to_schema[i],
+                    )
+                )
+            return output, {op._name: stats_list}
+
+    return fn
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/repartition.py b/lib/python3.12/site-packages/ray/data/_internal/planner/repartition.py
new file mode 100644
index 0000000000000000000000000000000000000000..37bb14c6f1ada71076cacecedf14501a4d3f069e
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/repartition.py
@@ -0,0 +1,88 @@
+from typing import List, Optional, Tuple
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.execution.operators.map_transformer import MapTransformer
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.planner.exchange.split_repartition_task_scheduler import (
+    SplitRepartitionTaskScheduler,
+)
+from ray.data._internal.stats import StatsDict
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_repartition_fn(
+    num_outputs: int,
+    shuffle: bool,
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to partition each records of blocks."""
+
+    def shuffle_repartition_fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        # If map_transformer is specified (e.g. from fusing
+        # MapOperator->AllToAllOperator), we pass a map function which
+        # is applied to each block before shuffling.
+        map_transformer: Optional["MapTransformer"] = ctx.upstream_map_transformer
+        upstream_map_fn = None
+        if map_transformer:
+            # NOTE: We override target max-block sizing of the previous
+            #       transformation to avoid unnecessary block shaping (if any)
+            map_transformer.override_target_max_block_size(None)
+
+            def upstream_map_fn(blocks):
+                return map_transformer.apply_transform(blocks, ctx)
+
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=False,
+            upstream_map_fn=upstream_map_fn,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(shuffle_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(shuffle_spec)
+
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    def split_repartition_fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=False,
+        )
+        scheduler = SplitRepartitionTaskScheduler(shuffle_spec)
+        return scheduler.execute(refs, num_outputs, ctx)
+
+    if shuffle:
+        return shuffle_repartition_fn
+    return split_repartition_fn
diff --git a/lib/python3.12/site-packages/ray/data/_internal/planner/sort.py b/lib/python3.12/site-packages/ray/data/_internal/planner/sort.py
new file mode 100644
index 0000000000000000000000000000000000000000..852154c66c368ad916dadc489a9317f334e96efa
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/_internal/planner/sort.py
@@ -0,0 +1,83 @@
+from functools import partial
+from typing import List, Optional, Tuple
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data._internal.stats import StatsDict
+from ray.data._internal.util import unify_ref_bundles_schema
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_sort_fn(
+    sort_key: SortKey,
+    batch_format: str,
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to sort blocks by the specified key column or key function."""
+
+    def fn(
+        sort_key: SortKey,
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> Tuple[List[RefBundle], StatsDict]:
+        blocks = []
+        for ref_bundle in refs:
+            blocks.extend(ref_bundle.block_refs)
+        if len(blocks) == 0:
+            return (blocks, {})
+
+        sort_key.validate_schema(unify_ref_bundles_schema(refs))
+
+        num_mappers = len(blocks)
+        # Use same number of output partitions.
+        num_outputs = num_mappers
+
+        # Sample boundaries for sort key.
+        if not sort_key.boundaries:
+            sample_bar = ctx.sub_progress_bar_dict[
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME
+            ]
+            boundaries = SortTaskSpec.sample_boundaries(
+                blocks, sort_key, num_outputs, sample_bar
+            )
+        else:
+            # For user-specified boundaries (which only partition by the primary
+            # sort key), reverse `boundaries` so that the partitions are produced
+            # in descending order, as desired.
+            boundaries = [(b,) for b in sort_key.boundaries]
+            if sort_key.get_descending()[0]:
+                boundaries = boundaries[::-1]
+            num_outputs = len(boundaries) + 1
+        sort_spec = SortTaskSpec(
+            boundaries=boundaries, sort_key=sort_key, batch_format=batch_format
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(sort_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(sort_spec)
+
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    # NOTE: use partial function to pass parameters to avoid error like
+    # "UnboundLocalError: local variable ... referenced before assignment",
+    # because `key` and `descending` variables are reassigned in `fn()`.
+    return partial(fn, sort_key)
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__init__.py b/lib/python3.12/site-packages/ray/data/datasource/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..753b3146fe730aa55893db4780c8ae10168c4461
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/__init__.py
@@ -0,0 +1,76 @@
+from ray.data._internal.datasource.delta_sharing_datasource import (
+    DeltaSharingDatasource,
+)
+from ray.data._internal.datasource.mcap_datasource import (
+    MCAPDatasource,
+    TimeRange,
+)
+from ray.data._internal.datasource.sql_datasource import Connection
+from ray.data._internal.savemode import SaveMode
+from ray.data.datasource.datasink import (
+    Datasink,
+    DummyOutputDatasink,
+    WriteResult,
+    WriteReturnType,
+)
+from ray.data.datasource.datasource import (
+    Datasource,
+    RandomIntRowDatasource,
+    Reader,
+    ReadTask,
+)
+from ray.data.datasource.file_based_datasource import (
+    FileBasedDatasource,
+    FileShuffleConfig,
+    _S3FileSystemWrapper,
+)
+from ray.data.datasource.file_datasink import (
+    BlockBasedFileDatasink,
+    RowBasedFileDatasink,
+)
+from ray.data.datasource.file_meta_provider import (
+    BaseFileMetadataProvider,
+    DefaultFileMetadataProvider,
+    FastFileMetadataProvider,
+    FileMetadataProvider,
+)
+from ray.data.datasource.filename_provider import FilenameProvider
+from ray.data.datasource.partitioning import (
+    Partitioning,
+    PartitionStyle,
+    PathPartitionFilter,
+    PathPartitionParser,
+)
+
+# Note: HuggingFaceDatasource should NOT be imported here, because
+# we want to only import the Hugging Face datasets library when we use
+# ray.data.from_huggingface() or HuggingFaceDatasource() directly.
+__all__ = [
+    "BaseFileMetadataProvider",
+    "Connection",
+    "Datasink",
+    "Datasource",
+    "DefaultFileMetadataProvider",
+    "DeltaSharingDatasource",
+    "DummyOutputDatasink",
+    "FastFileMetadataProvider",
+    "FileBasedDatasource",
+    "FileShuffleConfig",
+    "FileMetadataProvider",
+    "FilenameProvider",
+    "MCAPDatasource",
+    "PartitionStyle",
+    "PathPartitionFilter",
+    "PathPartitionParser",
+    "Partitioning",
+    "RandomIntRowDatasource",
+    "ReadTask",
+    "Reader",
+    "RowBasedFileDatasink",
+    "BlockBasedFileDatasink",
+    "_S3FileSystemWrapper",
+    "TimeRange",
+    "WriteResult",
+    "WriteReturnType",
+    "SaveMode",
+]
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..56c7f5945cd651542aee3f2d12814f37a603407d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e4893caf9ffb1fe0e61aa3f8fdec6d8c897405a5
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bfd7902939cde287f659df716ba60880e8e6de45
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_based_datasource.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_based_datasource.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8396ff5f3ace4ecdf51c4000bfe50168df45211e
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_based_datasource.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_datasink.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_datasink.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..32ad14bfbadfd0c824239936b196d6159f25c8b3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_datasink.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_meta_provider.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_meta_provider.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..af143bc7296a3fb77d1b1af1a7f65b6584e470c7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/file_meta_provider.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/filename_provider.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/filename_provider.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ab28ba4f6f446128ba31bc098d6ba4746359318a
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/filename_provider.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/partitioning.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/partitioning.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cddec3c2e98a95348c1c6b4361b8ab615e1d5aa4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/partitioning.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/path_util.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/path_util.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1efac1667d2a03b49fd72bf2b6190291018dd1f3
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/path_util.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/util.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7664d39b38a343588ab80d780d17f9c28552be28
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/datasource/__pycache__/util.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/datasource/datasink.py b/lib/python3.12/site-packages/ray/data/datasource/datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..2ef83344785d07c07a4e52abadcd8ff8f5e440c0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/datasink.py
@@ -0,0 +1,194 @@
+import itertools
+import logging
+from dataclasses import dataclass
+from typing import Generic, Iterable, List, Optional, TypeVar
+
+import ray
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.block import Block, BlockAccessor
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+WriteReturnType = TypeVar("WriteReturnType")
+"""Generic type for the return value of `Datasink.write`."""
+
+
+@dataclass
+@DeveloperAPI
+class WriteResult(Generic[WriteReturnType]):
+    """Aggregated result of the Datasink write operations."""
+
+    # Total number of written rows.
+    num_rows: int
+    # Total size in bytes of written data.
+    size_bytes: int
+    # All returned values of `Datasink.write`.
+    write_returns: List[WriteReturnType]
+
+    @classmethod
+    def combine(cls, *wrs: "WriteResult") -> "WriteResult":
+        num_rows = sum(wr.num_rows for wr in wrs)
+        size_bytes = sum(wr.size_bytes for wr in wrs)
+        write_returns = list(itertools.chain(*[wr.write_returns for wr in wrs]))
+
+        return WriteResult(
+            num_rows=num_rows,
+            size_bytes=size_bytes,
+            write_returns=write_returns,
+        )
+
+
+@DeveloperAPI
+class Datasink(Generic[WriteReturnType]):
+    """Interface for defining write-related logic.
+
+    If you want to write data to something that isn't built-in, subclass this class
+    and call :meth:`~ray.data.Dataset.write_datasink`.
+    """
+
+    def on_write_start(self) -> None:
+        """Callback for when a write job starts.
+
+        Use this method to perform setup for write tasks. For example, creating a
+        staging bucket in S3.
+        """
+        pass
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> WriteReturnType:
+        """Write blocks. This is used by a single write task.
+
+        Args:
+            blocks: Generator of data blocks.
+            ctx: ``TaskContext`` for the write task.
+
+        Returns:
+            Result of this write task. When the entire write operator finishes,
+            All returned values will be passed as `WriteResult.write_returns`
+            to `Datasink.on_write_complete`.
+        """
+        raise NotImplementedError
+
+    def on_write_complete(self, write_result: WriteResult[WriteReturnType]):
+        """Callback for when a write job completes.
+
+        This can be used to `commit` a write output. This method must
+        succeed prior to ``write_datasink()`` returning to the user. If this
+        method fails, then ``on_write_failed()`` is called.
+
+        Args:
+            write_result: Aggregated result of the
+               Write operator, containing write results and stats.
+        """
+        pass
+
+    def on_write_failed(self, error: Exception) -> None:
+        """Callback for when a write job fails.
+
+        This is called on a best-effort basis on write failures.
+
+        Args:
+            error: The first error encountered.
+        """
+        pass
+
+    def get_name(self) -> str:
+        """Return a human-readable name for this datasink.
+
+        This is used as the names of the write tasks.
+        """
+        name = type(self).__name__
+        datasink_suffix = "Datasink"
+        if name.startswith("_"):
+            name = name[1:]
+        if name.endswith(datasink_suffix):
+            name = name[: -len(datasink_suffix)]
+        return name
+
+    @property
+    def supports_distributed_writes(self) -> bool:
+        """If ``False``, only launch write tasks on the driver's node."""
+        return True
+
+    @property
+    def min_rows_per_write(self) -> Optional[int]:
+        """The target number of rows to pass to each :meth:`~ray.data.Datasink.write` call.
+
+        If ``None``, Ray Data passes a system-chosen number of rows.
+        """
+        return None
+
+
+@DeveloperAPI
+class DummyOutputDatasink(Datasink[None]):
+    """An example implementation of a writable datasource for testing.
+    Examples:
+        >>> import ray
+        >>> from ray.data.datasource import DummyOutputDatasink
+        >>> output = DummyOutputDatasink()
+        >>> ray.data.range(10).write_datasink(output)
+        >>> assert output.num_ok == 1
+    """
+
+    def __init__(self):
+        ctx = ray.data.DataContext.get_current()
+
+        # Setup a dummy actor to send the data. In a real datasource, write
+        # tasks would send data to an external system instead of a Ray actor.
+        @ray.remote(scheduling_strategy=ctx.scheduling_strategy)
+        class DataSink:
+            def __init__(self):
+                self.rows_written = 0
+                self.enabled = True
+
+            def write(self, block: Block) -> None:
+                block = BlockAccessor.for_block(block)
+                self.rows_written += block.num_rows()
+
+            def get_rows_written(self):
+                return self.rows_written
+
+        self.data_sink = DataSink.remote()
+        self.num_ok = 0
+        self.num_failed = 0
+        self.enabled = True
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        tasks = []
+        if not self.enabled:
+            raise ValueError("disabled")
+        for b in blocks:
+            tasks.append(self.data_sink.write.remote(b))
+        ray.get(tasks)
+
+    def on_write_complete(self, write_result: WriteResult[None]):
+        self.num_ok += 1
+
+    def on_write_failed(self, error: Exception) -> None:
+        self.num_failed += 1
+
+
+def _gen_datasink_write_result(
+    write_result_blocks: List[Block],
+) -> WriteResult:
+    import pandas as pd
+
+    assert all(
+        isinstance(block, pd.DataFrame) and len(block) == 1
+        for block in write_result_blocks
+    )
+
+    total_num_rows = sum(result["num_rows"].sum() for result in write_result_blocks)
+    total_size_bytes = sum(result["size_bytes"].sum() for result in write_result_blocks)
+
+    write_returns = [result["write_return"][0] for result in write_result_blocks]
+    return WriteResult(total_num_rows, total_size_bytes, write_returns)
diff --git a/lib/python3.12/site-packages/ray/data/datasource/datasource.py b/lib/python3.12/site-packages/ray/data/datasource/datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e5be8562977b185a7a8bbd4814e29aa7f6b5288
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/datasource.py
@@ -0,0 +1,502 @@
+import copy
+from typing import Callable, Dict, Generator, Iterable, List, Optional
+
+import numpy as np
+import pyarrow as pa
+
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.block import Block, BlockMetadata, Schema
+from ray.data.datasource.util import _iter_sliced_blocks
+from ray.data.expressions import Expr
+from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
+
+
+class _DatasourceProjectionPushdownMixin:
+    """Mixin for reading operators supporting projection pushdown"""
+
+    def supports_projection_pushdown(self) -> bool:
+        """Returns ``True`` in case ``Datasource`` supports projection operation
+        being pushed down into the reading layer"""
+        return False
+
+    def get_projection_map(self) -> Optional[Dict[str, str]]:
+        """Return the projection map (original column names -> final column names).
+
+        Returns:
+            Dict mapping original column names (in storage) to final column names
+            (after optional renames). Keys indicate which columns are selected.
+            None means all columns are selected with no renames.
+            Empty dict {} means no columns are selected.
+        """
+        return self._projection_map
+
+    def get_column_renames(self) -> Optional[Dict[str, str]]:
+        """Return the column renames from the projection map.
+
+        This is used by predicate pushdown to rewrite filter expressions
+        from renamed column names back to original column names.
+
+        Returns:
+            Dict mapping original column names to renamed names,
+            or None if no renaming has been applied.
+        """
+        if self._projection_map is None:
+            return None
+        # Only include actual renames (where key != value)
+        renames = {k: v for k, v in self._projection_map.items() if k != v}
+        return renames if renames else None
+
+    def _get_data_columns(self) -> Optional[List[str]]:
+        """Extract data columns from projection map.
+
+        Helper method for datasources that need to pass columns to legacy
+        read functions expecting separate columns and rename_map parameters.
+
+        Returns:
+            List of column names, or None if all columns should be read.
+            Empty list [] means no columns.
+        """
+        return (
+            list(self._projection_map.keys())
+            if self._projection_map is not None
+            else None
+        )
+
+    @staticmethod
+    def _combine_projection_map(
+        prev_projection_map: Optional[Dict[str, str]],
+        new_projection_map: Optional[Dict[str, str]],
+    ) -> Optional[Dict[str, str]]:
+        """Combine two projection maps via transitive composition.
+
+        Args:
+            prev_projection_map: Previous projection (original -> intermediate names)
+            new_projection_map: New projection to apply (intermediate -> final names)
+
+        Returns:
+            Combined projection map (original -> final names)
+
+        Examples:
+            >>> # Select columns a, b with no renames
+            >>> prev = {"a": "a", "b": "b"}
+            >>> # Select only 'a', rename to 'x'
+            >>> new = {"a": "x"}
+            >>> _DatasourceProjectionPushdownMixin._combine_projection_map(prev, new)
+            {'a': 'x'}
+
+            >>> # First rename a->temp
+            >>> prev = {"a": "temp"}
+            >>> # Then rename temp->final
+            >>> new = {"temp": "final"}
+            >>> _DatasourceProjectionPushdownMixin._combine_projection_map(prev, new)
+            {'a': 'final'}
+        """
+        # Handle None cases (None means "all columns, no renames")
+        if prev_projection_map is None:
+            return new_projection_map
+        elif new_projection_map is None:
+            return prev_projection_map
+
+        # Compose projections: for each original->intermediate mapping in prev,
+        # check if intermediate is selected by new projection
+        composed = {}
+        for orig_col, intermediate_name in prev_projection_map.items():
+            # If intermediate name is in new projection, follow the chain
+            if intermediate_name in new_projection_map:
+                final_name = new_projection_map[intermediate_name]
+                composed[orig_col] = final_name
+
+        # The composition already handles transitive chains correctly:
+        # prev {a: temp}, new {temp: final} -> composed {a: final}
+        # No need for collapse_transitive_map which would incorrectly remove
+        # identity mappings like {b: b}
+        return composed
+
+    def apply_projection(
+        self,
+        projection_map: Optional[Dict[str, str]],
+    ) -> "Datasource":
+        """Apply a projection to this datasource.
+
+        Args:
+            projection_map: Dict mapping original column names (in storage)
+                to final column names (after optional renames). Keys indicate
+                which columns to select. None means select all columns with no renames.
+
+        Returns:
+            A new datasource instance with the projection applied.
+        """
+        clone = copy.copy(self)
+
+        # Combine projections via transitive map composition
+        clone._projection_map = self._combine_projection_map(
+            self._projection_map, projection_map
+        )
+
+        return clone
+
+    @staticmethod
+    def _apply_rename(
+        table: "pa.Table",
+        column_rename_map: Optional[Dict[str, str]],
+    ) -> "pa.Table":
+        """Apply column renaming to a PyArrow table.
+
+        Args:
+            table: PyArrow table to rename
+            column_rename_map: Mapping from old column names to new names
+
+        Returns:
+            Table with renamed columns
+        """
+        if not column_rename_map:
+            return table
+
+        new_names = [column_rename_map.get(col, col) for col in table.schema.names]
+        return table.rename_columns(new_names)
+
+    @staticmethod
+    def _apply_rename_to_tables(
+        tables: Iterable["pa.Table"],
+        column_rename_map: Optional[Dict[str, str]],
+    ) -> Generator["pa.Table", None, None]:
+        """Wrap a table generator to apply column renaming to each table.
+
+        This helper eliminates duplication across datasources that need to apply
+        column renames to tables yielded from generators.
+
+        Args:
+            tables: Iterator/generator yielding PyArrow tables
+            column_rename_map: Mapping from old column names to new names
+
+        Yields:
+            pa.Table: Tables with renamed columns
+        """
+        for table in tables:
+            yield _DatasourceProjectionPushdownMixin._apply_rename(
+                table, column_rename_map
+            )
+
+
+class _DatasourcePredicatePushdownMixin:
+    """Mixin for reading operators supporting predicate pushdown"""
+
+    def __init__(self):
+        self._predicate_expr: Optional[Expr] = None
+
+    def supports_predicate_pushdown(self) -> bool:
+        return False
+
+    def get_current_predicate(self) -> Optional[Expr]:
+        return self._predicate_expr
+
+    def apply_predicate(
+        self,
+        predicate_expr: Expr,
+    ) -> "Datasource":
+        """Apply a predicate to this datasource.
+
+        Default implementation that combines predicates using AND.
+        Subclasses that support predicate pushdown should have a _predicate_expr
+        attribute to store the predicate.
+
+        Note: Column rebinding is handled by the PredicatePushdown rule
+        before this method is called, so the predicate_expr should already
+        reference the correct column names.
+        """
+        import copy
+
+        clone = copy.copy(self)
+
+        # Combine with existing predicate using AND
+        clone._predicate_expr = (
+            predicate_expr
+            if clone._predicate_expr is None
+            else clone._predicate_expr & predicate_expr
+        )
+
+        return clone
+
+
+@PublicAPI
+class Datasource(_DatasourceProjectionPushdownMixin, _DatasourcePredicatePushdownMixin):
+    """Interface for defining a custom :class:`~ray.data.Dataset` datasource.
+
+    To read a datasource into a dataset, use :meth:`~ray.data.read_datasource`.
+    """  # noqa: E501
+
+    def __init__(self):
+        """Initialize the datasource and its mixins."""
+        _DatasourcePredicatePushdownMixin.__init__(self)
+
+    @Deprecated
+    def create_reader(self, **read_args) -> "Reader":
+        """
+        Deprecated: Implement :meth:`~ray.data.Datasource.get_read_tasks` and
+        :meth:`~ray.data.Datasource.estimate_inmemory_data_size` instead.
+        """
+        return _LegacyDatasourceReader(self, **read_args)
+
+    @Deprecated
+    def prepare_read(self, parallelism: int, **read_args) -> List["ReadTask"]:
+        """
+        Deprecated: Implement :meth:`~ray.data.Datasource.get_read_tasks` and
+        :meth:`~ray.data.Datasource.estimate_inmemory_data_size` instead.
+        """
+        raise NotImplementedError
+
+    def get_name(self) -> str:
+        """Return a human-readable name for this datasource.
+        This will be used as the names of the read tasks.
+        """
+        name = type(self).__name__
+        datasource_suffix = "Datasource"
+        if name.endswith(datasource_suffix):
+            name = name[: -len(datasource_suffix)]
+        return name
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        """Return an estimate of the in-memory data size, or None if unknown.
+
+        Note that the in-memory data size may be larger than the on-disk data size.
+        """
+        raise NotImplementedError
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List["ReadTask"]:
+        """Execute the read and return read tasks.
+
+        Args:
+            parallelism: The requested read parallelism. The number of read
+                tasks should equal to this value if possible.
+            per_task_row_limit: The per-task row limit for the read tasks.
+        Returns:
+            A list of read tasks that can be executed to read blocks from the
+            datasource in parallel.
+        """
+        raise NotImplementedError
+
+    @property
+    def should_create_reader(self) -> bool:
+        has_implemented_get_read_tasks = (
+            type(self).get_read_tasks is not Datasource.get_read_tasks
+        )
+        has_implemented_estimate_inmemory_data_size = (
+            type(self).estimate_inmemory_data_size
+            is not Datasource.estimate_inmemory_data_size
+        )
+        return (
+            not has_implemented_get_read_tasks
+            or not has_implemented_estimate_inmemory_data_size
+        )
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        """If ``False``, only launch read tasks on the driver's node."""
+        return True
+
+
+@Deprecated
+class Reader:
+    """A bound read operation for a :class:`~ray.data.Datasource`.
+
+    This is a stateful class so that reads can be prepared in multiple stages.
+    For example, it is useful for :class:`Datasets <ray.data.Dataset>` to know the
+    in-memory size of the read prior to executing it.
+    """
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        """Return an estimate of the in-memory data size, or None if unknown.
+
+        Note that the in-memory data size may be larger than the on-disk data size.
+        """
+        raise NotImplementedError
+
+    def get_read_tasks(self, parallelism: int) -> List["ReadTask"]:
+        """Execute the read and return read tasks.
+
+        Args:
+            parallelism: The requested read parallelism. The number of read
+                tasks should equal to this value if possible.
+
+        Returns:
+            A list of read tasks that can be executed to read blocks from the
+            datasource in parallel.
+        """
+        raise NotImplementedError
+
+
+class _LegacyDatasourceReader(Reader):
+    def __init__(self, datasource: Datasource, **read_args):
+        self._datasource = datasource
+        self._read_args = read_args
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List["ReadTask"]:
+        """Execute the read and return read tasks.
+
+        Args:
+            parallelism: The requested read parallelism. The number of read
+                tasks should equal to this value if possible.
+            per_task_row_limit: The per-task row limit for the read tasks.
+
+        Returns:
+            A list of read tasks that can be executed to read blocks from the
+            datasource in parallel.
+        """
+        return self._datasource.prepare_read(parallelism, **self._read_args)
+
+
+@DeveloperAPI
+class ReadTask(Callable[[], Iterable[Block]]):
+    """A function used to read blocks from the :class:`~ray.data.Dataset`.
+
+    Read tasks are generated by :meth:`~ray.data.Datasource.get_read_tasks`,
+    and return a list of ``ray.data.Block`` when called. Initial metadata about the read
+    operation can be retrieved via the ``metadata`` attribute prior to executing the
+    read. Final metadata is returned after the read along with the blocks.
+
+    Ray will execute read tasks in remote functions to parallelize execution.
+    Note that the number of blocks returned can vary at runtime. For example,
+    if a task is reading a single large file it can return multiple blocks to
+    avoid running out of memory during the read.
+
+    The initial metadata should reflect all the blocks returned by the read,
+    e.g., if the metadata says ``num_rows=1000``, the read can return a single
+    block of 1000 rows, or multiple blocks with 1000 rows altogether.
+
+    The final metadata (returned with the actual block) reflects the exact
+    contents of the block itself.
+    """
+
+    def __init__(
+        self,
+        read_fn: Callable[[], Iterable[Block]],
+        metadata: BlockMetadata,
+        schema: Optional["Schema"] = None,
+        per_task_row_limit: Optional[int] = None,
+    ):
+        self._metadata = metadata
+        self._read_fn = read_fn
+        self._schema = schema
+        self._per_task_row_limit = per_task_row_limit
+
+    @property
+    def metadata(self) -> BlockMetadata:
+        return self._metadata
+
+    # TODO(justin): We want to remove schema from `ReadTask` later on
+    @property
+    def schema(self) -> Optional["Schema"]:
+        return self._schema
+
+    @property
+    def read_fn(self) -> Callable[[], Iterable[Block]]:
+        return self._read_fn
+
+    @property
+    def per_task_row_limit(self) -> Optional[int]:
+        """Get the per-task row limit for this read task."""
+        return self._per_task_row_limit
+
+    def __call__(self) -> Iterable[Block]:
+        result = self._read_fn()
+        if not hasattr(result, "__iter__"):
+            DeprecationWarning(
+                "Read function must return Iterable[Block], got {}. "
+                "Probably you need to return `[block]` instead of "
+                "`block`.".format(result)
+            )
+        if self._per_task_row_limit is None:
+            yield from result
+            return
+
+        yield from _iter_sliced_blocks(result, self._per_task_row_limit)
+
+
+@DeveloperAPI
+class RandomIntRowDatasource(Datasource):
+    """An example datasource that generates rows with random int64 columns.
+
+    Examples:
+        >>> import ray
+        >>> from ray.data.datasource import RandomIntRowDatasource
+        >>> source = RandomIntRowDatasource() # doctest: +SKIP
+        >>> ray.data.read_datasource( # doctest: +SKIP
+        ...     source, n=10, num_columns=2).take()
+        {'c_0': 1717767200176864416, 'c_1': 999657309586757214}
+        {'c_0': 4983608804013926748, 'c_1': 1160140066899844087}
+    """
+
+    def __init__(self, n: int, num_columns: int):
+        """Initialize the datasource that generates random-integer rows.
+
+        Args:
+            n: The number of rows to generate.
+            num_columns: The number of columns to generate.
+        """
+        self._n = n
+        self._num_columns = num_columns
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._n * self._num_columns * 8
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+    ) -> List[ReadTask]:
+        _check_pyarrow_version()
+        import pyarrow
+
+        read_tasks: List[ReadTask] = []
+        n = self._n
+        num_columns = self._num_columns
+        block_size = max(1, n // parallelism)
+
+        def make_block(count: int, num_columns: int) -> Block:
+            return pyarrow.Table.from_arrays(
+                np.random.randint(
+                    np.iinfo(np.int64).max, size=(num_columns, count), dtype=np.int64
+                ),
+                names=[f"c_{i}" for i in range(num_columns)],
+            )
+
+        schema = pyarrow.Table.from_pydict(
+            {f"c_{i}": [0] for i in range(num_columns)}
+        ).schema
+
+        i = 0
+        while i < n:
+            count = min(block_size, n - i)
+            meta = BlockMetadata(
+                num_rows=count,
+                size_bytes=8 * count * num_columns,
+                input_files=None,
+                exec_stats=None,
+            )
+            read_tasks.append(
+                ReadTask(
+                    lambda count=count, num_columns=num_columns: [
+                        make_block(count, num_columns)
+                    ],
+                    meta,
+                    schema=schema,
+                    per_task_row_limit=per_task_row_limit,
+                )
+            )
+            i += block_size
+
+        return read_tasks
+
+    def get_name(self) -> str:
+        """Return a human-readable name for this datasource.
+        This will be used as the names of the read tasks.
+        Note: overrides the base `Datasource` method.
+        """
+        return "RandomInt"
diff --git a/lib/python3.12/site-packages/ray/data/datasource/file_based_datasource.py b/lib/python3.12/site-packages/ray/data/datasource/file_based_datasource.py
new file mode 100644
index 0000000000000000000000000000000000000000..1cad8eb48a0b23a7b0077bf5961ee619db1634f0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/file_based_datasource.py
@@ -0,0 +1,559 @@
+import io
+import logging
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Union,
+)
+
+import numpy as np
+
+import ray
+from ray.data._internal.util import (
+    RetryingContextManager,
+    RetryingPyFileSystem,
+    _check_pyarrow_version,
+    _is_local_scheme,
+    infer_compression,
+    iterate_with_retry,
+    make_async_gen,
+)
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.datasource import Datasource, ReadTask
+from ray.data.datasource.file_meta_provider import (
+    BaseFileMetadataProvider,
+    DefaultFileMetadataProvider,
+)
+from ray.data.datasource.partitioning import (
+    Partitioning,
+    PathPartitionFilter,
+    PathPartitionParser,
+)
+from ray.data.datasource.path_util import (
+    _has_file_extension,
+    _resolve_paths_and_filesystem,
+)
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+
+# We should parallelize file size fetch operations beyond this threshold.
+FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD = 16
+
+# 16 file size fetches from S3 takes ~1.5 seconds with Arrow's S3FileSystem.
+PATHS_PER_FILE_SIZE_FETCH_TASK = 16
+
+
+@DeveloperAPI
+@dataclass
+class FileShuffleConfig:
+    """Configuration for file shuffling.
+
+    This configuration object controls how files are shuffled while reading file-based
+    datasets.
+
+    .. note::
+        Even if you provided a seed, you might still observe a non-deterministic row
+        order. This is because tasks are executed in parallel and their completion
+        order might vary. If you need to preserve the order of rows, set
+        `DataContext.get_current().execution_options.preserve_order`.
+
+    Args:
+        seed: An optional integer seed for the file shuffler. If provided, Ray Data
+            shuffles files deterministically based on this seed.
+
+    Example:
+        >>> import ray
+        >>> from ray.data import FileShuffleConfig
+        >>> shuffle = FileShuffleConfig(seed=42)
+        >>> ds = ray.data.read_images("s3://anonymous@ray-example-data/batoidea", shuffle=shuffle)
+    """  # noqa: E501
+
+    seed: Optional[int] = None
+
+    def __post_init__(self):
+        """Ensure that the seed is either None or an integer."""
+        if self.seed is not None and not isinstance(self.seed, int):
+            raise ValueError("Seed must be an integer or None.")
+
+
+@DeveloperAPI
+class FileBasedDatasource(Datasource):
+    """File-based datasource for reading files.
+
+    Don't use this class directly. Instead, subclass it and implement `_read_stream()`.
+    """
+
+    # If `_WRITE_FILE_PER_ROW` is `True`, this datasource calls `_write_row` and writes
+    # each row to a file. Otherwise, this datasource calls `_write_block` and writes
+    # each block to a file.
+    _WRITE_FILE_PER_ROW = False
+    _FILE_EXTENSIONS: Optional[Union[str, List[str]]] = None
+    # Number of threads for concurrent reading within each read task.
+    # If zero or negative, reading will be performed in the main thread.
+    _NUM_THREADS_PER_TASK = 0
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        *,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        schema: Optional[Union[type, "pyarrow.lib.Schema"]] = None,
+        open_stream_args: Optional[Dict[str, Any]] = None,
+        meta_provider: BaseFileMetadataProvider = DefaultFileMetadataProvider(),
+        partition_filter: PathPartitionFilter = None,
+        partitioning: Partitioning = None,
+        ignore_missing_paths: bool = False,
+        shuffle: Optional[Union[Literal["files"], FileShuffleConfig]] = None,
+        include_paths: bool = False,
+        file_extensions: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        _check_pyarrow_version()
+
+        self._supports_distributed_reads = not _is_local_scheme(paths)
+        if not self._supports_distributed_reads and ray.util.client.ray.is_connected():
+            raise ValueError(
+                "Because you're using Ray Client, read tasks scheduled on the Ray "
+                "cluster can't access your local files. To fix this issue, store "
+                "files in cloud storage or a distributed filesystem like NFS."
+            )
+
+        self._schema = schema
+        self._data_context = DataContext.get_current()
+        self._open_stream_args = open_stream_args
+        self._meta_provider = meta_provider
+        self._partition_filter = partition_filter
+        self._partitioning = partitioning
+        self._ignore_missing_paths = ignore_missing_paths
+        self._include_paths = include_paths
+        # Need this property for lineage tracking
+        self._source_paths = paths
+        paths, self._filesystem = _resolve_paths_and_filesystem(paths, filesystem)
+        self._filesystem = RetryingPyFileSystem.wrap(
+            self._filesystem, retryable_errors=self._data_context.retried_io_errors
+        )
+        paths, file_sizes = map(
+            list,
+            zip(
+                *meta_provider.expand_paths(
+                    paths,
+                    self._filesystem,
+                    partitioning,
+                    ignore_missing_paths=ignore_missing_paths,
+                )
+            ),
+        )
+
+        if ignore_missing_paths and len(paths) == 0:
+            raise ValueError(
+                "None of the provided paths exist. "
+                "The 'ignore_missing_paths' field is set to True."
+            )
+
+        if self._partition_filter is not None:
+            # Use partition filter to skip files which are not needed.
+            path_to_size = dict(zip(paths, file_sizes))
+            paths = self._partition_filter(paths)
+            file_sizes = [path_to_size[p] for p in paths]
+            if len(paths) == 0:
+                raise ValueError(
+                    "No input files found to read. Please double check that "
+                    "'partition_filter' field is set properly."
+                )
+
+        if file_extensions is not None:
+            path_to_size = dict(zip(paths, file_sizes))
+            paths = [p for p in paths if _has_file_extension(p, file_extensions)]
+            file_sizes = [path_to_size[p] for p in paths]
+            if len(paths) == 0:
+                raise ValueError(
+                    "No input files found to read with the following file extensions: "
+                    f"{file_extensions}. Please double check that "
+                    "'file_extensions' field is set properly."
+                )
+
+        _validate_shuffle_arg(shuffle)
+        self._file_metadata_shuffler = None
+        if shuffle == "files":
+            self._file_metadata_shuffler = np.random.default_rng()
+        elif isinstance(shuffle, FileShuffleConfig):
+            # Create a NumPy random generator with a fixed seed if provided
+            self._file_metadata_shuffler = np.random.default_rng(shuffle.seed)
+
+        # Read tasks serialize `FileBasedDatasource` instances, and the list of paths
+        # can be large. To avoid slow serialization speeds, we store a reference to
+        # the paths rather than the paths themselves.
+        self._paths_ref = ray.put(paths)
+        self._file_sizes_ref = ray.put(file_sizes)
+
+    def _paths(self) -> List[str]:
+        return ray.get(self._paths_ref)
+
+    def _file_sizes(self) -> List[float]:
+        return ray.get(self._file_sizes_ref)
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        total_size = 0
+        for sz in self._file_sizes():
+            if sz is not None:
+                total_size += sz
+        return total_size
+
+    def get_read_tasks(
+        self, parallelism: int, per_task_row_limit: Optional[int] = None
+    ) -> List[ReadTask]:
+        import numpy as np
+
+        open_stream_args = self._open_stream_args
+        partitioning = self._partitioning
+
+        paths = self._paths()
+        file_sizes = self._file_sizes()
+
+        if self._file_metadata_shuffler is not None:
+            files_metadata = list(zip(paths, file_sizes))
+            shuffled_files_metadata = [
+                files_metadata[i]
+                for i in self._file_metadata_shuffler.permutation(len(files_metadata))
+            ]
+            paths, file_sizes = list(map(list, zip(*shuffled_files_metadata)))
+
+        filesystem = _wrap_s3_serialization_workaround(self._filesystem)
+
+        if open_stream_args is None:
+            open_stream_args = {}
+
+        def read_files(
+            read_paths: Iterable[str],
+        ) -> Iterable[Block]:
+            nonlocal filesystem, open_stream_args, partitioning
+
+            fs = _unwrap_s3_serialization_workaround(filesystem)
+
+            for read_path in read_paths:
+                partitions: Dict[str, str] = {}
+                if partitioning is not None:
+                    parse = PathPartitionParser(partitioning)
+                    partitions = parse(read_path)
+
+                with RetryingContextManager(
+                    self._open_input_source(fs, read_path, **open_stream_args),
+                    context=self._data_context,
+                ) as f:
+                    for block in iterate_with_retry(
+                        lambda: self._read_stream(f, read_path),
+                        description="read stream iteratively",
+                        match=self._data_context.retried_io_errors,
+                    ):
+                        if partitions:
+                            block = _add_partitions(block, partitions)
+                        if self._include_paths:
+                            block_accessor = BlockAccessor.for_block(block)
+                            block = block_accessor.fill_column("path", read_path)
+                        yield block
+
+        def create_read_task_fn(read_paths, num_threads):
+            def read_task_fn():
+                nonlocal num_threads, read_paths
+
+                # TODO: We should refactor the code so that we can get the results in
+                # order even when using multiple threads.
+                if self._data_context.execution_options.preserve_order:
+                    num_threads = 0
+
+                if num_threads > 0:
+                    num_threads = min(num_threads, len(read_paths))
+
+                    logger.debug(
+                        f"Reading {len(read_paths)} files with {num_threads} threads."
+                    )
+
+                    yield from make_async_gen(
+                        iter(read_paths),
+                        read_files,
+                        num_workers=num_threads,
+                        preserve_ordering=True,
+                    )
+                else:
+                    logger.debug(f"Reading {len(read_paths)} files.")
+                    yield from read_files(read_paths)
+
+            return read_task_fn
+
+        # fix https://github.com/ray-project/ray/issues/24296
+        parallelism = min(parallelism, len(paths))
+
+        read_tasks = []
+        split_paths = np.array_split(paths, parallelism)
+        split_file_sizes = np.array_split(file_sizes, parallelism)
+
+        for read_paths, file_sizes in zip(split_paths, split_file_sizes):
+            if len(read_paths) <= 0:
+                continue
+
+            meta = self._meta_provider(
+                read_paths,
+                rows_per_file=self._rows_per_file(),
+                file_sizes=file_sizes,
+            )
+
+            read_task_fn = create_read_task_fn(read_paths, self._NUM_THREADS_PER_TASK)
+
+            read_task = ReadTask(
+                read_task_fn, meta, per_task_row_limit=per_task_row_limit
+            )
+
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+    def resolve_compression(
+        self, path: str, open_args: Dict[str, Any]
+    ) -> Optional[str]:
+        """Resolves the compression format for a stream.
+
+        Args:
+            path: The file path to resolve compression for.
+            open_args: kwargs passed to
+                `pyarrow.fs.FileSystem.open_input_stream <https://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html#pyarrow.fs.FileSystem.open_input_stream>`_
+                when opening input files to read.
+
+        Returns:
+            The compression format (e.g., "gzip", "snappy", "bz2") or None if
+            no compression is detected or specified.
+        """
+        compression = open_args.get("compression", None)
+        if compression is None:
+            compression = infer_compression(path)
+        return compression
+
+    def _resolve_buffer_size(self, open_args: Dict[str, Any]) -> Optional[int]:
+        buffer_size = open_args.pop("buffer_size", None)
+        if buffer_size is None:
+            buffer_size = self._data_context.streaming_read_buffer_size
+        return buffer_size
+
+    def _file_to_snappy_stream(
+        self,
+        file: "pyarrow.NativeFile",
+        filesystem: "RetryingPyFileSystem",
+    ) -> "pyarrow.PythonFile":
+        import pyarrow as pa
+        import snappy
+        from pyarrow.fs import HadoopFileSystem
+
+        stream = io.BytesIO()
+        if isinstance(filesystem.unwrap(), HadoopFileSystem):
+            snappy.hadoop_snappy.stream_decompress(src=file, dst=stream)
+        else:
+            snappy.stream_decompress(src=file, dst=stream)
+        stream.seek(0)
+
+        return pa.PythonFile(stream, mode="r")
+
+    def _open_input_source(
+        self,
+        filesystem: "RetryingPyFileSystem",
+        path: str,
+        **open_args,
+    ) -> "pyarrow.NativeFile":
+        """Opens a source path for reading and returns the associated Arrow NativeFile.
+
+        The default implementation opens the source path as a sequential input stream,
+        using self._data_context.streaming_read_buffer_size as the buffer size if none
+        is given by the caller.
+
+        Implementations that do not support streaming reads (e.g. that require random
+        access) should override this method.
+        """
+
+        compression = self.resolve_compression(path, open_args)
+        buffer_size = self._resolve_buffer_size(open_args)
+
+        if compression == "snappy":
+            # Arrow doesn't support streaming Snappy decompression since the canonical
+            # C++ Snappy library doesn't natively support streaming decompression. We
+            # works around this by manually decompressing the file with python-snappy.
+            open_args["compression"] = None
+            file = filesystem.open_input_stream(
+                path, buffer_size=buffer_size, **open_args
+            )
+            return self._file_to_snappy_stream(file, filesystem)
+
+        open_args["compression"] = compression
+        return filesystem.open_input_stream(path, buffer_size=buffer_size, **open_args)
+
+    def _rows_per_file(self):
+        """Returns the number of rows per file, or None if unknown."""
+        return None
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        """Streaming read a single file.
+
+        This method should be implemented by subclasses.
+        """
+        raise NotImplementedError(
+            "Subclasses of FileBasedDatasource must implement _read_stream()."
+        )
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        return self._supports_distributed_reads
+
+
+def _add_partitions(
+    data: Union["pyarrow.Table", "pd.DataFrame"], partitions: Dict[str, Any]
+) -> Union["pyarrow.Table", "pd.DataFrame"]:
+    import pandas as pd
+    import pyarrow as pa
+
+    assert isinstance(data, (pa.Table, pd.DataFrame))
+    if isinstance(data, pa.Table):
+        return _add_partitions_to_table(data, partitions)
+    if isinstance(data, pd.DataFrame):
+        return _add_partitions_to_dataframe(data, partitions)
+
+
+def _add_partitions_to_table(
+    table: "pyarrow.Table", partitions: Dict[str, Any]
+) -> "pyarrow.Table":
+    import pyarrow as pa
+    import pyarrow.compute as pc
+
+    column_names = set(table.column_names)
+    for field, value in partitions.items():
+        column = pa.array([value] * len(table))
+        if field in column_names:
+            # TODO: Handle cast error.
+            column_type = table.schema.field(field).type
+            column = column.cast(column_type)
+
+            values_are_equal = pc.all(pc.equal(column, table[field]))
+            values_are_equal = values_are_equal.as_py()
+
+            if not values_are_equal:
+                raise ValueError(
+                    f"Partition column {field} exists in table data, but partition "
+                    f"value '{value}' is different from in-data values: "
+                    f"{table[field].unique().to_pylist()}."
+                )
+
+            i = table.schema.get_field_index(field)
+            table = table.set_column(i, field, column)
+        else:
+            table = table.append_column(field, column)
+
+    return table
+
+
+def _add_partitions_to_dataframe(
+    df: "pd.DataFrame", partitions: Dict[str, Any]
+) -> "pd.DataFrame":
+    import pandas as pd
+
+    for field, value in partitions.items():
+        column = pd.Series(data=[value] * len(df), name=field)
+
+        if field in df:
+            column = column.astype(df[field].dtype)
+            mask = df[field].notna()
+            if not df[field][mask].equals(column[mask]):
+                raise ValueError(
+                    f"Partition column {field} exists in table data, but partition "
+                    f"value '{value}' is different from in-data values: "
+                    f"{list(df[field].unique())}."
+                )
+
+        df[field] = column
+
+    return df
+
+
+def _wrap_s3_serialization_workaround(filesystem: "pyarrow.fs.FileSystem"):
+    # This is needed because pa.fs.S3FileSystem assumes pa.fs is already
+    # imported before deserialization. See #17085.
+    import pyarrow as pa
+    import pyarrow.fs
+
+    base_fs = filesystem
+    if isinstance(filesystem, RetryingPyFileSystem):
+        base_fs = filesystem.unwrap()
+
+    if isinstance(base_fs, pa.fs.S3FileSystem):
+        return _S3FileSystemWrapper(filesystem)
+
+    return filesystem
+
+
+def _unwrap_s3_serialization_workaround(
+    filesystem: Union["pyarrow.fs.FileSystem", "_S3FileSystemWrapper"],
+):
+    if isinstance(filesystem, _S3FileSystemWrapper):
+        filesystem = filesystem.unwrap()
+    return filesystem
+
+
+class _S3FileSystemWrapper:
+    """pyarrow.fs.S3FileSystem wrapper that can be deserialized safely.
+
+    Importing pyarrow.fs during reconstruction triggers the pyarrow
+    S3 subsystem initialization.
+
+    NOTE: This is only needed for pyarrow<14.0.0 and should be removed
+        once the minimum supported pyarrow version exceeds that.
+        See https://github.com/apache/arrow/pull/38375 for context.
+    """
+
+    def __init__(self, fs: "pyarrow.fs.FileSystem"):
+        self._fs = fs
+
+    def unwrap(self):
+        return self._fs
+
+    @classmethod
+    def _reconstruct(cls, fs_reconstruct, fs_args):
+        # Implicitly trigger S3 subsystem initialization by importing
+        # pyarrow.fs.
+        import pyarrow.fs  # noqa: F401
+
+        return cls(fs_reconstruct(*fs_args))
+
+    def __reduce__(self):
+        return _S3FileSystemWrapper._reconstruct, self._fs.__reduce__()
+
+
+def _resolve_kwargs(
+    kwargs_fn: Callable[[], Dict[str, Any]], **kwargs
+) -> Dict[str, Any]:
+    if kwargs_fn:
+        kwarg_overrides = kwargs_fn()
+        kwargs.update(kwarg_overrides)
+    return kwargs
+
+
+def _validate_shuffle_arg(
+    shuffle: Union[Literal["files"], FileShuffleConfig, None],
+) -> None:
+    if not (
+        shuffle is None or shuffle == "files" or isinstance(shuffle, FileShuffleConfig)
+    ):
+        raise ValueError(
+            f"Invalid value for 'shuffle': {shuffle}. "
+            "Valid values are None, 'files', `FileShuffleConfig`."
+        )
diff --git a/lib/python3.12/site-packages/ray/data/datasource/file_datasink.py b/lib/python3.12/site-packages/ray/data/datasource/file_datasink.py
new file mode 100644
index 0000000000000000000000000000000000000000..954d04ccaa0256f2b7035243e414e4b0ab8c34c1
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/file_datasink.py
@@ -0,0 +1,285 @@
+import logging
+import posixpath
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional
+from urllib.parse import urlparse
+
+from ray._common.retry import call_with_retry
+from ray._private.arrow_utils import add_creatable_buckets_param_if_s3_uri
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME
+from ray.data._internal.savemode import SaveMode
+from ray.data._internal.util import (
+    RetryingPyFileSystem,
+    _is_local_scheme,
+)
+from ray.data.block import Block, BlockAccessor
+from ray.data.context import DataContext
+from ray.data.datasource.datasink import Datasink, WriteResult
+from ray.data.datasource.filename_provider import (
+    FilenameProvider,
+    _DefaultFilenameProvider,
+)
+from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow
+
+logger = logging.getLogger(__name__)
+
+
+class _FileDatasink(Datasink[None]):
+    def __init__(
+        self,
+        path: str,
+        *,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        try_create_dir: bool = True,
+        open_stream_args: Optional[Dict[str, Any]] = None,
+        filename_provider: Optional[FilenameProvider] = None,
+        dataset_uuid: Optional[str] = None,
+        file_format: Optional[str] = None,
+        mode: SaveMode = SaveMode.APPEND,
+    ):
+        """Initialize this datasink.
+
+        Args:
+            path: The folder to write files to.
+            filesystem: The filesystem to write files to. If not provided, the
+                filesystem is inferred from the path.
+            try_create_dir: Whether to create the directory to write files to.
+            open_stream_args: Arguments to pass to ``filesystem.open_output_stream``.
+            filename_provider: A :class:`ray.data.datasource.FilenameProvider` that
+                generates filenames for each row or block.
+            dataset_uuid: The UUID of the dataset being written. If specified, it's
+                included in the filename.
+            file_format: The file extension. If specified, files are written with this
+                extension.
+        """
+        if open_stream_args is None:
+            open_stream_args = {}
+
+        if filename_provider is None:
+            filename_provider = _DefaultFilenameProvider(
+                dataset_uuid=dataset_uuid, file_format=file_format
+            )
+
+        self._data_context = DataContext.get_current()
+        self.unresolved_path = path
+        paths, self.filesystem = _resolve_paths_and_filesystem(path, filesystem)
+        self.filesystem = RetryingPyFileSystem.wrap(
+            self.filesystem, retryable_errors=self._data_context.retried_io_errors
+        )
+        assert len(paths) == 1, len(paths)
+        self.path = paths[0]
+
+        self.try_create_dir = try_create_dir
+        self.open_stream_args = open_stream_args
+        self.filename_provider = filename_provider
+        self.dataset_uuid = dataset_uuid
+        self.file_format = file_format
+        self.mode = mode
+        self.has_created_dir = False
+
+    def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
+        return self.filesystem.open_output_stream(path, **self.open_stream_args)
+
+    def on_write_start(self) -> None:
+        from pyarrow.fs import FileType
+
+        dir_exists = (
+            self.filesystem.get_file_info(self.path).type is not FileType.NotFound
+        )
+        if dir_exists:
+            if self.mode == SaveMode.ERROR:
+                raise ValueError(
+                    f"Path {self.path} already exists. If this is unexpected, use mode='ignore' to ignore those files"
+                )
+            if self.mode == SaveMode.IGNORE:
+                logger.warning(f"[SaveMode={self.mode}] Skipping {self.path}")
+                return
+            if self.mode == SaveMode.OVERWRITE:
+                logger.warning(f"[SaveMode={self.mode}] Replacing contents {self.path}")
+                self.filesystem.delete_dir_contents(self.path)
+        self.has_created_dir = self._create_dir(self.path)
+
+    def _create_dir(self, dest) -> bool:
+        """Create a directory to write files to.
+
+        If ``try_create_dir`` is ``False``, this method is a no-op.
+        """
+        from pyarrow.fs import FileType
+
+        # We should skip creating directories in s3 unless the user specifically
+        # overrides this behavior. PyArrow's s3fs implementation for create_dir
+        # will attempt to check if the parent directory exists before trying to
+        # create the directory (with recursive=True it will try to do this to
+        # all of the directories until the root of the bucket). An IAM Policy that
+        # restricts access to a subset of prefixes within the bucket might cause
+        # the creation of the directory to fail even if the permissions should
+        # allow the data can be written to the specified path. For example if a
+        # a policy only allows users to write blobs prefixed with s3://bucket/foo
+        # a call to create_dir for s3://bucket/foo/bar will fail even though it
+        # should not.
+        parsed_uri = urlparse(dest)
+        is_s3_uri = parsed_uri.scheme == "s3"
+        skip_create_dir_for_s3 = is_s3_uri and not self._data_context.s3_try_create_dir
+
+        if self.try_create_dir and not skip_create_dir_for_s3:
+            if self.filesystem.get_file_info(dest).type is FileType.NotFound:
+                # Arrow's S3FileSystem doesn't allow creating buckets by default, so we
+                # add a query arg enabling bucket creation if an S3 URI is provided.
+                tmp = add_creatable_buckets_param_if_s3_uri(dest)
+                self.filesystem.create_dir(tmp, recursive=True)
+                return True
+
+        return False
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        builder = DelegatingBlockBuilder()
+        for block in blocks:
+            builder.add_block(block)
+        block = builder.build()
+        block_accessor = BlockAccessor.for_block(block)
+
+        if block_accessor.num_rows() == 0:
+            logger.warning(f"Skipped writing empty block to {self.path}")
+            return
+
+        self.write_block(block_accessor, 0, ctx)
+
+    def write_block(self, block: BlockAccessor, block_index: int, ctx: TaskContext):
+        raise NotImplementedError
+
+    def on_write_complete(self, write_result: WriteResult[None]):
+        # If no rows were written, we can delete the directory.
+        if self.has_created_dir and write_result.num_rows == 0:
+            self.filesystem.delete_dir(self.path)
+
+    @property
+    def supports_distributed_writes(self) -> bool:
+        return not _is_local_scheme(self.unresolved_path)
+
+
+@DeveloperAPI
+class RowBasedFileDatasink(_FileDatasink):
+    """A datasink that writes one row to each file.
+
+    Subclasses must implement ``write_row_to_file`` and call the superclass constructor.
+
+    Examples:
+        .. testcode::
+
+            import io
+            from typing import Any, Dict
+
+            import pyarrow
+            from PIL import Image
+
+            from ray.data.datasource import RowBasedFileDatasink
+
+            class ImageDatasink(RowBasedFileDatasink):
+                def __init__(self, path: str, *, column: str, file_format: str = "png"):
+                    super().__init__(path, file_format=file_format)
+                    self._file_format = file_format
+                    self._column = column
+
+                def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
+                    image = Image.fromarray(row[self._column])
+                    buffer = io.BytesIO()
+                    image.save(buffer, format=self._file_format)
+                    file.write(buffer.getvalue())
+    """  # noqa: E501
+
+    def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
+        """Write a row to a file.
+
+        Args:
+            row: The row to write.
+            file: The file to write the row to.
+        """
+        raise NotImplementedError
+
+    def write_block(self, block: BlockAccessor, block_index: int, ctx: TaskContext):
+        for row_index, row in enumerate(block.iter_rows(public_row_format=False)):
+            filename = self.filename_provider.get_filename_for_row(
+                row,
+                ctx.kwargs[WRITE_UUID_KWARG_NAME],
+                ctx.task_idx,
+                block_index,
+                row_index,
+            )
+            write_path = posixpath.join(self.path, filename)
+            logger.debug(f"Writing {write_path} file.")
+
+            def write_row_to_path():
+                with self.open_output_stream(write_path) as file:
+                    self.write_row_to_file(row, file)
+
+            call_with_retry(
+                write_row_to_path,
+                description=f"write '{write_path}'",
+                match=self._data_context.retried_io_errors,
+            )
+
+
+@DeveloperAPI
+class BlockBasedFileDatasink(_FileDatasink):
+    """A datasink that writes multiple rows to each file.
+
+    Subclasses must implement ``write_block_to_file`` and call the superclass
+    constructor.
+
+    Examples:
+        .. testcode::
+
+            class CSVDatasink(BlockBasedFileDatasink):
+                def __init__(self, path: str):
+                    super().__init__(path, file_format="csv")
+
+                def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+                    from pyarrow import csv
+                    csv.write_csv(block.to_arrow(), file)
+    """  # noqa: E501
+
+    def __init__(
+        self, path, *, min_rows_per_file: Optional[int] = None, **file_datasink_kwargs
+    ):
+        super().__init__(path, **file_datasink_kwargs)
+
+        self._min_rows_per_file = min_rows_per_file
+
+    def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
+        """Write a block of data to a file.
+
+        Args:
+            block: The block to write.
+            file: The file to write the block to.
+        """
+        raise NotImplementedError
+
+    def write_block(self, block: BlockAccessor, block_index: int, ctx: TaskContext):
+        filename = self.filename_provider.get_filename_for_block(
+            block, ctx.kwargs[WRITE_UUID_KWARG_NAME], ctx.task_idx, block_index
+        )
+        write_path = posixpath.join(self.path, filename)
+
+        def write_block_to_path():
+            with self.open_output_stream(write_path) as file:
+                self.write_block_to_file(block, file)
+
+        logger.debug(f"Writing {write_path} file.")
+        call_with_retry(
+            write_block_to_path,
+            description=f"write '{write_path}'",
+            match=self._data_context.retried_io_errors,
+        )
+
+    @property
+    def min_rows_per_write(self) -> Optional[int]:
+        return self._min_rows_per_file
diff --git a/lib/python3.12/site-packages/ray/data/datasource/file_meta_provider.py b/lib/python3.12/site-packages/ray/data/datasource/file_meta_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d3b2b55cb4576975621a899107a50c4b001e39e
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/file_meta_provider.py
@@ -0,0 +1,517 @@
+import itertools
+import logging
+import os
+import pathlib
+import re
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+import numpy as np
+
+from ray.data._internal.progress_bar import ProgressBar
+from ray.data._internal.remote_fn import cached_remote_fn
+from ray.data._internal.util import RetryingPyFileSystem
+from ray.data.block import BlockMetadata
+from ray.data.datasource.partitioning import Partitioning, PathPartitionFilter
+from ray.data.datasource.path_util import _has_file_extension
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class FileMetadataProvider:
+    """Abstract callable that provides metadata for the files of a single dataset block.
+
+    Current subclasses:
+        - :class:`BaseFileMetadataProvider`
+    """
+
+    def _get_block_metadata(
+        self,
+        paths: List[str],
+        **kwargs,
+    ) -> BlockMetadata:
+        """Resolves and returns block metadata for files in the given paths.
+
+        All file paths provided should belong to a single dataset block.
+
+        Args:
+            paths: The file paths for a single dataset block.
+            **kwargs: Additional kwargs used to determine block metadata.
+
+        Returns:
+            BlockMetadata aggregated across the given paths.
+        """
+        raise NotImplementedError
+
+    def __call__(
+        self,
+        paths: List[str],
+        **kwargs,
+    ) -> BlockMetadata:
+        return self._get_block_metadata(paths, **kwargs)
+
+
+@DeveloperAPI
+class BaseFileMetadataProvider(FileMetadataProvider):
+    """Abstract callable that provides metadata for
+    :class:`~ray.data.datasource.file_based_datasource.FileBasedDatasource`
+    implementations that reuse the base :meth:`~ray.data.Datasource.prepare_read`
+    method.
+
+    Also supports file and file size discovery in input directory paths.
+
+    Current subclasses:
+        - :class:`DefaultFileMetadataProvider`
+    """
+
+    def _get_block_metadata(
+        self,
+        paths: List[str],
+        *,
+        rows_per_file: Optional[int],
+        file_sizes: List[Optional[int]],
+    ) -> BlockMetadata:
+        """Resolves and returns block metadata for files of a single dataset block.
+
+        Args:
+            paths: The file paths for a single dataset block. These
+                paths will always be a subset of those previously returned from
+                :meth:`.expand_paths`.
+            rows_per_file: The fixed number of rows per input file, or None.
+            file_sizes: Optional file size per input file previously returned
+                from :meth:`.expand_paths`, where `file_sizes[i]` holds the size of
+                the file at `paths[i]`.
+
+        Returns:
+            BlockMetadata aggregated across the given file paths.
+        """
+        raise NotImplementedError
+
+    def expand_paths(
+        self,
+        paths: List[str],
+        filesystem: Optional["RetryingPyFileSystem"],
+        partitioning: Optional[Partitioning] = None,
+        ignore_missing_paths: bool = False,
+    ) -> Iterator[Tuple[str, int]]:
+        """Expands all paths into concrete file paths by walking directories.
+
+        Also returns a sidecar of file sizes.
+
+        The input paths must be normalized for compatibility with the input
+        filesystem prior to invocation.
+
+        Args:
+            paths: A list of file and/or directory paths compatible with the
+                given filesystem.
+            filesystem: The filesystem implementation that should be used for
+                expanding all paths and reading their files.
+            ignore_missing_paths: If True, ignores any file paths in ``paths`` that
+                are not found. Defaults to False.
+
+        Returns:
+            An iterator of `(file_path, file_size)` pairs. None may be returned for the
+            file size if it is either unknown or will be fetched later by
+            `_get_block_metadata()`, but the length of
+            both lists must be equal.
+        """
+        raise NotImplementedError
+
+
+@DeveloperAPI
+class DefaultFileMetadataProvider(BaseFileMetadataProvider):
+    """Default metadata provider for
+    :class:`~ray.data.datasource.file_based_datasource.FileBasedDatasource`
+    implementations that reuse the base `prepare_read` method.
+
+    Calculates block size in bytes as the sum of its constituent file sizes,
+    and assumes a fixed number of rows per file.
+    """
+
+    def _get_block_metadata(
+        self,
+        paths: List[str],
+        *,
+        rows_per_file: Optional[int],
+        file_sizes: List[Optional[int]],
+    ) -> BlockMetadata:
+        if rows_per_file is None:
+            num_rows = None
+        else:
+            num_rows = len(paths) * rows_per_file
+        return BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=None if None in file_sizes else int(sum(file_sizes)),
+            input_files=paths,
+            exec_stats=None,
+        )  # Exec stats filled in later.
+
+    def expand_paths(
+        self,
+        paths: List[str],
+        filesystem: "RetryingPyFileSystem",
+        partitioning: Optional[Partitioning] = None,
+        ignore_missing_paths: bool = False,
+    ) -> Iterator[Tuple[str, int]]:
+        yield from _expand_paths(paths, filesystem, partitioning, ignore_missing_paths)
+
+
+@DeveloperAPI
+class FastFileMetadataProvider(DefaultFileMetadataProvider):
+    """Fast Metadata provider for
+    :class:`~ray.data.datasource.file_based_datasource.FileBasedDatasource`
+    implementations.
+
+    Offers improved performance vs.
+    :class:`DefaultFileMetadataProvider`
+    by skipping directory path expansion and file size collection.
+    While this performance improvement may be negligible for local filesystems,
+    it can be substantial for cloud storage service providers.
+
+    This should only be used when all input paths exist and are known to be files.
+    """
+
+    def expand_paths(
+        self,
+        paths: List[str],
+        filesystem: "RetryingPyFileSystem",
+        partitioning: Optional[Partitioning] = None,
+        ignore_missing_paths: bool = False,
+    ) -> Iterator[Tuple[str, int]]:
+        if ignore_missing_paths:
+            raise ValueError(
+                "`ignore_missing_paths` cannot be set when used with "
+                "`FastFileMetadataProvider`. All paths must exist when "
+                "using `FastFileMetadataProvider`."
+            )
+
+        logger.warning(
+            f"Skipping expansion of {len(paths)} path(s). If your paths contain "
+            f"directories or if file size collection is required, try rerunning this "
+            f"read with `meta_provider=DefaultFileMetadataProvider()`."
+        )
+
+        yield from zip(paths, itertools.repeat(None, len(paths)))
+
+
+def _handle_read_os_error(error: OSError, paths: Union[str, List[str]]) -> str:
+    # NOTE: this is not comprehensive yet, and should be extended as more errors arise.
+    # NOTE: The latter patterns are raised in Arrow 10+, while the former is raised in
+    # Arrow < 10.
+    aws_error_pattern = (
+        r"^(?:(.*)AWS Error \[code \d+\]: No response body\.(.*))|"
+        r"(?:(.*)AWS Error UNKNOWN \(HTTP status 400\) during HeadObject operation: "
+        r"No response body\.(.*))|"
+        r"(?:(.*)AWS Error ACCESS_DENIED during HeadObject operation: No response "
+        r"body\.(.*))$"
+    )
+    if re.match(aws_error_pattern, str(error)):
+        # Specially handle AWS error when reading files, to give a clearer error
+        # message to avoid confusing users. The real issue is most likely that the AWS
+        # S3 file credentials have not been properly configured yet.
+        if isinstance(paths, str):
+            # Quote to highlight single file path in error message for better
+            # readability. List of file paths will be shown up as ['foo', 'boo'],
+            # so only quote single file path here.
+            paths = f'"{paths}"'
+        raise OSError(
+            (
+                f"Failing to read AWS S3 file(s): {paths}. "
+                "Please check that file exists and has properly configured access. "
+                "You can also run AWS CLI command to get more detailed error message "
+                "(e.g., aws s3 ls <file-name>). "
+                "See https://awscli.amazonaws.com/v2/documentation/api/latest/reference/s3/index.html "  # noqa
+                "and https://docs.ray.io/en/latest/data/creating-datasets.html#reading-from-remote-storage "  # noqa
+                "for more information."
+            )
+        )
+    else:
+        raise error
+
+
+def _list_files(
+    paths: List[str],
+    filesystem: "RetryingPyFileSystem",
+    *,
+    partition_filter: Optional[PathPartitionFilter],
+    file_extensions: Optional[List[str]],
+) -> List[Tuple[str, int]]:
+    return list(
+        _list_files_internal(
+            paths,
+            filesystem,
+            partition_filter=partition_filter,
+            file_extensions=file_extensions,
+        )
+    )
+
+
+def _list_files_internal(
+    paths: List[str],
+    filesystem: "RetryingPyFileSystem",
+    *,
+    partition_filter: Optional[PathPartitionFilter],
+    file_extensions: Optional[List[str]],
+) -> Iterator[Tuple[str, int]]:
+    default_meta_provider = DefaultFileMetadataProvider()
+
+    for path, file_size in default_meta_provider.expand_paths(paths, filesystem):
+        # HACK: PyArrow's `ParquetDataset` errors if input paths contain non-parquet
+        # files. To avoid this, we expand the input paths with the default metadata
+        # provider and then apply the partition filter or file extensions.
+        if (
+            partition_filter
+            and not partition_filter.apply(path)
+            or not _has_file_extension(path, file_extensions)
+        ):
+            continue
+
+        yield path, file_size
+
+
+def _expand_paths(
+    paths: List[str],
+    filesystem: "RetryingPyFileSystem",
+    partitioning: Optional[Partitioning],
+    ignore_missing_paths: bool = False,
+) -> Iterator[Tuple[str, int]]:
+    """Get the file sizes for all provided file paths."""
+    from pyarrow.fs import LocalFileSystem
+
+    from ray.data.datasource.file_based_datasource import (
+        FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD,
+    )
+    from ray.data.datasource.path_util import _is_http_url, _unwrap_protocol
+
+    # We break down our processing paths into a few key cases:
+    # 1. If len(paths) < threshold, fetch the file info for the individual files/paths
+    #    serially.
+    # 2. If all paths are contained under the same parent directory (or base directory,
+    #    if using partitioning), fetch all file infos at this prefix and filter to the
+    #    provided paths on the client; this should be a single file info request.
+    # 3. If more than threshold requests required, parallelize them via Ray tasks.
+    # 1. Small # of paths case.
+    is_local = isinstance(filesystem, LocalFileSystem)
+    if isinstance(filesystem, RetryingPyFileSystem):
+        is_local = isinstance(filesystem.unwrap(), LocalFileSystem)
+
+    if (
+        len(paths) < FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
+        # Local file systems are very fast to hit.
+        or is_local
+    ):
+        yield from _get_file_infos_serial(paths, filesystem, ignore_missing_paths)
+    else:
+        # 2. Common path prefix case.
+        # Get longest common path of all paths.
+        common_path = os.path.commonpath(paths)
+        # If parent directory (or base directory, if using partitioning) is common to
+        # all paths, fetch all file infos at that prefix and filter the response to the
+        # provided paths.
+        if not _is_http_url(common_path) and (
+            (
+                partitioning is not None
+                and common_path == _unwrap_protocol(partitioning.base_dir)
+            )
+            or all(str(pathlib.Path(path).parent) == common_path for path in paths)
+        ):
+            yield from _get_file_infos_common_path_prefix(
+                paths, common_path, filesystem, ignore_missing_paths
+            )
+        # 3. Parallelization case.
+        else:
+            # Parallelize requests via Ray tasks.
+            yield from _get_file_infos_parallel(paths, filesystem, ignore_missing_paths)
+
+
+def _get_file_infos_serial(
+    paths: List[str],
+    filesystem: "RetryingPyFileSystem",
+    ignore_missing_paths: bool = False,
+) -> Iterator[Tuple[str, int]]:
+    for path in paths:
+        yield from _get_file_infos(path, filesystem, ignore_missing_paths)
+
+
+def _get_file_infos_common_path_prefix(
+    paths: List[str],
+    common_path: str,
+    filesystem: "pyarrow.fs.FileSystem",
+    ignore_missing_paths: bool = False,
+) -> Iterator[Tuple[str, int]]:
+    path_to_size = {path: None for path in paths}
+    for path, file_size in _get_file_infos(
+        common_path, filesystem, ignore_missing_paths
+    ):
+        if path in path_to_size:
+            path_to_size[path] = file_size
+
+    # Check if all `paths` have file size metadata.
+    # If any of paths has no file size, fall back to get files metadata in parallel.
+    # This can happen when path is a directory, but not a file.
+    have_missing_path = False
+    for path in paths:
+        if path_to_size[path] is None:
+            logger.debug(
+                f"Finding path {path} not have file size metadata. "
+                "Fall back to get files metadata in parallel for all paths."
+            )
+            have_missing_path = True
+            break
+
+    if have_missing_path:
+        # Parallelize requests via Ray tasks.
+        yield from _get_file_infos_parallel(paths, filesystem, ignore_missing_paths)
+    else:
+        # Iterate over `paths` to yield each path in original order.
+        # NOTE: do not iterate over `path_to_size` because the dictionary skips
+        # duplicated path, while `paths` might contain duplicated path if one wants
+        # to read same file multiple times.
+        for path in paths:
+            yield path, path_to_size[path]
+
+
+def _get_file_infos_parallel(
+    paths: List[str],
+    filesystem: "RetryingPyFileSystem",
+    ignore_missing_paths: bool = False,
+) -> Iterator[Tuple[str, int]]:
+    from ray.data.datasource.file_based_datasource import (
+        PATHS_PER_FILE_SIZE_FETCH_TASK,
+        _unwrap_s3_serialization_workaround,
+        _wrap_s3_serialization_workaround,
+    )
+
+    logger.warning(
+        f"Expanding {len(paths)} path(s). This may be a HIGH LATENCY "
+        f"operation on some cloud storage services. Moving all the "
+        "paths to a common parent directory will lead to faster "
+        "metadata fetching."
+    )
+
+    # Capture the filesystem in the fetcher func closure, but wrap it in our
+    # serialization workaround to make sure that the pickle roundtrip works as expected.
+    filesystem = _wrap_s3_serialization_workaround(filesystem)
+
+    def _file_infos_fetcher(paths: List[str]) -> List[Tuple[str, int]]:
+        fs = _unwrap_s3_serialization_workaround(filesystem)
+        return list(
+            itertools.chain.from_iterable(
+                _get_file_infos(path, fs, ignore_missing_paths) for path in paths
+            )
+        )
+
+    yield from _fetch_metadata_parallel(
+        paths, _file_infos_fetcher, PATHS_PER_FILE_SIZE_FETCH_TASK
+    )
+
+
+Uri = TypeVar("Uri")
+Meta = TypeVar("Meta")
+
+
+def _fetch_metadata_parallel(
+    uris: List[Uri],
+    fetch_func: Callable[[List[Uri]], List[Meta]],
+    desired_uris_per_task: int,
+    **ray_remote_args,
+) -> Iterator[Meta]:
+    """Fetch file metadata in parallel using Ray tasks."""
+    remote_fetch_func = cached_remote_fn(fetch_func)
+    if ray_remote_args:
+        remote_fetch_func = remote_fetch_func.options(**ray_remote_args)
+    # Choose a parallelism that results in a # of metadata fetches per task that
+    # dominates the Ray task overhead while ensuring good parallelism.
+    # Always launch at least 2 parallel fetch tasks.
+    parallelism = max(len(uris) // desired_uris_per_task, 2)
+    metadata_fetch_bar = ProgressBar(
+        "Metadata Fetch Progress", total=parallelism, unit="task"
+    )
+    fetch_tasks = []
+    for uri_chunk in np.array_split(uris, parallelism):
+        if len(uri_chunk) == 0:
+            continue
+        fetch_tasks.append(remote_fetch_func.remote(uri_chunk))
+    results = metadata_fetch_bar.fetch_until_complete(fetch_tasks)
+    yield from itertools.chain.from_iterable(results)
+
+
+def _get_file_infos(
+    path: str, filesystem: "RetryingPyFileSystem", ignore_missing_path: bool = False
+) -> List[Tuple[str, int]]:
+    """Get the file info for all files at or under the provided path."""
+    from pyarrow.fs import FileType
+
+    file_infos = []
+    try:
+        file_info = filesystem.get_file_info(path)
+    except OSError as e:
+        _handle_read_os_error(e, path)
+    if file_info.type == FileType.Directory:
+        for file_path, file_size in _expand_directory(path, filesystem):
+            file_infos.append((file_path, file_size))
+    elif file_info.type == FileType.File:
+        file_infos.append((path, file_info.size))
+    elif file_info.type == FileType.NotFound and ignore_missing_path:
+        pass
+    else:
+        raise FileNotFoundError(path)
+
+    return file_infos
+
+
+def _expand_directory(
+    path: str,
+    filesystem: "RetryingPyFileSystem",
+    exclude_prefixes: Optional[List[str]] = None,
+    ignore_missing_path: bool = False,
+) -> List[Tuple[str, int]]:
+    """
+    Expand the provided directory path to a list of file paths.
+
+    Args:
+        path: The directory path to expand.
+        filesystem: The filesystem implementation that should be used for
+            reading these files.
+        exclude_prefixes: The file relative path prefixes that should be
+            excluded from the returned file set. Default excluded prefixes are
+            "." and "_".
+
+    Returns:
+        An iterator of (file_path, file_size) tuples.
+    """
+    if exclude_prefixes is None:
+        exclude_prefixes = [".", "_"]
+
+    from pyarrow.fs import FileSelector
+
+    selector = FileSelector(path, recursive=True, allow_not_found=ignore_missing_path)
+    files = filesystem.get_file_info(selector)
+    base_path = selector.base_dir
+    out = []
+    for file_ in files:
+        if not file_.is_file:
+            continue
+        file_path = file_.path
+        if not file_path.startswith(base_path):
+            continue
+        relative = file_path[len(base_path) :]
+        if any(relative.startswith(prefix) for prefix in exclude_prefixes):
+            continue
+        out.append((file_path, file_.size))
+    # We sort the paths to guarantee a stable order.
+    return sorted(out)
diff --git a/lib/python3.12/site-packages/ray/data/datasource/filename_provider.py b/lib/python3.12/site-packages/ray/data/datasource/filename_provider.py
new file mode 100644
index 0000000000000000000000000000000000000000..526ff8952d371bc12638532c5ee358ef5e2e0a21
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/filename_provider.py
@@ -0,0 +1,135 @@
+from typing import Any, Dict, Optional
+
+from ray.data.block import Block
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class FilenameProvider:
+    """Generates filenames when you write a :class:`~ray.data.Dataset`.
+
+    Use this class to customize the filenames used when writing a Dataset.
+
+    Some methods write each row to a separate file, while others write each block to a
+    separate file. For example, :meth:`ray.data.Dataset.write_images` writes individual
+    rows, and :func:`ray.data.Dataset.write_parquet` writes blocks of data. For more
+    information about blocks, see :ref:`Data internals <datasets_scheduling>`.
+
+    If you're writing each row to a separate file, implement
+    :meth:`~FilenameProvider.get_filename_for_row`. Otherwise, implement
+    :meth:`~FilenameProvider.get_filename_for_block`.
+
+    Example:
+
+        This snippet shows you how to encode labels in written files. For example, if
+        `"cat"` is a label, you might write a file named `cat_000000_000000_000000.png`.
+
+        .. testcode::
+
+            import ray
+            from ray.data.datasource import FilenameProvider
+
+            class ImageFilenameProvider(FilenameProvider):
+
+                def __init__(self, file_format: str):
+                    self.file_format = file_format
+
+                def get_filename_for_row(self, row, write_uuid, task_index, block_index, row_index):
+                    return (
+                        f"{row['label']}_{write_uuid}_{task_index:06}_{block_index:06}"
+                        f"_{row_index:06}.{self.file_format}"
+                    )
+
+            ds = ray.data.read_parquet("s3://anonymous@ray-example-data/images.parquet")
+            ds.write_images(
+                "/tmp/results",
+                column="image",
+                filename_provider=ImageFilenameProvider("png")
+            )
+    """  # noqa: E501
+
+    def get_filename_for_block(
+        self, block: Block, write_uuid: str, task_index: int, block_index: int
+    ) -> str:
+        """Generate a filename for a block of data.
+
+        .. note::
+            Filenames must be unique and deterministic for a given write UUID, and
+            task and block index.
+
+            A block consists of multiple rows and corresponds to a single output file.
+            Each task might produce a different number of blocks.
+
+        Args:
+            block: The block that will be written to a file.
+            write_uuid: The UUID of the write operation.
+            task_index: The index of the write task.
+            block_index: The index of the block *within* the write task.
+        """
+        raise NotImplementedError
+
+    def get_filename_for_row(
+        self,
+        row: Dict[str, Any],
+        write_uuid: str,
+        task_index: int,
+        block_index: int,
+        row_index: int,
+    ) -> str:
+        """Generate a filename for a row.
+
+        .. note::
+            Filenames must be unique and deterministic for a given write UUID, and
+            task, block, and row index.
+
+            A block consists of multiple rows, and each row corresponds to a single
+            output file. Each task might produce a different number of blocks, and each
+            block might contain a different number of rows.
+
+        .. tip::
+            If you require a contiguous row index into the global dataset, use
+            :meth:`~ray.data.Dataset.iter_rows`. This method is single-threaded and
+            isn't recommended for large datasets.
+
+        Args:
+            row: The row that will be written to a file.
+            write_uuid: The UUID of the write operation.
+            task_index: The index of the write task.
+            block_index: The index of the block *within* the write task.
+            row_index: The index of the row *within* the block.
+        """
+        raise NotImplementedError
+
+
+class _DefaultFilenameProvider(FilenameProvider):
+    def __init__(
+        self, dataset_uuid: Optional[str] = None, file_format: Optional[str] = None
+    ):
+        self._dataset_uuid = dataset_uuid
+        self._file_format = file_format
+
+    def get_filename_for_block(
+        self, block: Block, write_uuid: str, task_index: int, block_index: int
+    ) -> str:
+        file_id = f"{write_uuid}_{task_index:06}_{block_index:06}"
+        return self._generate_filename(file_id)
+
+    def get_filename_for_row(
+        self,
+        row: Dict[str, Any],
+        write_uuid: str,
+        task_index: int,
+        block_index: int,
+        row_index: int,
+    ) -> str:
+        file_id = f"{write_uuid}_{task_index:06}_{block_index:06}_{row_index:06}"
+        return self._generate_filename(file_id)
+
+    def _generate_filename(self, file_id: str) -> str:
+        filename = ""
+        if self._dataset_uuid is not None:
+            filename += f"{self._dataset_uuid}_"
+        filename += file_id
+        if self._file_format is not None:
+            filename += f".{self._file_format}"
+        return filename
diff --git a/lib/python3.12/site-packages/ray/data/datasource/partitioning.py b/lib/python3.12/site-packages/ray/data/datasource/partitioning.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8d32abb02ffe938c7b77ac89c28dd1644941184
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/partitioning.py
@@ -0,0 +1,463 @@
+import posixpath
+import urllib.parse
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
+
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+PartitionDataType = Type[Union[int, float, str, bool]]
+
+
+@DeveloperAPI
+class PartitionStyle(str, Enum):
+    """Supported dataset partition styles.
+
+    Inherits from `str` to simplify plain text serialization/deserialization.
+
+    Examples:
+        >>> # Serialize to JSON text.
+        >>> json.dumps(PartitionStyle.HIVE)  # doctest: +SKIP
+        '"hive"'
+
+        >>> # Deserialize from JSON text.
+        >>> PartitionStyle(json.loads('"hive"'))  # doctest: +SKIP
+        <PartitionStyle.HIVE: 'hive'>
+    """
+
+    HIVE = "hive"
+    DIRECTORY = "dir"
+
+
+@DeveloperAPI
+@dataclass
+class Partitioning:
+    """Partition scheme used to describe path-based partitions.
+
+    Path-based partition formats embed all partition keys and values directly in
+    their dataset file paths.
+
+    For example, to read a dataset with
+    `Hive-style partitions <https://athena.guide/articles/hive-style-partitioning>`_:
+
+        >>> import ray
+        >>> from ray.data.datasource.partitioning import Partitioning
+        >>> ds = ray.data.read_csv(
+        ...     "s3://anonymous@ray-example-data/iris.csv",
+        ...     partitioning=Partitioning("hive"),
+        ... )
+
+    Instead, if your files are arranged in a directory structure such as:
+
+    .. code::
+
+        root/dog/dog_0.jpeg
+        root/dog/dog_1.jpeg
+        ...
+
+        root/cat/cat_0.jpeg
+        root/cat/cat_1.jpeg
+        ...
+
+    Then you can use directory-based partitioning:
+
+        >>> import ray
+        >>> from ray.data.datasource.partitioning import Partitioning
+        >>> root = "s3://anonymous@air-example-data/cifar-10/images"
+        >>> partitioning = Partitioning("dir", field_names=["class"], base_dir=root)
+        >>> ds = ray.data.read_images(root, partitioning=partitioning)
+    """
+
+    #: The partition style - may be either HIVE or DIRECTORY.
+    style: PartitionStyle
+    #: "/"-delimited base directory that all partitioned paths should
+    #: exist under (exclusive). File paths either outside of, or at the first
+    #: level of, this directory will be considered unpartitioned. Specify
+    #: `None` or an empty string to search for partitions in all file path
+    #: directories.
+    base_dir: Optional[str] = None
+    #: The partition key field names (i.e. column names for tabular
+    #: datasets). When non-empty, the order and length of partition key
+    #: field names must match the order and length of partition values.
+    #: Required when parsing DIRECTORY partitioned paths or generating
+    #: HIVE partitioned paths.
+    field_names: Optional[List[str]] = None
+    #: A dictionary that maps partition key names to their desired data type. If not
+    #: provided, the data type defaults to string.
+    field_types: Optional[Dict[str, PartitionDataType]] = None
+    #: Filesystem that will be used for partition path file I/O.
+    filesystem: Optional["pyarrow.fs.FileSystem"] = None
+
+    def __post_init__(self):
+        if self.base_dir is None:
+            self.base_dir = ""
+
+        if self.field_types is None:
+            self.field_types = {}
+
+        self._normalized_base_dir = None
+        self._resolved_filesystem = None
+
+    @property
+    def normalized_base_dir(self) -> str:
+        """Returns the base directory normalized for compatibility with a filesystem."""
+        if self._normalized_base_dir is None:
+            self._normalize_base_dir()
+        return self._normalized_base_dir
+
+    @property
+    def resolved_filesystem(self) -> "pyarrow.fs.FileSystem":
+        """Returns the filesystem resolved for compatibility with a base directory."""
+        if self._resolved_filesystem is None:
+            self._normalize_base_dir()
+        return self._resolved_filesystem
+
+    def _normalize_base_dir(self):
+        """Normalizes the partition base directory for compatibility with the
+        given filesystem.
+
+        This should be called once a filesystem has been resolved to ensure that this
+        base directory is correctly discovered at the root of all partitioned file
+        paths.
+        """
+        from ray.data.datasource.path_util import _resolve_paths_and_filesystem
+
+        paths, self._resolved_filesystem = _resolve_paths_and_filesystem(
+            self.base_dir,
+            self.filesystem,
+        )
+        assert (
+            len(paths) == 1
+        ), f"Expected 1 normalized base directory, but found {len(paths)}"
+        normalized_base_dir = paths[0]
+        if len(normalized_base_dir) and not normalized_base_dir.endswith("/"):
+            normalized_base_dir += "/"
+        self._normalized_base_dir = normalized_base_dir
+
+
+@DeveloperAPI
+class PathPartitionParser:
+    """Partition parser for path-based partition formats.
+
+    Path-based partition formats embed all partition keys and values directly in
+    their dataset file paths.
+
+    Two path partition formats are currently supported - `HIVE` and `DIRECTORY`.
+
+    For `HIVE` Partitioning, all partition directories under the base directory
+    will be discovered based on `{key1}={value1}/{key2}={value2}` naming
+    conventions. Key/value pairs do not need to be presented in the same
+    order across all paths. Directory names nested under the base directory that
+    don't follow this naming condition will be considered unpartitioned. If a
+    partition filter is defined, then it will be called with an empty input
+    dictionary for each unpartitioned file.
+
+    For `DIRECTORY` Partitioning, all directories under the base directory will
+    be interpreted as partition values of the form `{value1}/{value2}`. An
+    accompanying ordered list of partition field names must also be provided,
+    where the order and length of all partition values must match the order and
+    length of field names. Files stored directly in the base directory will
+    be considered unpartitioned. If a partition filter is defined, then it will
+    be called with an empty input dictionary for each unpartitioned file. For
+    example, if the base directory is `"foo"`, then `"foo.csv"` and `"foo/bar.csv"`
+    would be considered unpartitioned files but `"foo/bar/baz.csv"` would be associated
+    with partition `"bar"`. If the base directory is undefined, then `"foo.csv"` would
+    be unpartitioned, `"foo/bar.csv"` would be associated with partition `"foo"`, and
+    "foo/bar/baz.csv" would be associated with partition `("foo", "bar")`.
+    """
+
+    @staticmethod
+    def of(
+        style: PartitionStyle = PartitionStyle.HIVE,
+        base_dir: Optional[str] = None,
+        field_names: Optional[List[str]] = None,
+        field_types: Optional[Dict[str, PartitionDataType]] = None,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+    ) -> "PathPartitionParser":
+        """Creates a path-based partition parser using a flattened argument list.
+
+        Args:
+            style: The partition style - may be either HIVE or DIRECTORY.
+            base_dir: "/"-delimited base directory to start searching for partitions
+                (exclusive). File paths outside of this directory will be considered
+                unpartitioned. Specify `None` or an empty string to search for
+                partitions in all file path directories.
+            field_names: The partition key names. Required for DIRECTORY partitioning.
+                Optional for HIVE partitioning. When non-empty, the order and length of
+                partition key field names must match the order and length of partition
+                directories discovered. Partition key field names are not required to
+                exist in the dataset schema.
+            field_types: A dictionary that maps partition key names to their desired
+                data type. If not provided, the data type default to string.
+            filesystem: Filesystem that will be used for partition path file I/O.
+
+        Returns:
+            The new path-based partition parser.
+        """
+        scheme = Partitioning(style, base_dir, field_names, field_types, filesystem)
+        return PathPartitionParser(scheme)
+
+    def __init__(self, partitioning: Partitioning):
+        """Creates a path-based partition parser.
+
+        Args:
+            partitioning: The path-based partition scheme. The parser starts
+                searching for partitions from this scheme's base directory. File paths
+                outside the base directory will be considered unpartitioned. If the
+                base directory is `None` or an empty string then this will search for
+                partitions in all file path directories. Field names are required for
+                DIRECTORY partitioning, and optional for HIVE partitioning. When
+                non-empty, the order and length of partition key field names must match
+                the order and length of partition directories discovered.
+        """
+        style = partitioning.style
+        field_names = partitioning.field_names
+        if style == PartitionStyle.DIRECTORY and not field_names:
+            raise ValueError(
+                "Directory partitioning requires a corresponding list of "
+                "partition key field names. Please retry your request with one "
+                "or more field names specified."
+            )
+        parsers = {
+            PartitionStyle.HIVE: self._parse_hive_path,
+            PartitionStyle.DIRECTORY: self._parse_dir_path,
+        }
+        self._parser_fn: Callable[[str], Dict[str, str]] = parsers.get(style)
+        if self._parser_fn is None:
+            raise ValueError(
+                f"Unsupported partition style: {style}. "
+                f"Supported styles: {parsers.keys()}"
+            )
+        self._scheme = partitioning
+
+    def __call__(self, path: str) -> Dict[str, str]:
+        """Parses partition keys and values from a single file path.
+
+        Args:
+            path: Input file path to parse.
+
+        Returns:
+            Dictionary mapping directory partition keys to values from the input file
+            path. Returns an empty dictionary for unpartitioned files.
+        """
+        dir_path = self._dir_path_trim_base(path)
+        if dir_path is None:
+            return {}
+        partitions: Dict[str, str] = self._parser_fn(dir_path)
+
+        for field, data_type in self._scheme.field_types.items():
+            partitions[field] = _cast_value(partitions[field], data_type)
+
+        return partitions
+
+    @property
+    def scheme(self) -> Partitioning:
+        """Returns the partitioning for this parser."""
+        return self._scheme
+
+    def _dir_path_trim_base(self, path: str) -> Optional[str]:
+        """Trims the normalized base directory and returns the directory path.
+
+        Returns None if the path does not start with the normalized base directory.
+        Simply returns the directory path if the base directory is undefined.
+        """
+        if not path.startswith(self._scheme.normalized_base_dir):
+            return None
+        path = path[len(self._scheme.normalized_base_dir) :]
+        return posixpath.dirname(path)
+
+    def _parse_hive_path(self, dir_path: str) -> Dict[str, str]:
+        """Hive partition path parser.
+
+        Returns a dictionary mapping partition keys to values given a hive-style
+        partition path of the form "{key1}={value1}/{key2}={value2}/..." or an empty
+        dictionary for unpartitioned files.
+        """
+        dirs = [d for d in dir_path.split("/") if d and (d.count("=") == 1)]
+        kv_pairs = [d.split("=") for d in dirs] if dirs else []
+        # NOTE: PyArrow URL-encodes partition values when writing to cloud storage. To
+        #       ensure the values are consistent when you read them back, we need to
+        #       URL-decode them. See https://github.com/apache/arrow/issues/34905.
+        kv_pairs = [[key, urllib.parse.unquote(value)] for key, value in kv_pairs]
+
+        field_names = self._scheme.field_names
+        if field_names and kv_pairs:
+            if len(kv_pairs) != len(field_names):
+                raise ValueError(
+                    f"Expected {len(field_names)} partition value(s) but found "
+                    f"{len(kv_pairs)}: {kv_pairs}."
+                )
+            for i, field_name in enumerate(field_names):
+                if kv_pairs[i][0] != field_name:
+                    raise ValueError(
+                        f"Expected partition key {field_name} but found "
+                        f"{kv_pairs[i][0]}"
+                    )
+        return dict(kv_pairs)
+
+    def _parse_dir_path(self, dir_path: str) -> Dict[str, str]:
+        """Directory partition path parser.
+
+        Returns a dictionary mapping directory partition keys to values from a
+        partition path of the form "{value1}/{value2}/..." or an empty dictionary for
+        unpartitioned files.
+
+        Requires a corresponding ordered list of partition key field names to map the
+        correct key to each value.
+        """
+        dirs = [d for d in dir_path.split("/") if d]
+        field_names = self._scheme.field_names
+
+        if dirs and len(dirs) != len(field_names):
+            raise ValueError(
+                f"Expected {len(field_names)} partition value(s) but found "
+                f"{len(dirs)}: {dirs}."
+            )
+
+        if not dirs:
+            return {}
+        return {
+            field: directory
+            for field, directory in zip(field_names, dirs)
+            if field is not None
+        }
+
+
+@PublicAPI(stability="beta")
+class PathPartitionFilter:
+    """Partition filter for path-based partition formats.
+
+    Used to explicitly keep or reject files based on a custom filter function that
+    takes partition keys and values parsed from the file's path as input.
+    """
+
+    @staticmethod
+    def of(
+        filter_fn: Callable[[Dict[str, str]], bool],
+        style: PartitionStyle = PartitionStyle.HIVE,
+        base_dir: Optional[str] = None,
+        field_names: Optional[List[str]] = None,
+        field_types: Optional[Dict[str, PartitionDataType]] = None,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+    ) -> "PathPartitionFilter":
+        """Creates a path-based partition filter using a flattened argument list.
+
+        Args:
+            filter_fn: Callback used to filter partitions. Takes a dictionary mapping
+                partition keys to values as input. Unpartitioned files are denoted with
+                an empty input dictionary. Returns `True` to read a file for that
+                partition or `False` to skip it. Partition keys and values are always
+                strings read from the filesystem path. For example, this removes all
+                unpartitioned files:
+
+                .. code:: python
+
+                    lambda d: True if d else False
+
+                This raises an assertion error for any unpartitioned file found:
+
+                .. code:: python
+
+                    def do_assert(val, msg):
+                        assert val, msg
+
+                    lambda d: do_assert(d, "Expected all files to be partitioned!")
+
+                And this only reads files from January, 2022 partitions:
+
+                .. code:: python
+
+                    lambda d: d["month"] == "January" and d["year"] == "2022"
+
+            style: The partition style - may be either HIVE or DIRECTORY.
+            base_dir: "/"-delimited base directory to start searching for partitions
+                (exclusive). File paths outside of this directory will be considered
+                unpartitioned. Specify `None` or an empty string to search for
+                partitions in all file path directories.
+            field_names: The partition key names. Required for DIRECTORY partitioning.
+                Optional for HIVE partitioning. When non-empty, the order and length of
+                partition key field names must match the order and length of partition
+                directories discovered. Partition key field names are not required to
+                exist in the dataset schema.
+            field_types: A dictionary that maps partition key names to their desired
+                data type. If not provided, the data type defaults to string.
+            filesystem: Filesystem that will be used for partition path file I/O.
+
+        Returns:
+            The new path-based partition filter.
+        """
+        scheme = Partitioning(style, base_dir, field_names, field_types, filesystem)
+        path_partition_parser = PathPartitionParser(scheme)
+        return PathPartitionFilter(path_partition_parser, filter_fn)
+
+    def __init__(
+        self,
+        path_partition_parser: PathPartitionParser,
+        filter_fn: Callable[[Dict[str, str]], bool],
+    ):
+        """Creates a new path-based partition filter based on a parser.
+
+        Args:
+            path_partition_parser: The path-based partition parser.
+            filter_fn: Callback used to filter partitions. Takes a dictionary mapping
+                partition keys to values as input. Unpartitioned files are denoted with
+                an empty input dictionary. Returns `True` to read a file for that
+                partition or `False` to skip it. Partition keys and values are always
+                strings read from the filesystem path. For example, this removes all
+                unpartitioned files:
+                ``lambda d: True if d else False``
+                This raises an assertion error for any unpartitioned file found:
+                ``lambda d: assert d, "Expected all files to be partitioned!"``
+                And this only reads files from January, 2022 partitions:
+                ``lambda d: d["month"] == "January" and d["year"] == "2022"``
+        """
+        self._parser = path_partition_parser
+        self._filter_fn = filter_fn
+
+    def __call__(self, paths: List[str]) -> List[str]:
+        """Returns all paths that pass this partition scheme's partition filter.
+
+        If no partition filter is set, then returns all input paths. If a base
+        directory is set, then only paths under this base directory will be parsed
+        for partitions. All paths outside of this base directory will automatically
+        be considered unpartitioned, and passed into the filter function as empty
+        dictionaries.
+
+        Also normalizes the partition base directory for compatibility with the
+        given filesystem before applying the filter.
+
+        Args:
+            paths: Paths to pass through the partition filter function. All
+                paths should be normalized for compatibility with the given
+                filesystem.
+        Returns:
+            List of paths that pass the partition filter, or all paths if no
+            partition filter is defined.
+        """
+        filtered_paths = paths
+        if self._filter_fn is not None:
+            filtered_paths = [path for path in paths if self.apply(path)]
+        return filtered_paths
+
+    def apply(self, path: str) -> bool:
+        return self._filter_fn(self._parser(path))
+
+    @property
+    def parser(self) -> PathPartitionParser:
+        """Returns the path partition parser for this filter."""
+        return self._parser
+
+
+def _cast_value(value: str, data_type: PartitionDataType) -> Any:
+    if data_type is int:
+        return int(value)
+    elif data_type is float:
+        return float(value)
+    elif data_type is bool:
+        return value.lower() == "true"
+    else:
+        return value
diff --git a/lib/python3.12/site-packages/ray/data/datasource/path_util.py b/lib/python3.12/site-packages/ray/data/datasource/path_util.py
new file mode 100644
index 0000000000000000000000000000000000000000..d120f3b8b09027f81badd5592f515b70eff573c6
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/path_util.py
@@ -0,0 +1,238 @@
+import pathlib
+import sys
+import urllib
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+from ray.data._internal.util import RetryingPyFileSystem, _resolve_custom_scheme
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+def _has_file_extension(path: str, extensions: Optional[List[str]]) -> bool:
+    """Check if a path has a file extension in the provided list.
+
+    Examples:
+        >>> _has_file_extension("foo.csv", ["csv"])
+        True
+        >>> _has_file_extension("foo.CSV", ["csv"])
+        True
+        >>> _has_file_extension("foo.CSV", [".csv"])
+        True
+        >>> _has_file_extension("foo.csv", ["json", "jsonl"])
+        False
+        >>> _has_file_extension("foo.csv", None)
+        True
+
+    Args:
+        path: The path to check.
+        extensions: A list of extensions to check against. If `None`, any extension is
+            considered valid.
+    """
+    assert extensions is None or isinstance(extensions, list), type(extensions)
+
+    if extensions is None:
+        return True
+
+    # If the user-specified extensions don't contain a leading dot, we add it here
+    extensions = [
+        f".{ext.lower()}" if not ext.startswith(".") else ext.lower()
+        for ext in extensions
+    ]
+    return any(path.lower().endswith(ext) for ext in extensions)
+
+
+def _resolve_paths_and_filesystem(
+    paths: Union[str, List[str]],
+    filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+) -> Tuple[List[str], "pyarrow.fs.FileSystem"]:
+    """
+    Resolves and normalizes all provided paths, infers a filesystem from the
+    paths and ensures that all paths use the same filesystem.
+
+    Args:
+        paths: A single file/directory path or a list of file/directory paths.
+            A list of paths can contain both files and directories.
+        filesystem: The filesystem implementation that should be used for
+            reading these files. If None, a filesystem will be inferred. If not
+            None, the provided filesystem will still be validated against all
+            filesystems inferred from the provided paths to ensure
+            compatibility.
+    """
+    import pyarrow as pa
+    from pyarrow.fs import (
+        FileSystem,
+        FSSpecHandler,
+        PyFileSystem,
+        _resolve_filesystem_and_path,
+    )
+
+    if isinstance(paths, str):
+        paths = [paths]
+    if isinstance(paths, pathlib.Path):
+        paths = [str(paths)]
+    elif not isinstance(paths, list) or any(not isinstance(p, str) for p in paths):
+        raise ValueError(
+            "Expected `paths` to be a `str`, `pathlib.Path`, or `list[str]`, but got "
+            f"`{paths}`"
+        )
+    elif len(paths) == 0:
+        raise ValueError("Must provide at least one path.")
+
+    if filesystem and not isinstance(filesystem, FileSystem):
+        err_msg = (
+            f"The filesystem passed must either conform to "
+            f"pyarrow.fs.FileSystem, or "
+            f"fsspec.spec.AbstractFileSystem. The provided "
+            f"filesystem was: {filesystem}"
+        )
+        try:
+            import fsspec  # noqa: F401
+            from fsspec.implementations.http import HTTPFileSystem
+        except ModuleNotFoundError:
+            # If filesystem is not a pyarrow filesystem and fsspec isn't
+            # installed, then filesystem is neither a pyarrow filesystem nor
+            # an fsspec filesystem, so we raise a TypeError.
+            raise TypeError(err_msg) from None
+        if not isinstance(filesystem, fsspec.spec.AbstractFileSystem):
+            raise TypeError(err_msg) from None
+
+        filesystem = PyFileSystem(FSSpecHandler(filesystem))
+
+    resolved_paths = []
+    for path in paths:
+        path = _resolve_custom_scheme(path)
+        try:
+            resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
+                path, filesystem
+            )
+        except pa.lib.ArrowInvalid as e:
+            if "Cannot parse URI" in str(e):
+                resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
+                    _encode_url(path), filesystem
+                )
+                resolved_path = _decode_url(resolved_path)
+            elif "Unrecognized filesystem type in URI" in str(e):
+                scheme = urllib.parse.urlparse(path, allow_fragments=False).scheme
+                if scheme in ["http", "https"]:
+                    # If scheme of path is HTTP and filesystem is not resolved,
+                    # try to use fsspec HTTPFileSystem. This expects fsspec is
+                    # installed.
+                    try:
+                        import fsspec  # noqa: F401
+                        from fsspec.implementations.http import HTTPFileSystem
+                    except ModuleNotFoundError:
+                        raise ImportError(
+                            "Please install fsspec to read files from HTTP."
+                        ) from None
+
+                    resolved_filesystem = PyFileSystem(FSSpecHandler(HTTPFileSystem()))
+                    resolved_path = path
+                else:
+                    raise
+            else:
+                raise
+
+        if filesystem is None:
+            filesystem = resolved_filesystem
+
+        # If the PyArrow filesystem is handled by a fsspec HTTPFileSystem, the protocol/
+        # scheme of paths should not be unwrapped/removed, because HTTPFileSystem
+        # expects full file paths including protocol/scheme. This is different behavior
+        # compared to other file system implementation in pyarrow.fs.FileSystem.
+        if not _is_http_filesystem(filesystem):
+            resolved_path = _unwrap_protocol(resolved_path)
+
+        resolved_path = filesystem.normalize_path(resolved_path)
+        resolved_paths.append(resolved_path)
+
+    return resolved_paths, filesystem
+
+
+def _is_http_filesystem(fs: "pyarrow.fs.FileSystem") -> bool:
+    """Return whether ``fs`` is a PyFileSystem handled by a fsspec HTTPFileSystem."""
+    from pyarrow.fs import FSSpecHandler, PyFileSystem
+
+    try:
+        import fsspec  # noqa: F401
+        from fsspec.implementations.http import HTTPFileSystem
+    except ModuleNotFoundError:
+        return False
+
+    if isinstance(fs, RetryingPyFileSystem):
+        fs = fs.unwrap()
+
+    if not isinstance(fs, PyFileSystem):
+        return False
+
+    return isinstance(fs.handler, FSSpecHandler) and isinstance(
+        fs.handler.fs, HTTPFileSystem
+    )
+
+
+def _unwrap_protocol(path):
+    """
+    Slice off any protocol prefixes on path.
+    """
+    if sys.platform == "win32" and _is_local_windows_path(path):
+        # Represent as posix path such that downstream functions properly handle it.
+        # This is executed when 'file://' is NOT included in the path.
+        return pathlib.Path(path).as_posix()
+
+    parsed = urllib.parse.urlparse(path, allow_fragments=False)  # support '#' in path
+    params = ";" + parsed.params if parsed.params else ""  # support ';' in path
+    query = "?" + parsed.query if parsed.query else ""  # support '?' in path
+    netloc = parsed.netloc
+    if parsed.scheme == "s3" and "@" in parsed.netloc:
+        # If the path contains an @, it is assumed to be an anonymous
+        # credentialed path, and we need to strip off the credentials.
+        netloc = parsed.netloc.split("@")[-1]
+
+    parsed_path = parsed.path
+    # urlparse prepends the path with a '/'. This does not work on Windows
+    # so if this is the case strip the leading slash.
+    if (
+        sys.platform == "win32"
+        and not netloc
+        and len(parsed_path) >= 3
+        and parsed_path[0] == "/"  # The problematic leading slash
+        and parsed_path[1].isalpha()  # Ensure it is a drive letter.
+        and parsed_path[2:4] in (":", ":/")
+    ):
+        parsed_path = parsed_path[1:]
+
+    return netloc + parsed_path + params + query
+
+
+def _is_url(path) -> bool:
+    return urllib.parse.urlparse(path).scheme != ""
+
+
+def _is_http_url(path) -> bool:
+    parsed = urllib.parse.urlparse(path)
+    return parsed.scheme in ("http", "https")
+
+
+def _is_local_windows_path(path: str) -> bool:
+    """Determines if path is a Windows file-system location."""
+    if sys.platform != "win32":
+        return False
+
+    if len(path) >= 1 and path[0] == "\\":
+        return True
+    if (
+        len(path) >= 3
+        and path[1] == ":"
+        and (path[2] == "/" or path[2] == "\\")
+        and path[0].isalpha()
+    ):
+        return True
+    return False
+
+
+def _encode_url(path):
+    return urllib.parse.quote(path, safe="/:")
+
+
+def _decode_url(path):
+    return urllib.parse.unquote(path)
diff --git a/lib/python3.12/site-packages/ray/data/datasource/util.py b/lib/python3.12/site-packages/ray/data/datasource/util.py
new file mode 100644
index 0000000000000000000000000000000000000000..49784ba7a8a17b61d1c381a32bb00719fc68cae0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/datasource/util.py
@@ -0,0 +1,28 @@
+from typing import Iterable
+
+from ray.data.block import Block
+
+
+def _iter_sliced_blocks(
+    blocks: Iterable[Block], per_task_row_limit: int
+) -> Iterable[Block]:
+    """Iterate over blocks, accumulating rows up to the per-task row limit."""
+    rows_read = 0
+    for block in blocks:
+        if rows_read >= per_task_row_limit:
+            break
+
+        from ray.data.block import BlockAccessor
+
+        accessor = BlockAccessor.for_block(block)
+        block_rows = accessor.num_rows()
+
+        if rows_read + block_rows <= per_task_row_limit:
+            yield block
+            rows_read += block_rows
+        else:
+            # Slice the block to meet the limit exactly
+            remaining_rows = per_task_row_limit - rows_read
+            sliced_block = accessor.slice(0, remaining_rows, copy=True)
+            yield sliced_block
+            break
diff --git a/lib/python3.12/site-packages/ray/data/extensions/__init__.py b/lib/python3.12/site-packages/ray/data/extensions/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..517b4fe7a3a2b6274f35714c709625fade8ef46e
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/extensions/__init__.py
@@ -0,0 +1,45 @@
+from ray.air.util.tensor_extensions.arrow import (
+    ArrowTensorTypeV2,
+    get_arrow_extension_tensor_types,
+)
+from ray.data.extensions.object_extension import (
+    ArrowPythonObjectArray,
+    ArrowPythonObjectScalar,
+    ArrowPythonObjectType,
+    PythonObjectArray,
+    PythonObjectDtype,
+    _object_extension_type_allowed,
+)
+from ray.data.extensions.tensor_extension import (
+    ArrowConversionError,
+    ArrowTensorArray,
+    ArrowTensorType,
+    ArrowVariableShapedTensorArray,
+    ArrowVariableShapedTensorType,
+    TensorArray,
+    TensorArrayElement,
+    TensorDtype,
+    column_needs_tensor_extension,
+)
+
+__all__ = [
+    # Tensor array extension.
+    "TensorDtype",
+    "TensorArray",
+    "TensorArrayElement",
+    "ArrowTensorType",
+    "ArrowTensorTypeV2",
+    "ArrowTensorArray",
+    "ArrowVariableShapedTensorType",
+    "ArrowVariableShapedTensorArray",
+    "column_needs_tensor_extension",
+    "ArrowConversionError",
+    # Object array extension
+    "ArrowPythonObjectArray",
+    "ArrowPythonObjectType",
+    "ArrowPythonObjectScalar",
+    "PythonObjectArray",
+    "PythonObjectDtype",
+    "_object_extension_type_allowed",
+    "get_arrow_extension_tensor_types",
+]
diff --git a/lib/python3.12/site-packages/ray/data/extensions/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/extensions/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a462230cf476258cb63b861b1555b11e2f1ca3c5
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/extensions/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/extensions/__pycache__/object_extension.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/extensions/__pycache__/object_extension.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9e4af3c4c86d921ad9eab8127868b8846da6dec9
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/extensions/__pycache__/object_extension.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/extensions/__pycache__/tensor_extension.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/extensions/__pycache__/tensor_extension.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a4cdda16e0c4465d5bb74598585137da7e8ac472
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/extensions/__pycache__/tensor_extension.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/extensions/object_extension.py b/lib/python3.12/site-packages/ray/data/extensions/object_extension.py
new file mode 100644
index 0000000000000000000000000000000000000000..42ab20a231c62e607ba531e1f25f0b2d91c96178
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/extensions/object_extension.py
@@ -0,0 +1,10 @@
+from ray.air.util.object_extensions.arrow import (  # noqa: F401
+    ArrowPythonObjectArray,
+    ArrowPythonObjectScalar,
+    ArrowPythonObjectType,
+    _object_extension_type_allowed,
+)
+from ray.air.util.object_extensions.pandas import (  # noqa: F401
+    PythonObjectArray,
+    PythonObjectDtype,
+)
diff --git a/lib/python3.12/site-packages/ray/data/extensions/tensor_extension.py b/lib/python3.12/site-packages/ray/data/extensions/tensor_extension.py
new file mode 100644
index 0000000000000000000000000000000000000000..121685e4c5ad07f95b6d24f0b721e28c65590482
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/extensions/tensor_extension.py
@@ -0,0 +1,15 @@
+from ray.air.util.tensor_extensions.arrow import (  # noqa: F401
+    ArrowConversionError,
+    ArrowTensorArray,
+    ArrowTensorType,
+    ArrowTensorTypeV2,
+    ArrowVariableShapedTensorArray,
+    ArrowVariableShapedTensorType,
+)
+from ray.air.util.tensor_extensions.pandas import (  # noqa: F401
+    TensorArray,
+    TensorArrayElement,
+    TensorDtype,
+    column_needs_tensor_extension,
+)
+from ray.air.util.tensor_extensions.utils import create_ragged_ndarray  # noqa: F401
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/__init__.py b/lib/python3.12/site-packages/ray/data/namespace_expressions/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a20192dbe6cd4621deadc841b52e51f872ae76bb
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/list_namespace.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/list_namespace.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f17d1178083ac2246ee1231243afd1a4429b821f
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/list_namespace.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/string_namespace.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/string_namespace.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..473492cd4cbb0501faaf306885300191b00748bf
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/string_namespace.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/struct_namespace.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/struct_namespace.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1bf6c96d044fc46d2fff28d373f46f8ca439c919
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/namespace_expressions/__pycache__/struct_namespace.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/list_namespace.py b/lib/python3.12/site-packages/ray/data/namespace_expressions/list_namespace.py
new file mode 100644
index 0000000000000000000000000000000000000000..f44b04ac8e79fce71cf436dd8f030d1625d95a37
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/namespace_expressions/list_namespace.py
@@ -0,0 +1,122 @@
+"""List namespace for expression operations on list-typed columns."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Union
+
+import pyarrow
+import pyarrow.compute as pc
+
+from ray.data.datatype import DataType
+from ray.data.expressions import pyarrow_udf
+
+if TYPE_CHECKING:
+    from ray.data.expressions import Expr, UDFExpr
+
+
+@dataclass
+class _ListNamespace:
+    """Namespace for list operations on expression columns.
+
+    This namespace provides methods for operating on list-typed columns using
+    PyArrow compute functions.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> # Get length of list column
+        >>> expr = col("items").list.len()
+        >>> # Get first item using method
+        >>> expr = col("items").list.get(0)
+        >>> # Get first item using indexing
+        >>> expr = col("items").list[0]
+        >>> # Slice list
+        >>> expr = col("items").list[1:3]
+    """
+
+    _expr: Expr
+
+    def len(self) -> "UDFExpr":
+        """Get the length of each list."""
+
+        @pyarrow_udf(return_dtype=DataType.int32())
+        def _list_len(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.list_value_length(arr)
+
+        return _list_len(self._expr)
+
+    def __getitem__(self, key: Union[int, slice]) -> "UDFExpr":
+        """Get element or slice using bracket notation.
+
+        Args:
+            key: An integer for element access or slice for list slicing.
+
+        Returns:
+            UDFExpr that extracts the element or slice.
+
+        Example:
+            >>> col("items").list[0]      # Get first item  # doctest: +SKIP
+            >>> col("items").list[1:3]    # Get slice [1, 3)  # doctest: +SKIP
+            >>> col("items").list[-1]     # Get last item  # doctest: +SKIP
+        """
+        if isinstance(key, int):
+            return self.get(key)
+        elif isinstance(key, slice):
+            return self.slice(key.start, key.stop, key.step)
+        else:
+            raise TypeError(
+                f"List indices must be integers or slices, not {type(key).__name__}"
+            )
+
+    def get(self, index: int) -> "UDFExpr":
+        """Get element at the specified index from each list.
+
+        Args:
+            index: The index of the element to retrieve. Negative indices are supported.
+
+        Returns:
+            UDFExpr that extracts the element at the given index.
+        """
+        # Infer return type from the list's value type
+        return_dtype = DataType(object)  # fallback
+        if self._expr.data_type.is_arrow_type():
+            arrow_type = self._expr.data_type.to_arrow_dtype()
+            if pyarrow.types.is_list(arrow_type) or pyarrow.types.is_large_list(
+                arrow_type
+            ):
+                return_dtype = DataType.from_arrow(arrow_type.value_type)
+            elif pyarrow.types.is_fixed_size_list(arrow_type):
+                return_dtype = DataType.from_arrow(arrow_type.value_type)
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _list_get(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.list_element(arr, index)
+
+        return _list_get(self._expr)
+
+    def slice(
+        self, start: int | None = None, stop: int | None = None, step: int | None = None
+    ) -> "UDFExpr":
+        """Slice each list.
+
+        Args:
+            start: Start index (inclusive). Defaults to 0.
+            stop: Stop index (exclusive). Defaults to list length.
+            step: Step size. Defaults to 1.
+
+        Returns:
+            UDFExpr that extracts a slice from each list.
+        """
+        # Return type is the same as the input list type
+        return_dtype = self._expr.data_type
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _list_slice(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.list_slice(
+                arr,
+                start=0 if start is None else start,
+                stop=stop,
+                step=1 if step is None else step,
+            )
+
+        return _list_slice(self._expr)
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/string_namespace.py b/lib/python3.12/site-packages/ray/data/namespace_expressions/string_namespace.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a177621a2f8e563ca5b5f21288bac5a7f170bba
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/namespace_expressions/string_namespace.py
@@ -0,0 +1,355 @@
+"""String namespace for expression operations on string-typed columns."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Literal
+
+import pyarrow
+import pyarrow.compute as pc
+
+from ray.data.datatype import DataType
+from ray.data.expressions import pyarrow_udf
+
+if TYPE_CHECKING:
+    from ray.data.expressions import Expr, UDFExpr
+
+
+def _create_str_udf(
+    pc_func: Callable[..., pyarrow.Array], return_dtype: DataType
+) -> Callable[..., "UDFExpr"]:
+    """Helper to create a string UDF that wraps a PyArrow compute function.
+
+    This helper handles all types of PyArrow compute operations:
+    - Unary operations (no args): upper(), lower(), reverse()
+    - Pattern operations (pattern + args): starts_with(), contains()
+    - Multi-argument operations: replace(), replace_slice()
+
+    Args:
+        pc_func: PyArrow compute function that takes (array, *positional, **kwargs)
+        return_dtype: The return data type
+
+    Returns:
+        A callable that creates UDFExpr instances
+    """
+
+    def wrapper(expr: Expr, *positional: Any, **kwargs: Any) -> "UDFExpr":
+        @pyarrow_udf(return_dtype=return_dtype)
+        def udf(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc_func(arr, *positional, **kwargs)
+
+        return udf(expr)
+
+    return wrapper
+
+
+@dataclass
+class _StringNamespace:
+    """Namespace for string operations on expression columns.
+
+    This namespace provides methods for operating on string-typed columns using
+    PyArrow compute functions.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> # Convert to uppercase
+        >>> expr = col("name").str.upper()
+        >>> # Get string length
+        >>> expr = col("name").str.len()
+        >>> # Check if string starts with a prefix
+        >>> expr = col("name").str.starts_with("A")
+    """
+
+    _expr: Expr
+
+    # Length methods
+    def len(self) -> "UDFExpr":
+        """Get the length of each string in characters."""
+        return _create_str_udf(pc.utf8_length, DataType.int32())(self._expr)
+
+    def byte_len(self) -> "UDFExpr":
+        """Get the length of each string in bytes."""
+        return _create_str_udf(pc.binary_length, DataType.int32())(self._expr)
+
+    # Case methods
+    def upper(self) -> "UDFExpr":
+        """Convert strings to uppercase."""
+        return _create_str_udf(pc.utf8_upper, DataType.string())(self._expr)
+
+    def lower(self) -> "UDFExpr":
+        """Convert strings to lowercase."""
+        return _create_str_udf(pc.utf8_lower, DataType.string())(self._expr)
+
+    def capitalize(self) -> "UDFExpr":
+        """Capitalize the first character of each string."""
+        return _create_str_udf(pc.utf8_capitalize, DataType.string())(self._expr)
+
+    def title(self) -> "UDFExpr":
+        """Convert strings to title case."""
+        return _create_str_udf(pc.utf8_title, DataType.string())(self._expr)
+
+    def swapcase(self) -> "UDFExpr":
+        """Swap the case of each character."""
+        return _create_str_udf(pc.utf8_swapcase, DataType.string())(self._expr)
+
+    # Predicate methods
+    def is_alpha(self) -> "UDFExpr":
+        """Check if strings contain only alphabetic characters."""
+        return _create_str_udf(pc.utf8_is_alpha, DataType.bool())(self._expr)
+
+    def is_alnum(self) -> "UDFExpr":
+        """Check if strings contain only alphanumeric characters."""
+        return _create_str_udf(pc.utf8_is_alnum, DataType.bool())(self._expr)
+
+    def is_digit(self) -> "UDFExpr":
+        """Check if strings contain only digits."""
+        return _create_str_udf(pc.utf8_is_digit, DataType.bool())(self._expr)
+
+    def is_decimal(self) -> "UDFExpr":
+        """Check if strings contain only decimal characters."""
+        return _create_str_udf(pc.utf8_is_decimal, DataType.bool())(self._expr)
+
+    def is_numeric(self) -> "UDFExpr":
+        """Check if strings contain only numeric characters."""
+        return _create_str_udf(pc.utf8_is_numeric, DataType.bool())(self._expr)
+
+    def is_space(self) -> "UDFExpr":
+        """Check if strings contain only whitespace."""
+        return _create_str_udf(pc.utf8_is_space, DataType.bool())(self._expr)
+
+    def is_lower(self) -> "UDFExpr":
+        """Check if strings are lowercase."""
+        return _create_str_udf(pc.utf8_is_lower, DataType.bool())(self._expr)
+
+    def is_upper(self) -> "UDFExpr":
+        """Check if strings are uppercase."""
+        return _create_str_udf(pc.utf8_is_upper, DataType.bool())(self._expr)
+
+    def is_title(self) -> "UDFExpr":
+        """Check if strings are title-cased."""
+        return _create_str_udf(pc.utf8_is_title, DataType.bool())(self._expr)
+
+    def is_printable(self) -> "UDFExpr":
+        """Check if strings contain only printable characters."""
+        return _create_str_udf(pc.utf8_is_printable, DataType.bool())(self._expr)
+
+    def is_ascii(self) -> "UDFExpr":
+        """Check if strings contain only ASCII characters."""
+        return _create_str_udf(pc.string_is_ascii, DataType.bool())(self._expr)
+
+    # Searching methods
+    def starts_with(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Check if strings start with a pattern."""
+        return _create_str_udf(pc.starts_with, DataType.bool())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def ends_with(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Check if strings end with a pattern."""
+        return _create_str_udf(pc.ends_with, DataType.bool())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def contains(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Check if strings contain a substring."""
+        return _create_str_udf(pc.match_substring, DataType.bool())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def match(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Match strings against a SQL LIKE pattern."""
+        return _create_str_udf(pc.match_like, DataType.bool())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def find(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Find the first occurrence of a substring."""
+        return _create_str_udf(pc.find_substring, DataType.int32())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def count(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Count occurrences of a substring."""
+        return _create_str_udf(pc.count_substring, DataType.int32())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def find_regex(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Find the first occurrence matching a regex pattern."""
+        return _create_str_udf(pc.find_substring_regex, DataType.int32())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def count_regex(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Count occurrences matching a regex pattern."""
+        return _create_str_udf(pc.count_substring_regex, DataType.int32())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def match_regex(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Check if strings match a regex pattern."""
+        return _create_str_udf(pc.match_substring_regex, DataType.bool())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    # Transformation methods
+    def reverse(self) -> "UDFExpr":
+        """Reverse each string."""
+        return _create_str_udf(pc.utf8_reverse, DataType.string())(self._expr)
+
+    def slice(self, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Slice strings by codeunit indices."""
+        return _create_str_udf(pc.utf8_slice_codeunits, DataType.string())(
+            self._expr, *args, **kwargs
+        )
+
+    def replace(
+        self, pattern: str, replacement: str, *args: Any, **kwargs: Any
+    ) -> "UDFExpr":
+        """Replace occurrences of a substring."""
+        return _create_str_udf(pc.replace_substring, DataType.string())(
+            self._expr, pattern, replacement, *args, **kwargs
+        )
+
+    def replace_regex(
+        self, pattern: str, replacement: str, *args: Any, **kwargs: Any
+    ) -> "UDFExpr":
+        """Replace occurrences matching a regex pattern."""
+        return _create_str_udf(pc.replace_substring_regex, DataType.string())(
+            self._expr, pattern, replacement, *args, **kwargs
+        )
+
+    def replace_slice(
+        self, start: int, stop: int, replacement: str, *args: Any, **kwargs: Any
+    ) -> "UDFExpr":
+        """Replace a slice with a string."""
+        return _create_str_udf(pc.binary_replace_slice, DataType.string())(
+            self._expr, start, stop, replacement, *args, **kwargs
+        )
+
+    def split(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Split strings by a pattern."""
+        return _create_str_udf(pc.split_pattern, DataType(object))(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def split_regex(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Split strings by a regex pattern."""
+        return _create_str_udf(pc.split_pattern_regex, DataType(object))(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def split_whitespace(self, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Split strings on whitespace."""
+        return _create_str_udf(pc.utf8_split_whitespace, DataType(object))(
+            self._expr, *args, **kwargs
+        )
+
+    def extract(self, pattern: str, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Extract a substring matching a regex pattern."""
+        return _create_str_udf(pc.extract_regex, DataType.string())(
+            self._expr, pattern, *args, **kwargs
+        )
+
+    def repeat(self, n: int, *args: Any, **kwargs: Any) -> "UDFExpr":
+        """Repeat each string n times."""
+        return _create_str_udf(pc.binary_repeat, DataType.string())(
+            self._expr, n, *args, **kwargs
+        )
+
+    def center(
+        self, width: int, padding: str = " ", *args: Any, **kwargs: Any
+    ) -> "UDFExpr":
+        """Center strings in a field of given width."""
+        return _create_str_udf(pc.utf8_center, DataType.string())(
+            self._expr, width, padding, *args, **kwargs
+        )
+
+    # Custom methods that need special logic beyond simple PyArrow function calls
+    def strip(self, characters: str | None = None) -> "UDFExpr":
+        """Remove leading and trailing whitespace or specified characters.
+
+        Args:
+            characters: Characters to remove. If None, removes whitespace.
+
+        Returns:
+            UDFExpr that strips characters from both ends.
+        """
+
+        @pyarrow_udf(return_dtype=DataType.string())
+        def _str_strip(arr: pyarrow.Array) -> pyarrow.Array:
+            if characters is None:
+                return pc.utf8_trim_whitespace(arr)
+            else:
+                return pc.utf8_trim(arr, characters=characters)
+
+        return _str_strip(self._expr)
+
+    def lstrip(self, characters: str | None = None) -> "UDFExpr":
+        """Remove leading whitespace or specified characters.
+
+        Args:
+            characters: Characters to remove. If None, removes whitespace.
+
+        Returns:
+            UDFExpr that strips characters from the left.
+        """
+
+        @pyarrow_udf(return_dtype=DataType.string())
+        def _str_lstrip(arr: pyarrow.Array) -> pyarrow.Array:
+            if characters is None:
+                return pc.utf8_ltrim_whitespace(arr)
+            else:
+                return pc.utf8_ltrim(arr, characters=characters)
+
+        return _str_lstrip(self._expr)
+
+    def rstrip(self, characters: str | None = None) -> "UDFExpr":
+        """Remove trailing whitespace or specified characters.
+
+        Args:
+            characters: Characters to remove. If None, removes whitespace.
+
+        Returns:
+            UDFExpr that strips characters from the right.
+        """
+
+        @pyarrow_udf(return_dtype=DataType.string())
+        def _str_rstrip(arr: pyarrow.Array) -> pyarrow.Array:
+            if characters is None:
+                return pc.utf8_rtrim_whitespace(arr)
+            else:
+                return pc.utf8_rtrim(arr, characters=characters)
+
+        return _str_rstrip(self._expr)
+
+    # Padding
+    def pad(
+        self,
+        width: int,
+        fillchar: str = " ",
+        side: Literal["left", "right", "both"] = "right",
+    ) -> "UDFExpr":
+        """Pad strings to a specified width.
+
+        Args:
+            width: Target width.
+            fillchar: Character to use for padding.
+            side: "left", "right", or "both" for padding side.
+
+        Returns:
+            UDFExpr that pads strings.
+        """
+
+        @pyarrow_udf(return_dtype=DataType.string())
+        def _str_pad(arr: pyarrow.Array) -> pyarrow.Array:
+            if side == "right":
+                return pc.utf8_rpad(arr, width=width, padding=fillchar)
+            elif side == "left":
+                return pc.utf8_lpad(arr, width=width, padding=fillchar)
+            elif side == "both":
+                return pc.utf8_center(arr, width=width, padding=fillchar)
+            else:
+                raise ValueError("side must be 'left', 'right', or 'both'")
+
+        return _str_pad(self._expr)
diff --git a/lib/python3.12/site-packages/ray/data/namespace_expressions/struct_namespace.py b/lib/python3.12/site-packages/ray/data/namespace_expressions/struct_namespace.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2c6570cb7e2462a98f1365c6d45869ffd7dcb09
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/namespace_expressions/struct_namespace.py
@@ -0,0 +1,77 @@
+"""Struct namespace for expression operations on struct-typed columns."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import pyarrow
+import pyarrow.compute as pc
+
+from ray.data.datatype import DataType
+from ray.data.expressions import pyarrow_udf
+
+if TYPE_CHECKING:
+    from ray.data.expressions import Expr, UDFExpr
+
+
+@dataclass
+class _StructNamespace:
+    """Namespace for struct operations on expression columns.
+
+    This namespace provides methods for operating on struct-typed columns using
+    PyArrow compute functions.
+
+    Example:
+        >>> from ray.data.expressions import col
+        >>> # Access a field using method
+        >>> expr = col("user_record").struct.field("age")
+        >>> # Access a field using bracket notation
+        >>> expr = col("user_record").struct["age"]
+        >>> # Access nested field
+        >>> expr = col("user_record").struct["address"].struct["city"]
+    """
+
+    _expr: Expr
+
+    def __getitem__(self, field_name: str) -> "UDFExpr":
+        """Extract a field using bracket notation.
+
+        Args:
+            field_name: The name of the field to extract.
+
+        Returns:
+            UDFExpr that extracts the specified field from each struct.
+
+        Example:
+            >>> col("user").struct["age"]  # Get age field  # doctest: +SKIP
+            >>> col("user").struct["address"].struct["city"]  # Get nested city field  # doctest: +SKIP
+        """
+        return self.field(field_name)
+
+    def field(self, field_name: str) -> "UDFExpr":
+        """Extract a field from a struct.
+
+        Args:
+            field_name: The name of the field to extract.
+
+        Returns:
+            UDFExpr that extracts the specified field from each struct.
+        """
+        # Infer return type from the struct field type
+        return_dtype = DataType(object)  # fallback
+        if self._expr.data_type.is_arrow_type():
+            arrow_type = self._expr.data_type.to_arrow_dtype()
+            if pyarrow.types.is_struct(arrow_type):
+                try:
+                    field_type = arrow_type.field(field_name).type
+                    return_dtype = DataType.from_arrow(field_type)
+                except KeyError:
+                    # Field not found in schema, fallback to object
+                    pass
+
+        @pyarrow_udf(return_dtype=return_dtype)
+        def _struct_field(arr: pyarrow.Array) -> pyarrow.Array:
+            return pc.struct_field(arr, field_name)
+
+        return _struct_field(self._expr)
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__init__.py b/lib/python3.12/site-packages/ray/data/preprocessors/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef99a58ee113934da0162c86945c1275cb85455b
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/__init__.py
@@ -0,0 +1,50 @@
+from ray.data.preprocessors.chain import Chain
+from ray.data.preprocessors.concatenator import Concatenator
+from ray.data.preprocessors.discretizer import (
+    CustomKBinsDiscretizer,
+    UniformKBinsDiscretizer,
+)
+from ray.data.preprocessors.encoder import (
+    Categorizer,
+    LabelEncoder,
+    MultiHotEncoder,
+    OneHotEncoder,
+    OrdinalEncoder,
+)
+from ray.data.preprocessors.hasher import FeatureHasher
+from ray.data.preprocessors.imputer import SimpleImputer
+from ray.data.preprocessors.normalizer import Normalizer
+from ray.data.preprocessors.scaler import (
+    MaxAbsScaler,
+    MinMaxScaler,
+    RobustScaler,
+    StandardScaler,
+)
+from ray.data.preprocessors.tokenizer import Tokenizer
+from ray.data.preprocessors.torch import TorchVisionPreprocessor
+from ray.data.preprocessors.transformer import PowerTransformer
+from ray.data.preprocessors.vectorizer import CountVectorizer, HashingVectorizer
+
+__all__ = [
+    "Categorizer",
+    "CountVectorizer",
+    "Chain",
+    "FeatureHasher",
+    "HashingVectorizer",
+    "LabelEncoder",
+    "MaxAbsScaler",
+    "MinMaxScaler",
+    "MultiHotEncoder",
+    "Normalizer",
+    "OneHotEncoder",
+    "OrdinalEncoder",
+    "PowerTransformer",
+    "RobustScaler",
+    "SimpleImputer",
+    "StandardScaler",
+    "Concatenator",
+    "Tokenizer",
+    "TorchVisionPreprocessor",
+    "CustomKBinsDiscretizer",
+    "UniformKBinsDiscretizer",
+]
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bb10611747ff4672622de4d6cb8b6c9d39e6a259
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/chain.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/chain.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..98466fa5527ca5df736d1fab2873b73ea2df465d
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/chain.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/concatenator.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/concatenator.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4dddcdffcbfbec1351094fb626f790f63782f6d1
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/concatenator.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/discretizer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/discretizer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6e18a99057fb8d9bcfecff924bfd5f71e048e8ee
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/discretizer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/encoder.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/encoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1f7de5c3d5e69b78c9e060dbbfbe99181875e2f4
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/encoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/hasher.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/hasher.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6832fa4612f958b08969d97c8c52895f0d81f8c0
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/hasher.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/imputer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/imputer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b5d8dfe7371ffb4d50520f70fedab2c17e619d0c
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/imputer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/normalizer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/normalizer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..220eecb6a2ca5c01c950076f249d1f803be4a22b
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/normalizer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/scaler.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/scaler.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f2c8fcd43cf8b200abdef6fbb703d00ac8192e44
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/scaler.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/serialization_handlers.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/serialization_handlers.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..56616bf8f215adca25f55dbed04a4cefcd3838cf
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/serialization_handlers.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/tokenizer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/tokenizer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6de21a9b649eda1ff19c16b5ce013426349465b2
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/tokenizer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/torch.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/torch.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..13f52da9c6f8f7944d15f60625e88a7119ea9e15
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/torch.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/transformer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/transformer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..20fec45ab7e6222ef15fde2ad06a8a3401860222
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/transformer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e64a7d21042c657e5ae8fd2526d0a25a6bd629e1
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/vectorizer.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/vectorizer.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c7c83415502f02afb17d2fd4a7f31d73eeeacfff
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/vectorizer.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/version_support.cpython-312.pyc b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/version_support.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2ebfe72261e70d824cd02c68863c1e9a019165c7
Binary files /dev/null and b/lib/python3.12/site-packages/ray/data/preprocessors/__pycache__/version_support.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/chain.py b/lib/python3.12/site-packages/ray/data/preprocessors/chain.py
new file mode 100644
index 0000000000000000000000000000000000000000..bfe53ca06ac3a06a917902ddde203e3364dac774
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/chain.py
@@ -0,0 +1,115 @@
+from typing import TYPE_CHECKING, Optional
+
+from ray.air.util.data_batch_conversion import BatchFormat
+from ray.data.preprocessor import Preprocessor
+
+if TYPE_CHECKING:
+    from ray.air.data_batch_type import DataBatchType
+    from ray.data.dataset import Dataset
+
+
+class Chain(Preprocessor):
+    """Combine multiple preprocessors into a single :py:class:`Preprocessor`.
+
+    When you call ``fit``, each preprocessor is fit on the dataset produced by the
+    preceeding preprocessor's ``fit_transform``.
+
+    Example:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import *
+        >>>
+        >>> df = pd.DataFrame({
+        ...     "X0": [0, 1, 2],
+        ...     "X1": [3, 4, 5],
+        ...     "Y": ["orange", "blue", "orange"],
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>>
+        >>> preprocessor = Chain(
+        ...     StandardScaler(columns=["X0", "X1"]),
+        ...     Concatenator(columns=["X0", "X1"], output_column_name="X"),
+        ...     LabelEncoder(label_column="Y")
+        ... )
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           Y                                         X
+        0  1  [-1.224744871391589, -1.224744871391589]
+        1  0                                [0.0, 0.0]
+        2  1    [1.224744871391589, 1.224744871391589]
+
+    Args:
+        preprocessors: The preprocessors to sequentially compose.
+    """
+
+    def fit_status(self):
+        fittable_count = 0
+        fitted_count = 0
+        for p in self.preprocessors:
+            if p.fit_status() == Preprocessor.FitStatus.FITTED:
+                fittable_count += 1
+                fitted_count += 1
+            elif p.fit_status() in (
+                Preprocessor.FitStatus.NOT_FITTED,
+                Preprocessor.FitStatus.PARTIALLY_FITTED,
+            ):
+                fittable_count += 1
+            else:
+                assert p.fit_status() == Preprocessor.FitStatus.NOT_FITTABLE
+        if fittable_count > 0:
+            if fitted_count == fittable_count:
+                return Preprocessor.FitStatus.FITTED
+            elif fitted_count > 0:
+                return Preprocessor.FitStatus.PARTIALLY_FITTED
+            else:
+                return Preprocessor.FitStatus.NOT_FITTED
+        else:
+            return Preprocessor.FitStatus.NOT_FITTABLE
+
+    def __init__(self, *preprocessors: Preprocessor):
+        super().__init__()
+        self.preprocessors = preprocessors
+
+    def _fit(self, ds: "Dataset") -> Preprocessor:
+        for preprocessor in self.preprocessors[:-1]:
+            ds = preprocessor.fit_transform(ds)
+        self.preprocessors[-1].fit(ds)
+        return self
+
+    def fit_transform(self, ds: "Dataset") -> "Dataset":
+        for preprocessor in self.preprocessors:
+            ds = preprocessor.fit_transform(ds)
+        return ds
+
+    def _transform(
+        self,
+        ds: "Dataset",
+        batch_size: Optional[int],
+        num_cpus: Optional[float] = None,
+        memory: Optional[float] = None,
+        concurrency: Optional[int] = None,
+    ) -> "Dataset":
+        for preprocessor in self.preprocessors:
+            ds = preprocessor.transform(
+                ds,
+                batch_size=batch_size,
+                num_cpus=num_cpus,
+                memory=memory,
+                concurrency=concurrency,
+            )
+        return ds
+
+    def _transform_batch(self, df: "DataBatchType") -> "DataBatchType":
+        for preprocessor in self.preprocessors:
+            df = preprocessor.transform_batch(df)
+        return df
+
+    def __repr__(self):
+        arguments = ", ".join(repr(preprocessor) for preprocessor in self.preprocessors)
+        return f"{self.__class__.__name__}({arguments})"
+
+    def _determine_transform_to_use(self) -> BatchFormat:
+        # This is relevant for BatchPrediction.
+        # For Chain preprocessor, we picked the first one as entry point.
+        # TODO (jiaodong): We should revisit if our Chain preprocessor is
+        # still optimal with context of lazy execution.
+        return self.preprocessors[0]._determine_transform_to_use()
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/concatenator.py b/lib/python3.12/site-packages/ray/data/preprocessors/concatenator.py
new file mode 100644
index 0000000000000000000000000000000000000000..acb4f95abeeb64656bdcbd2a5ff405ca6f4122cd
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/concatenator.py
@@ -0,0 +1,162 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+
+from ray.data.preprocessor import Preprocessor
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+class Concatenator(Preprocessor):
+    """Combine numeric columns into a column of type
+    :class:`~ray.air.util.tensor_extensions.pandas.TensorDtype`. Only columns
+    specified in ``columns`` will be concatenated.
+
+    This preprocessor concatenates numeric columns and stores the result in a new
+    column. The new column contains
+    :class:`~ray.air.util.tensor_extensions.pandas.TensorArrayElement` objects of
+    shape :math:`(m,)`, where :math:`m` is the number of columns concatenated.
+    The :math:`m` concatenated columns are dropped after concatenation.
+    The preprocessor preserves the order of the columns provided in the ``colummns``
+    argument and will use that order when calling ``transform()`` and ``transform_batch()``.
+
+    Examples:
+        >>> import numpy as np
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import Concatenator
+
+        :py:class:`Concatenator` combines numeric columns into a column of
+        :py:class:`~ray.air.util.tensor_extensions.pandas.TensorDtype`.
+
+        >>> df = pd.DataFrame({"X0": [0, 3, 1], "X1": [0.5, 0.2, 0.9]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> concatenator = Concatenator(columns=["X0", "X1"])
+        >>> concatenator.transform(ds).to_pandas()  # doctest: +SKIP
+           concat_out
+        0  [0.0, 0.5]
+        1  [3.0, 0.2]
+        2  [1.0, 0.9]
+
+        By default, the created column is called `"concat_out"`, but you can specify
+        a different name.
+
+        >>> concatenator = Concatenator(columns=["X0", "X1"], output_column_name="tensor")
+        >>> concatenator.transform(ds).to_pandas()  # doctest: +SKIP
+               tensor
+        0  [0.0, 0.5]
+        1  [3.0, 0.2]
+        2  [1.0, 0.9]
+
+        >>> concatenator = Concatenator(columns=["X0", "X1"], dtype=np.float32)
+        >>> concatenator.transform(ds)  # doctest: +SKIP
+        Dataset(num_rows=3, schema={Y: object, concat_out: TensorDtype(shape=(2,), dtype=float32)})
+
+        When ``flatten=True``, nested vectors in the columns will be flattened during concatenation:
+
+        >>> df = pd.DataFrame({"X0": [[1, 2], [3, 4]], "X1": [0.5, 0.2]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> concatenator = Concatenator(columns=["X0", "X1"], flatten=True)
+        >>> concatenator.transform(ds).to_pandas()  # doctest: +SKIP
+           concat_out
+        0  [1.0, 2.0, 0.5]
+        1  [3.0, 4.0, 0.2]
+
+    Args:
+        columns: A list of columns to concatenate. The provided order of the columns
+             will be retained during concatenation.
+        output_column_name: The desired name for the new column.
+            Defaults to ``"concat_out"``.
+        dtype: The ``dtype`` to convert the output tensors to. If unspecified,
+            the ``dtype`` is determined by standard coercion rules.
+        raise_if_missing: If ``True``, an error is raised if any
+            of the columns in ``columns`` don't exist.
+            Defaults to ``False``.
+        flatten: If ``True``, nested vectors in the columns will be flattened during
+            concatenation. Defaults to ``False``.
+
+    Raises:
+        ValueError: if `raise_if_missing` is `True` and a column in `columns` or
+            doesn't exist in the dataset.
+    """  # noqa: E501
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        output_column_name: str = "concat_out",
+        dtype: Optional[np.dtype] = None,
+        raise_if_missing: bool = False,
+        flatten: bool = False,
+    ):
+        self.columns = columns
+
+        self.output_column_name = output_column_name
+        self.dtype = dtype
+        self.raise_if_missing = raise_if_missing
+        self.flatten = flatten
+
+    def _validate(self, df: pd.DataFrame) -> None:
+        missing_columns = set(self.columns) - set(df)
+        if missing_columns:
+            message = (
+                f"Missing columns specified in '{self.columns}': {missing_columns}"
+            )
+            if self.raise_if_missing:
+                raise ValueError(message)
+            else:
+                logger.warning(message)
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        self._validate(df)
+
+        if self.flatten:
+            concatenated = df[self.columns].to_numpy()
+            concatenated = [
+                np.concatenate(
+                    [
+                        np.atleast_1d(elem)
+                        if self.dtype is None
+                        else np.atleast_1d(elem).astype(self.dtype)
+                        for elem in row
+                    ]
+                )
+                for row in concatenated
+            ]
+        else:
+            concatenated = df[self.columns].to_numpy(dtype=self.dtype)
+
+        df = df.drop(columns=self.columns)
+        # Use a Pandas Series for column assignment to get more consistent
+        # behavior across Pandas versions.
+        df.loc[:, self.output_column_name] = pd.Series(list(concatenated))
+        return df
+
+    def __repr__(self):
+        default_values = {
+            "output_column_name": "concat_out",
+            "columns": None,
+            "dtype": None,
+            "raise_if_missing": False,
+            "flatten": False,
+        }
+
+        non_default_arguments = []
+        for parameter, default_value in default_values.items():
+            value = getattr(self, parameter)
+            if value != default_value:
+                non_default_arguments.append(f"{parameter}={value}")
+
+        return f"{self.__class__.__name__}({', '.join(non_default_arguments)})"
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        super().__setstate__(state)
+        # flatten is a recent field, to ensure backwards compatibility
+        # assign a default in case it is missing in the serialized state
+        if not hasattr(self, "flatten"):
+            self.flatten = False
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/discretizer.py b/lib/python3.12/site-packages/ray/data/preprocessors/discretizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..a19f0c2c4bd2666556131ab1361166192b52d340
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/discretizer.py
@@ -0,0 +1,427 @@
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Type, Union
+
+import numpy as np
+import pandas as pd
+
+from ray.data.aggregate import Max, Min
+from ray.data.preprocessor import Preprocessor
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Dataset
+
+
+class _AbstractKBinsDiscretizer(Preprocessor):
+    """Abstract base class for all KBinsDiscretizers.
+
+    Essentially a thin wraper around ``pd.cut``.
+
+    Expects either ``self.stats_`` or ``self.bins`` to be set and
+    contain {column:list_of_bin_intervals}.
+    """
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def bin_values(s: pd.Series) -> pd.Series:
+            if s.name not in self.columns:
+                return s
+            labels = self.dtypes.get(s.name) if self.dtypes else False
+            ordered = True
+            if labels:
+                if isinstance(labels, pd.CategoricalDtype):
+                    ordered = labels.ordered
+                    labels = list(labels.categories)
+                else:
+                    labels = False
+
+            bins = self.stats_ if self._is_fittable else self.bins
+            return pd.cut(
+                s,
+                bins[s.name] if isinstance(bins, dict) else bins,
+                right=self.right,
+                labels=labels,
+                ordered=ordered,
+                retbins=False,
+                include_lowest=self.include_lowest,
+                duplicates=self.duplicates,
+            )
+
+        binned_df = df.apply(bin_values, axis=0)
+        df[self.output_columns] = binned_df[self.columns]
+        return df
+
+    def _validate_bins_columns(self):
+        if isinstance(self.bins, dict) and not all(
+            col in self.bins for col in self.columns
+        ):
+            raise ValueError(
+                "If `bins` is a dictionary, all elements of `columns` must be present "
+                "in it."
+            )
+
+    def __repr__(self):
+        attr_str = ", ".join(
+            [
+                f"{attr_name}={attr_value!r}"
+                for attr_name, attr_value in vars(self).items()
+                if not attr_name.startswith("_")
+            ]
+        )
+        return f"{self.__class__.__name__}({attr_str})"
+
+
+@PublicAPI(stability="alpha")
+class CustomKBinsDiscretizer(_AbstractKBinsDiscretizer):
+    """Bin values into discrete intervals using custom bin edges.
+
+    Columns must contain numerical values.
+
+    Examples:
+        Use :class:`CustomKBinsDiscretizer` to bin continuous features.
+
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import CustomKBinsDiscretizer
+        >>> df = pd.DataFrame({
+        ...     "value_1": [0.2, 1.4, 2.5, 6.2, 9.7, 2.1],
+        ...     "value_2": [10, 15, 13, 12, 23, 25],
+        ... })
+        >>> ds = ray.data.from_pandas(df)
+        >>> discretizer = CustomKBinsDiscretizer(
+        ...     columns=["value_1", "value_2"],
+        ...     bins=[0, 1, 4, 10, 25]
+        ... )
+        >>> discretizer.transform(ds).to_pandas()
+           value_1  value_2
+        0        0        2
+        1        1        3
+        2        1        3
+        3        2        3
+        4        2        3
+        5        1        3
+
+        :class:`CustomKBinsDiscretizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> discretizer = CustomKBinsDiscretizer(
+        ...     columns=["value_1", "value_2"],
+        ...     bins=[0, 1, 4, 10, 25],
+        ...     output_columns=["value_1_discretized", "value_2_discretized"]
+        ... )
+        >>> discretizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           value_1  value_2  value_1_discretized  value_2_discretized
+        0      0.2       10                    0                    2
+        1      1.4       15                    1                    3
+        2      2.5       13                    1                    3
+        3      6.2       12                    2                    3
+        4      9.7       23                    2                    3
+        5      2.1       25                    1                    3
+
+        You can also specify different bin edges per column.
+
+        >>> discretizer = CustomKBinsDiscretizer(
+        ...     columns=["value_1", "value_2"],
+        ...     bins={"value_1": [0, 1, 4], "value_2": [0, 18, 35, 70]},
+        ... )
+        >>> discretizer.transform(ds).to_pandas()
+           value_1  value_2
+        0      0.0        0
+        1      1.0        0
+        2      1.0        0
+        3      NaN        0
+        4      NaN        1
+        5      1.0        1
+
+
+    Args:
+        columns: The columns to discretize.
+        bins: Defines custom bin edges. Can be an iterable of numbers,
+            a ``pd.IntervalIndex``, or a dict mapping columns to either of them.
+            Note that ``pd.IntervalIndex`` for bins must be non-overlapping.
+        right: Indicates whether bins include the rightmost edge.
+        include_lowest: Indicates whether the first interval should be left-inclusive.
+        duplicates: Can be either 'raise' or 'drop'. If bin edges are not unique,
+            raise ``ValueError`` or drop non-uniques.
+        dtypes: An optional dictionary that maps columns to ``pd.CategoricalDtype``
+            objects or ``np.integer`` types. If you don't include a column in ``dtypes``
+            or specify it as an integer dtype, the outputted column will consist of
+            ordered integers corresponding to bins. If you use a
+            ``pd.CategoricalDtype``, the outputted column will be a
+            ``pd.CategoricalDtype`` with the categories being mapped to bins.
+            You can use ``pd.CategoricalDtype(categories, ordered=True)`` to
+            preserve information about bin order.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    .. seealso::
+
+        :class:`UniformKBinsDiscretizer`
+            If you want to bin data into uniform width bins.
+    """
+
+    def __init__(
+        self,
+        columns: List[str],
+        bins: Union[
+            Iterable[float],
+            pd.IntervalIndex,
+            Dict[str, Union[Iterable[float], pd.IntervalIndex]],
+        ],
+        *,
+        right: bool = True,
+        include_lowest: bool = False,
+        duplicates: str = "raise",
+        dtypes: Optional[
+            Dict[str, Union[pd.CategoricalDtype, Type[np.integer]]]
+        ] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        self.columns = columns
+        self.bins = bins
+        self.right = right
+        self.include_lowest = include_lowest
+        self.duplicates = duplicates
+        self.dtypes = dtypes
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+        self._validate_bins_columns()
+
+    _is_fittable = False
+
+
+@PublicAPI(stability="alpha")
+class UniformKBinsDiscretizer(_AbstractKBinsDiscretizer):
+    """Bin values into discrete intervals (bins) of uniform width.
+
+    Columns must contain numerical values.
+
+    Examples:
+        Use :class:`UniformKBinsDiscretizer` to bin continuous features.
+
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import UniformKBinsDiscretizer
+        >>> df = pd.DataFrame({
+        ...     "value_1": [0.2, 1.4, 2.5, 6.2, 9.7, 2.1],
+        ...     "value_2": [10, 15, 13, 12, 23, 25],
+        ... })
+        >>> ds = ray.data.from_pandas(df)
+        >>> discretizer = UniformKBinsDiscretizer(
+        ...     columns=["value_1", "value_2"], bins=4
+        ... )
+        >>> discretizer.fit_transform(ds).to_pandas()
+           value_1  value_2
+        0        0        0
+        1        0        1
+        2        0        0
+        3        2        0
+        4        3        3
+        5        0        3
+
+        :class:`UniformKBinsDiscretizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> discretizer = UniformKBinsDiscretizer(
+        ...     columns=["value_1", "value_2"],
+        ...     bins=4,
+        ...     output_columns=["value_1_discretized", "value_2_discretized"]
+        ... )
+        >>> discretizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           value_1  value_2  value_1_discretized  value_2_discretized
+        0      0.2       10                    0                    0
+        1      1.4       15                    0                    1
+        2      2.5       13                    0                    0
+        3      6.2       12                    2                    0
+        4      9.7       23                    3                    3
+        5      2.1       25                    0                    3
+
+        You can also specify different number of bins per column.
+
+        >>> discretizer = UniformKBinsDiscretizer(
+        ...     columns=["value_1", "value_2"], bins={"value_1": 4, "value_2": 3}
+        ... )
+        >>> discretizer.fit_transform(ds).to_pandas()
+           value_1  value_2
+        0        0        0
+        1        0        0
+        2        0        0
+        3        2        0
+        4        3        2
+        5        0        2
+
+
+    Args:
+        columns: The columns to discretize.
+        bins: Defines the number of equal-width bins.
+            Can be either an integer (which will be applied to all columns),
+            or a dict that maps columns to integers.
+            The range is extended by .1% on each side to include
+            the minimum and maximum values.
+        right: Indicates whether bins includes the rightmost edge or not.
+        include_lowest: Whether the first interval should be left-inclusive
+            or not.
+        duplicates: Can be either 'raise' or 'drop'. If bin edges are not unique,
+            raise ``ValueError`` or drop non-uniques.
+        dtypes: An optional dictionary that maps columns to ``pd.CategoricalDtype``
+            objects or ``np.integer`` types. If you don't include a column in ``dtypes``
+            or specify it as an integer dtype, the outputted column will consist of
+            ordered integers corresponding to bins. If you use a
+            ``pd.CategoricalDtype``, the outputted column will be a
+            ``pd.CategoricalDtype`` with the categories being mapped to bins.
+            You can use ``pd.CategoricalDtype(categories, ordered=True)`` to
+            preserve information about bin order.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    .. seealso::
+
+        :class:`CustomKBinsDiscretizer`
+            If you want to specify your own bin edges.
+    """
+
+    def __init__(
+        self,
+        columns: List[str],
+        bins: Union[int, Dict[str, int]],
+        *,
+        right: bool = True,
+        include_lowest: bool = False,
+        duplicates: str = "raise",
+        dtypes: Optional[
+            Dict[str, Union[pd.CategoricalDtype, Type[np.integer]]]
+        ] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        self.columns = columns
+        self.bins = bins
+        self.right = right
+        self.include_lowest = include_lowest
+        self.duplicates = duplicates
+        self.dtypes = dtypes
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        self._validate_on_fit()
+
+        if isinstance(self.bins, dict):
+            columns = self.bins.keys()
+        else:
+            columns = self.columns
+
+        for column in columns:
+            bins = self.bins[column] if isinstance(self.bins, dict) else self.bins
+            if not isinstance(bins, int):
+                raise TypeError(
+                    f"`bins` must be an integer or a dict of integers, got {bins}"
+                )
+
+        self.stat_computation_plan.add_aggregator(
+            aggregator_fn=Min,
+            columns=columns,
+        )
+        self.stat_computation_plan.add_aggregator(
+            aggregator_fn=Max,
+            columns=columns,
+        )
+
+        return self
+
+    def _validate_on_fit(self):
+        self._validate_bins_columns()
+
+    def _fit_execute(self, dataset: "Dataset"):
+        stats = self.stat_computation_plan.compute(dataset)
+        self.stats_ = post_fit_processor(stats, self.bins, self.right)
+        return self
+
+
+def post_fit_processor(aggregate_stats: dict, bins: Union[str, Dict], right: bool):
+    mins, maxes, stats = {}, {}, {}
+    for key, value in aggregate_stats.items():
+        column_name = key[4:-1]  # min(column) -> column
+        if key.startswith("min"):
+            mins[column_name] = value
+        if key.startswith("max"):
+            maxes[column_name] = value
+
+    for column in mins.keys():
+        stats[column] = _translate_min_max_number_of_bins_to_bin_edges(
+            mn=mins[column],
+            mx=maxes[column],
+            bins=bins[column] if isinstance(bins, dict) else bins,
+            right=right,
+        )
+
+    return stats
+
+
+# Copied from
+# https://github.com/pandas-dev/pandas/blob/v1.4.4/pandas/core/reshape/tile.py#L257
+# under
+# BSD 3-Clause License
+#
+# Copyright (c) 2008-2011, AQR Capital Management, LLC, Lambda Foundry, Inc.
+# and PyData Development Team
+# All rights reserved.
+#
+# Copyright (c) 2011-2022, Open source contributors.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+def _translate_min_max_number_of_bins_to_bin_edges(
+    mn: float, mx: float, bins: int, right: bool
+) -> List[float]:
+    """Translates a range and desired number of bins into list of bin edges."""
+    rng = (mn, mx)
+    mn, mx = (mi + 0.0 for mi in rng)
+
+    if np.isinf(mn) or np.isinf(mx):
+        raise ValueError(
+            "Cannot specify integer `bins` when input data contains infinity."
+        )
+    elif mn == mx:  # adjust end points before binning
+        mn -= 0.001 * abs(mn) if mn != 0 else 0.001
+        mx += 0.001 * abs(mx) if mx != 0 else 0.001
+        bins = np.linspace(mn, mx, bins + 1, endpoint=True)
+    else:  # adjust end points after binning
+        bins = np.linspace(mn, mx, bins + 1, endpoint=True)
+        adj = (mx - mn) * 0.001  # 0.1% of the range
+        if right:
+            bins[0] -= adj
+        else:
+            bins[-1] += adj
+    return bins
+
+
+# TODO(ml-team)
+# Add QuantileKBinsDiscretizer
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/encoder.py b/lib/python3.12/site-packages/ray/data/preprocessors/encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1d112bc00e6d471854fac90e1448b2eb27379f0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/encoder.py
@@ -0,0 +1,948 @@
+import logging
+from collections import Counter
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Hashable, List, Optional, Set
+
+import numpy as np
+import pandas as pd
+import pandas.api.types
+
+from ray.air.util.data_batch_conversion import BatchFormat
+from ray.data.preprocessor import (
+    Preprocessor,
+    PreprocessorNotFittedException,
+    SerializablePreprocessorBase,
+)
+from ray.data.preprocessors.utils import make_post_processor
+from ray.data.preprocessors.version_support import SerializablePreprocessor
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Dataset
+
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.ordinal_encoder")
+class OrdinalEncoder(SerializablePreprocessorBase):
+    r"""Encode values within columns as ordered integer values.
+
+    :class:`OrdinalEncoder` encodes categorical features as integers that range from
+    :math:`0` to :math:`n - 1`, where :math:`n` is the number of categories.
+
+    If you transform a value that isn't in the fitted datset, then the value is encoded
+    as ``float("nan")``.
+
+    Columns must contain either hashable values or lists of hashable values. Also, you
+    can't have both scalars and lists in the same column.
+
+    Examples:
+        Use :class:`OrdinalEncoder` to encode categorical features as integers.
+
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import OrdinalEncoder
+        >>> df = pd.DataFrame({
+        ...     "sex": ["male", "female", "male", "female"],
+        ...     "level": ["L4", "L5", "L3", "L4"],
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> encoder = OrdinalEncoder(columns=["sex", "level"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           sex  level
+        0    1      1
+        1    0      2
+        2    1      0
+        3    0      1
+
+        :class:`OrdinalEncoder` can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> encoder = OrdinalEncoder(columns=["sex", "level"], output_columns=["sex_encoded", "level_encoded"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+              sex level  sex_encoded  level_encoded
+        0    male    L4            1              1
+        1  female    L5            0              2
+        2    male    L3            1              0
+        3  female    L4            0              1
+
+
+        If you transform a value not present in the original dataset, then the value
+        is encoded as ``float("nan")``.
+
+        >>> df = pd.DataFrame({"sex": ["female"], "level": ["L6"]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> encoder.transform(ds).to_pandas()  # doctest: +SKIP
+           sex  level
+        0    0    NaN
+
+        :class:`OrdinalEncoder` can also encode categories in a list.
+
+        >>> df = pd.DataFrame({
+        ...     "name": ["Shaolin Soccer", "Moana", "The Smartest Guys in the Room"],
+        ...     "genre": [
+        ...         ["comedy", "action", "sports"],
+        ...         ["animation", "comedy",  "action"],
+        ...         ["documentary"],
+        ...     ],
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> encoder = OrdinalEncoder(columns=["genre"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                                    name      genre
+        0                 Shaolin Soccer  [2, 0, 4]
+        1                          Moana  [1, 2, 0]
+        2  The Smartest Guys in the Room        [3]
+
+    Args:
+        columns: The columns to separately encode.
+        encode_lists: If ``True``, encode list elements.  If ``False``, encode
+            whole lists (i.e., replace each list with an integer). ``True``
+            by default.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    .. seealso::
+
+        :class:`OneHotEncoder`
+            Another preprocessor that encodes categorical data.
+    """
+
+    def __init__(
+        self,
+        columns: List[str],
+        *,
+        encode_lists: bool = True,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        # TODO: allow user to specify order of values within each column.
+        self.columns = columns
+        self.encode_lists = encode_lists
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        self.stat_computation_plan.add_callable_stat(
+            stat_fn=lambda key_gen: compute_unique_value_indices(
+                dataset=dataset,
+                columns=self.columns,
+                encode_lists=self.encode_lists,
+                key_gen=key_gen,
+            ),
+            post_process_fn=unique_post_fn(),
+            stat_key_fn=lambda col: f"unique({col})",
+            post_key_fn=lambda col: f"unique_values({col})",
+            columns=self.columns,
+        )
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        _validate_df(df, *self.columns)
+
+        def encode_list(element: list, *, name: str):
+            return [self.stats_[f"unique_values({name})"].get(x) for x in element]
+
+        def column_ordinal_encoder(s: pd.Series):
+            if _is_series_composed_of_lists(s):
+                if self.encode_lists:
+                    return s.map(partial(encode_list, name=s.name))
+
+                def list_as_category(element):
+                    key = tuple(element)
+                    return self.stats_[f"unique_values({s.name})"].get(key)
+
+                return s.apply(list_as_category)
+
+            s_values = self.stats_[f"unique_values({s.name})"]
+            return s.map(s_values)
+
+        df[self.output_columns] = df[self.columns].apply(column_ordinal_encoder)
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "encode_lists": self.encode_lists,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        self.encode_lists = fields["encode_lists"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"encode_lists={self.encode_lists!r}, "
+            f"output_columns={self.output_columns!r})"
+        )
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.one_hot_encoder")
+class OneHotEncoder(SerializablePreprocessorBase):
+    r"""`One-hot encode <https://en.wikipedia.org/wiki/One-hot#Machine_learning_and_statistics>`_
+    categorical data.
+
+    This preprocessor transforms each specified column into a one-hot encoded vector.
+    Each element in the vector corresponds to a unique category in the column, with a
+    value of 1 if the category matches and 0 otherwise.
+
+    If a category is infrequent (based on ``max_categories``) or not present in the
+    fitted dataset, it is encoded as all 0s.
+
+    Columns must contain hashable objects or lists of hashable objects.
+
+    .. note::
+        Lists are treated as categories. If you want to encode individual list
+        elements, use :class:`MultiHotEncoder`.
+
+    Example:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import OneHotEncoder
+        >>>
+        >>> df = pd.DataFrame({"color": ["red", "green", "red", "red", "blue", "green"]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> encoder = OneHotEncoder(columns=["color"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+               color
+        0  [0, 0, 1]
+        1  [0, 1, 0]
+        2  [0, 0, 1]
+        3  [0, 0, 1]
+        4  [1, 0, 0]
+        5  [0, 1, 0]
+
+        OneHotEncoder can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> encoder = OneHotEncoder(columns=["color"], output_columns=["color_encoded"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           color color_encoded
+        0    red     [0, 0, 1]
+        1  green     [0, 1, 0]
+        2    red     [0, 0, 1]
+        3    red     [0, 0, 1]
+        4   blue     [1, 0, 0]
+        5  green     [0, 1, 0]
+
+        If you one-hot encode a value that isn't in the fitted dataset, then the
+        value is encoded with zeros.
+
+        >>> df = pd.DataFrame({"color": ["yellow"]})
+        >>> batch = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> encoder.transform(batch).to_pandas()  # doctest: +SKIP
+            color color_encoded
+        0  yellow     [0, 0, 0]
+
+        Likewise, if you one-hot encode an infrequent value, then the value is encoded
+        with zeros.
+
+        >>> encoder = OneHotEncoder(columns=["color"], max_categories={"color": 2})
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+            color
+        0  [1, 0]
+        1  [0, 1]
+        2  [1, 0]
+        3  [1, 0]
+        4  [0, 0]
+        5  [0, 1]
+
+    Args:
+        columns: The columns to separately encode.
+        max_categories: The maximum number of features to create for each column.
+            If a value isn't specified for a column, then a feature is created
+            for every category in that column.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    .. seealso::
+
+        :class:`MultiHotEncoder`
+            If you want to encode individual list elements, use
+            :class:`MultiHotEncoder`.
+
+        :class:`OrdinalEncoder`
+            If your categories are ordered, you may want to use
+            :class:`OrdinalEncoder`.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        columns: List[str],
+        *,
+        max_categories: Optional[Dict[str, int]] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        # TODO: add `drop` parameter.
+        self.columns = columns
+        self.max_categories = max_categories or {}
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        self.stat_computation_plan.add_callable_stat(
+            stat_fn=lambda key_gen: compute_unique_value_indices(
+                dataset=dataset,
+                columns=self.columns,
+                encode_lists=False,
+                key_gen=key_gen,
+                max_categories=self.max_categories,
+            ),
+            post_process_fn=unique_post_fn(),
+            stat_key_fn=lambda col: f"unique({col})",
+            post_key_fn=lambda col: f"unique_values({col})",
+            columns=self.columns,
+        )
+        return self
+
+    def safe_get(self, v: Any, stats: Dict[str, int]):
+        if isinstance(v, (list, np.ndarray)):
+            v = tuple(v)
+        if isinstance(v, Hashable):
+            return stats.get(v, -1)
+        else:
+            return -1  # Unhashable type treated as a missing category
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        _validate_df(df, *self.columns)
+
+        # Compute new one-hot encoded columns
+        for column, output_column in zip(self.columns, self.output_columns):
+            stats = self.stats_[f"unique_values({column})"]
+            num_categories = len(stats)
+            one_hot = np.zeros((len(df), num_categories), dtype=np.uint8)
+            # Integer indices for each category in the column
+            codes = df[column].apply(lambda v: self.safe_get(v, stats)).to_numpy()
+            # Filter to only the rows that have a valid category
+            valid_category_mask = codes != -1
+            # Dimension should be (num_rows, ) - 1D boolean array
+            non_zero_indices = np.nonzero(valid_category_mask)[0]
+            # Mark the corresponding categories as 1
+            one_hot[
+                non_zero_indices,
+                codes[valid_category_mask],
+            ] = 1
+            df[output_column] = one_hot.tolist()
+
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "max_categories": self.max_categories,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        self.max_categories = fields["max_categories"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"max_categories={self.max_categories!r}, "
+            f"output_columns={self.output_columns!r})"
+        )
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(
+    version=1, identifier="io.ray.preprocessors.multi_hot_encoder"
+)
+class MultiHotEncoder(SerializablePreprocessorBase):
+    r"""Multi-hot encode categorical data.
+
+    This preprocessor replaces each list of categories with an :math:`m`-length binary
+    list, where :math:`m` is the number of unique categories in the column or the value
+    specified in ``max_categories``. The :math:`i\\text{-th}` element of the binary list
+    is :math:`1` if category :math:`i` is in the input list and :math:`0` otherwise.
+
+    Columns must contain hashable objects or lists of hashable objects.
+    Also, you can't have both types in the same column.
+
+    .. note::
+        The logic is similar to scikit-learn's [MultiLabelBinarizer][1]
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import MultiHotEncoder
+        >>>
+        >>> df = pd.DataFrame({
+        ...     "name": ["Shaolin Soccer", "Moana", "The Smartest Guys in the Room"],
+        ...     "genre": [
+        ...         ["comedy", "action", "sports"],
+        ...         ["animation", "comedy",  "action"],
+        ...         ["documentary"],
+        ...     ],
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>>
+        >>> encoder = MultiHotEncoder(columns=["genre"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                                    name            genre
+        0                 Shaolin Soccer  [1, 0, 1, 0, 1]
+        1                          Moana  [1, 1, 1, 0, 0]
+        2  The Smartest Guys in the Room  [0, 0, 0, 1, 0]
+
+        :class:`MultiHotEncoder` can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> encoder = MultiHotEncoder(columns=["genre"], output_columns=["genre_encoded"])
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                                    name                        genre    genre_encoded
+        0                 Shaolin Soccer     [comedy, action, sports]  [1, 0, 1, 0, 1]
+        1                          Moana  [animation, comedy, action]  [1, 1, 1, 0, 0]
+        2  The Smartest Guys in the Room                [documentary]  [0, 0, 0, 1, 0]
+
+        If you specify ``max_categories``, then :class:`MultiHotEncoder`
+        creates features for only the most frequent categories.
+
+        >>> encoder = MultiHotEncoder(columns=["genre"], max_categories={"genre": 3})
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                                    name      genre
+        0                 Shaolin Soccer  [1, 1, 1]
+        1                          Moana  [1, 1, 0]
+        2  The Smartest Guys in the Room  [0, 0, 0]
+        >>> encoder.stats_  # doctest: +SKIP
+        OrderedDict([('unique_values(genre)', {'comedy': 0, 'action': 1, 'sports': 2})])
+
+    Args:
+        columns: The columns to separately encode.
+        max_categories: The maximum number of features to create for each column.
+            If a value isn't specified for a column, then a feature is created
+            for every unique category in that column.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    .. seealso::
+
+        :class:`OneHotEncoder`
+            If you're encoding individual categories instead of lists of
+            categories, use :class:`OneHotEncoder`.
+
+        :class:`OrdinalEncoder`
+            If your categories are ordered, you may want to use
+            :class:`OrdinalEncoder`.
+
+    [1]: https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.MultiLabelBinarizer.html
+    """
+
+    def __init__(
+        self,
+        columns: List[str],
+        *,
+        max_categories: Optional[Dict[str, int]] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        # TODO: add `drop` parameter.
+        self.columns = columns
+        self.max_categories = max_categories or {}
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        self.stat_computation_plan.add_callable_stat(
+            stat_fn=lambda key_gen: compute_unique_value_indices(
+                dataset=dataset,
+                columns=self.columns,
+                encode_lists=True,
+                key_gen=key_gen,
+                max_categories=self.max_categories,
+            ),
+            post_process_fn=unique_post_fn(),
+            stat_key_fn=lambda col: f"unique({col})",
+            post_key_fn=lambda col: f"unique_values({col})",
+            columns=self.columns,
+        )
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        _validate_df(df, *self.columns)
+
+        def encode_list(element: list, *, name: str):
+            if isinstance(element, np.ndarray):
+                element = element.tolist()
+            elif not isinstance(element, list):
+                element = [element]
+            stats = self.stats_[f"unique_values({name})"]
+            counter = Counter(element)
+            return [counter.get(x, 0) for x in stats]
+
+        for column, output_column in zip(self.columns, self.output_columns):
+            df[output_column] = df[column].map(partial(encode_list, name=column))
+
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "max_categories": self.max_categories,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        self.max_categories = fields["max_categories"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"max_categories={self.max_categories!r}, "
+            f"output_columns={self.output_columns})"
+        )
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.label_encoder")
+class LabelEncoder(SerializablePreprocessorBase):
+    r"""Encode labels as integer targets.
+
+    :class:`LabelEncoder` encodes labels as integer targets that range from
+    :math:`0` to :math:`n - 1`, where :math:`n` is the number of unique labels.
+
+    If you transform a label that isn't in the fitted datset, then the label is encoded
+    as ``float("nan")``.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> df = pd.DataFrame({
+        ...     "sepal_width": [5.1, 7, 4.9, 6.2],
+        ...     "sepal_height": [3.5, 3.2, 3, 3.4],
+        ...     "species": ["setosa", "versicolor", "setosa", "virginica"]
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>>
+        >>> from ray.data.preprocessors import LabelEncoder
+        >>> encoder = LabelEncoder(label_column="species")
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           sepal_width  sepal_height  species
+        0          5.1           3.5        0
+        1          7.0           3.2        1
+        2          4.9           3.0        0
+        3          6.2           3.4        2
+
+        You can also provide the name of the output column that should hold the encoded
+        labels if you want to use :class:`LabelEncoder` in append mode.
+
+        >>> encoder = LabelEncoder(label_column="species", output_column="species_encoded")
+        >>> encoder.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           sepal_width  sepal_height     species  species_encoded
+        0          5.1           3.5      setosa                0
+        1          7.0           3.2  versicolor                1
+        2          4.9           3.0      setosa                0
+        3          6.2           3.4   virginica                2
+
+        If you transform a label not present in the original dataset, then the new
+        label is encoded as ``float("nan")``.
+
+        >>> df = pd.DataFrame({
+        ...     "sepal_width": [4.2],
+        ...     "sepal_height": [2.7],
+        ...     "species": ["bracteata"]
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> encoder.transform(ds).to_pandas()  # doctest: +SKIP
+           sepal_width  sepal_height  species
+        0          4.2           2.7      NaN
+
+    Args:
+        label_column: A column containing labels that you want to encode.
+        output_column: The name of the column that will contain the encoded
+            labels. If None, the output column will have the same name as the
+            input column.
+
+    .. seealso::
+
+        :class:`OrdinalEncoder`
+            If you're encoding ordered features, use :class:`OrdinalEncoder` instead of
+            :class:`LabelEncoder`.
+    """
+
+    def __init__(self, label_column: str, *, output_column: Optional[str] = None):
+        super().__init__()
+        self.label_column = label_column
+        self.output_column = output_column or label_column
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        self.stat_computation_plan.add_callable_stat(
+            stat_fn=lambda key_gen: compute_unique_value_indices(
+                dataset=dataset,
+                columns=[self.label_column],
+                key_gen=key_gen,
+            ),
+            post_process_fn=unique_post_fn(),
+            stat_key_fn=lambda col: f"unique({col})",
+            post_key_fn=lambda col: f"unique_values({col})",
+            columns=[self.label_column],
+        )
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        _validate_df(df, self.label_column)
+
+        def column_label_encoder(s: pd.Series):
+            s_values = self.stats_[f"unique_values({s.name})"]
+            return s.map(s_values)
+
+        df[self.output_column] = df[self.label_column].transform(column_label_encoder)
+        return df
+
+    def inverse_transform(self, ds: "Dataset") -> "Dataset":
+        """Inverse transform the given dataset.
+
+        Args:
+            ds: Input Dataset that has been fitted and/or transformed.
+
+        Returns:
+            ray.data.Dataset: The inverse transformed Dataset.
+
+        Raises:
+            PreprocessorNotFittedException: if ``fit`` is not called yet.
+        """
+
+        fit_status = self.fit_status()
+
+        if fit_status in (
+            Preprocessor.FitStatus.PARTIALLY_FITTED,
+            Preprocessor.FitStatus.NOT_FITTED,
+        ):
+            raise PreprocessorNotFittedException(
+                "`fit` must be called before `inverse_transform`, "
+            )
+
+        kwargs = self._get_transform_config()
+
+        return ds.map_batches(
+            self._inverse_transform_pandas, batch_format=BatchFormat.PANDAS, **kwargs
+        )
+
+    def _inverse_transform_pandas(self, df: pd.DataFrame):
+        def column_label_decoder(s: pd.Series):
+            inverse_values = {
+                value: key
+                for key, value in self.stats_[
+                    f"unique_values({self.label_column})"
+                ].items()
+            }
+            return s.map(inverse_values)
+
+        df[self.label_column] = df[self.output_column].transform(column_label_decoder)
+        return df
+
+    def get_input_columns(self) -> List[str]:
+        return [self.label_column]
+
+    def get_output_columns(self) -> List[str]:
+        return [self.output_column]
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "label_column": self.label_column,
+            "output_column": self.output_column,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.label_column = fields["label_column"]
+        self.output_column = fields["output_column"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(label_column={self.label_column!r}, output_column={self.output_column!r})"
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.categorizer")
+class Categorizer(SerializablePreprocessorBase):
+    r"""Convert columns to ``pd.CategoricalDtype``.
+
+    Use this preprocessor with frameworks that have built-in support for
+    ``pd.CategoricalDtype`` like LightGBM.
+
+    .. warning::
+
+        If you don't specify ``dtypes``, fit this preprocessor before splitting
+        your dataset into train and test splits. This ensures categories are
+        consistent across splits.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import Categorizer
+        >>>
+        >>> df = pd.DataFrame(
+        ... {
+        ...     "sex": ["male", "female", "male", "female"],
+        ...     "level": ["L4", "L5", "L3", "L4"],
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> categorizer = Categorizer(columns=["sex", "level"])
+        >>> categorizer.fit_transform(ds).schema().types  # doctest: +SKIP
+        [CategoricalDtype(categories=['female', 'male'], ordered=False), CategoricalDtype(categories=['L3', 'L4', 'L5'], ordered=False)]
+
+        :class:`Categorizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the categorized values.
+
+        >>> categorizer = Categorizer(columns=["sex", "level"], output_columns=["sex_cat", "level_cat"])
+        >>> categorizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+              sex level sex_cat level_cat
+        0    male    L4    male        L4
+        1  female    L5  female        L5
+        2    male    L3    male        L3
+        3  female    L4  female        L4
+
+        If you know the categories in advance, you can specify the categories with the
+        ``dtypes`` parameter.
+
+        >>> categorizer = Categorizer(
+        ...     columns=["sex", "level"],
+        ...     dtypes={"level": pd.CategoricalDtype(["L3", "L4", "L5", "L6"], ordered=True)},
+        ... )
+        >>> categorizer.fit_transform(ds).schema().types  # doctest: +SKIP
+        [CategoricalDtype(categories=['female', 'male'], ordered=False), CategoricalDtype(categories=['L3', 'L4', 'L5', 'L6'], ordered=True)]
+
+    Args:
+        columns: The columns to convert to ``pd.CategoricalDtype``.
+        dtypes: An optional dictionary that maps columns to ``pd.CategoricalDtype``
+            objects. If you don't include a column in ``dtypes``, the categories
+            are inferred.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        columns: List[str],
+        dtypes: Optional[Dict[str, pd.CategoricalDtype]] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        if not dtypes:
+            dtypes = {}
+
+        self.columns = columns
+        self.dtypes = dtypes
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        columns_to_get = [
+            column for column in self.columns if column not in self.dtypes
+        ]
+        self.stats_ |= self.dtypes
+        if not columns_to_get:
+            return self
+
+        def callback(unique_indices: Dict[str, Dict]) -> pd.CategoricalDtype:
+            return pd.CategoricalDtype(unique_indices.keys())
+
+        self.stat_computation_plan.add_callable_stat(
+            stat_fn=lambda key_gen: compute_unique_value_indices(
+                dataset=dataset,
+                columns=columns_to_get,
+                key_gen=key_gen,
+            ),
+            post_process_fn=make_post_processor(
+                base_fn=unique_post_fn(drop_na_values=True),
+                callbacks=[callback],
+            ),
+            stat_key_fn=lambda col: f"unique({col})",
+            post_key_fn=lambda col: col,
+            columns=columns_to_get,
+        )
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        df[self.output_columns] = df[self.columns].astype(self.stats_)
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "_fitted": getattr(self, "_fitted", None),
+            "dtypes": {
+                col: {"categories": list(dtype.categories), "ordered": dtype.ordered}
+                for col, dtype in self.dtypes.items()
+            }
+            if hasattr(self, "dtypes") and self.dtypes
+            else None,
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        # Handle dtypes field specially
+        self.dtypes = (
+            {
+                col: pd.CategoricalDtype(
+                    categories=dtype_data["categories"], ordered=dtype_data["ordered"]
+                )
+                for col, dtype_data in fields["dtypes"].items()
+            }
+            if fields.get("dtypes")
+            else {}
+        )
+
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"dtypes={self.dtypes!r}, output_columns={self.output_columns!r})"
+        )
+
+
+def compute_unique_value_indices(
+    *,
+    dataset: "Dataset",
+    columns: List[str],
+    key_gen: Callable,
+    encode_lists: bool = True,
+    max_categories: Optional[Dict[str, int]] = None,
+):
+    if max_categories is None:
+        max_categories = {}
+    columns_set = set(columns)
+    for column in max_categories:
+        if column not in columns_set:
+            raise ValueError(
+                f"You set `max_categories` for {column}, which is not present in "
+                f"{columns}."
+            )
+
+    def get_pd_value_counts_per_column(col: pd.Series) -> Dict:
+
+        # special handling for lists
+        if _is_series_composed_of_lists(col):
+            if encode_lists:
+                counter = Counter()
+
+                def update_counter(element):
+                    counter.update(element)
+                    return element
+
+                col.map(update_counter)
+                return counter
+            else:
+                # convert to tuples to make lists hashable
+                col = col.map(lambda x: tuple(x))
+        return Counter(col.value_counts(dropna=False).to_dict())
+
+    def get_pd_value_counts(df: pd.DataFrame) -> Dict[str, List[Dict]]:
+
+        df_columns = df.columns.tolist()
+        result = {}
+        for col in columns:
+            if col in df_columns:
+                result[col] = [get_pd_value_counts_per_column(df[col])]
+            else:
+                raise ValueError(
+                    f"Column '{col}' does not exist in DataFrame, which has columns: {df_columns}"  # noqa: E501
+                )
+        return result
+
+    value_counts_ds = dataset.map_batches(get_pd_value_counts, batch_format="pandas")
+    unique_values_by_col: Dict[str, Set] = {key_gen(col): set() for col in columns}
+    for batch in value_counts_ds.iter_batches(batch_size=None):
+        for col, counters in batch.items():
+            for counter in counters:
+                counter: Dict[Any, int] = {
+                    k: v for k, v in counter.items() if v is not None
+                }
+                if col in max_categories:
+                    counter: Dict[Any, int] = dict(
+                        Counter(counter).most_common(max_categories[col])
+                    )
+                # add only column values since frequencies are needed beyond this point
+                unique_values_by_col[key_gen(col)].update(counter.keys())
+
+    return unique_values_by_col
+
+
+def unique_post_fn(drop_na_values: bool = False) -> Callable[[Set], Dict[str, int]]:
+    """
+    Returns a post-processing function that generates an encoding map by
+    sorting the unique values produced during aggregation or stats computation.
+
+    Args:
+        drop_na_values: If True, NA/null values will be silently dropped from the
+            encoding map. If False, raises an error if any NA/null values are present.
+
+    Returns:
+        A callable that takes a set of unique values and returns a dictionary
+        mapping each value to a unique integer index.
+    """
+
+    def gen_value_index(values: Set) -> Dict[str, int]:
+        if drop_na_values:
+            values = {k for k in values if not pd.isnull(k)}
+        else:
+            if any(pd.isnull(k) for k in values):
+                raise ValueError(
+                    "Unable to fit column because it contains null"
+                    " values. Consider imputing missing values first."
+                )
+        return {k: j for j, k in enumerate(sorted(values))}
+
+    return gen_value_index
+
+
+def _validate_df(df: pd.DataFrame, *columns: str) -> None:
+    null_columns = [column for column in columns if df[column].isnull().values.any()]
+    if null_columns:
+        raise ValueError(
+            f"Unable to transform columns {null_columns} because they contain "
+            f"null values. Consider imputing missing values first."
+        )
+
+
+def _is_series_composed_of_lists(series: pd.Series) -> bool:
+    # we assume that all elements are a list here
+    first_not_none_element = next(
+        (element for element in series if element is not None), None
+    )
+    return pandas.api.types.is_object_dtype(series.dtype) and isinstance(
+        first_not_none_element, (list, np.ndarray)
+    )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/hasher.py b/lib/python3.12/site-packages/ray/data/preprocessors/hasher.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae9ae769ad9b5353693cbc3760a62eda529b41a0
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/hasher.py
@@ -0,0 +1,118 @@
+import collections
+from typing import List
+
+import pandas as pd
+
+from ray.data.preprocessor import Preprocessor
+from ray.data.preprocessors.utils import simple_hash
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class FeatureHasher(Preprocessor):
+    r"""Apply the `hashing trick <https://en.wikipedia.org/wiki/Feature_hashing>`_ to a
+    table that describes token frequencies.
+
+    :class:`FeatureHasher` creates ``num_features`` columns named ``hash_{index}``,
+    where ``index`` ranges from :math:`0` to ``num_features``:math:`- 1`. The column
+    ``hash_{index}`` describes the frequency of tokens that hash to ``index``.
+
+    Distinct tokens can correspond to the same index. However, if ``num_features`` is
+    large enough, then columns probably correspond to a unique token.
+
+    This preprocessor is memory efficient and quick to pickle. However, given a
+    transformed column, you can't know which tokens correspond to it. This might make it
+    hard to determine which tokens are important to your model.
+
+    .. warning::
+        Sparse matrices aren't supported. If you use a large ``num_features``, this
+        preprocessor might behave poorly.
+
+    Examples:
+
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import FeatureHasher
+
+        The data below describes the frequencies of tokens in ``"I like Python"`` and
+        ``"I dislike Python"``.
+
+        >>> df = pd.DataFrame({
+        ...     "I": [1, 1],
+        ...     "like": [1, 0],
+        ...     "dislike": [0, 1],
+        ...     "Python": [1, 1]
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+
+        :class:`FeatureHasher` hashes each token to determine its index. For example,
+        the index of ``"I"`` is :math:`hash(\\texttt{"I"}) \pmod 8 = 5`.
+
+        >>> hasher = FeatureHasher(columns=["I", "like", "dislike", "Python"], num_features=8, output_column = "hashed")
+        >>> hasher.fit_transform(ds)["hashed"].to_pandas().to_numpy()  # doctest: +SKIP
+        array([[0, 0, 0, 2, 0, 1, 0, 0],
+               [0, 0, 0, 1, 0, 1, 1, 0]])
+
+        Notice the hash collision: both ``"like"`` and ``"Python"`` correspond to index
+        :math:`3`. You can avoid hash collisions like these by increasing
+        ``num_features``.
+
+    Args:
+        columns: The columns to apply the hashing trick to. Each column should describe
+            the frequency of a token.
+        num_features: The number of features used to represent the vocabulary. You
+            should choose a value large enough to prevent hash collisions between
+            distinct tokens.
+        output_column: The name of the column that contains the hashed features.
+
+    .. seealso::
+        :class:`~ray.data.preprocessors.CountVectorizer`
+            Use this preprocessor to generate inputs for :class:`FeatureHasher`.
+
+        :class:`ray.data.preprocessors.HashingVectorizer`
+            If your input data describes documents rather than token frequencies,
+            use :class:`~ray.data.preprocessors.HashingVectorizer`.
+    """  # noqa: E501
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        num_features: int,
+        output_column: str,
+    ):
+        self.columns = columns
+        # TODO(matt): Set default number of features.
+        # This likely requires sparse matrix support to avoid explosion of columns.
+        self.num_features = num_features
+        self.output_column = output_column
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        # TODO(matt): Use sparse matrix for efficiency.
+        def row_feature_hasher(row):
+            hash_counts = collections.defaultdict(int)
+            for column in self.columns:
+                hashed_value = simple_hash(column, self.num_features)
+                hash_counts[hashed_value] += row[column]
+            return {f"hash_{i}": hash_counts[i] for i in range(self.num_features)}
+
+        feature_columns = df.loc[:, self.columns].apply(
+            row_feature_hasher, axis=1, result_type="expand"
+        )
+
+        # Concatenate the hash columns
+        hash_columns = [f"hash_{i}" for i in range(self.num_features)]
+        concatenated = feature_columns[hash_columns].to_numpy()
+        # Use a Pandas Series for column assignment to get more consistent
+        # behavior across Pandas versions.
+        df.loc[:, self.output_column] = pd.Series(list(concatenated))
+
+        return df
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"num_features={self.num_features!r}, "
+            f"output_column={self.output_column!r})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/imputer.py b/lib/python3.12/site-packages/ray/data/preprocessors/imputer.py
new file mode 100644
index 0000000000000000000000000000000000000000..4321f44ba89b07f5fb139c1f760326d0de1dc8d9
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/imputer.py
@@ -0,0 +1,255 @@
+import logging
+from collections import Counter
+from numbers import Number
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import is_categorical_dtype
+
+from ray.data.aggregate import Mean
+from ray.data.preprocessor import SerializablePreprocessorBase
+from ray.data.preprocessors.version_support import (
+    SerializablePreprocessor as Serializable,
+)
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Dataset
+
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+@Serializable(version=1, identifier="io.ray.preprocessors.simple_imputer")
+class SimpleImputer(SerializablePreprocessorBase):
+    """Replace missing values with imputed values. If the column is missing from a
+    batch, it will be filled with the imputed value.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import SimpleImputer
+        >>> df = pd.DataFrame({"X": [0, None, 3, 3], "Y": [None, "b", "c", "c"]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> ds.to_pandas()  # doctest: +SKIP
+             X     Y
+        0  0.0  None
+        1  NaN     b
+        2  3.0     c
+        3  3.0     c
+
+        The `"mean"` strategy imputes missing values with the mean of non-missing
+        values. This strategy doesn't work with categorical data.
+
+        >>> preprocessor = SimpleImputer(columns=["X"], strategy="mean")
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+             X     Y
+        0  0.0  None
+        1  2.0     b
+        2  3.0     c
+        3  3.0     c
+
+        The `"most_frequent"` strategy imputes missing values with the most frequent
+        value in each column.
+
+        >>> preprocessor = SimpleImputer(columns=["X", "Y"], strategy="most_frequent")
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+             X  Y
+        0  0.0  c
+        1  3.0  b
+        2  3.0  c
+        3  3.0  c
+
+        The `"constant"` strategy imputes missing values with the value specified by
+        `fill_value`.
+
+        >>> preprocessor = SimpleImputer(
+        ...     columns=["Y"],
+        ...     strategy="constant",
+        ...     fill_value="?",
+        ... )
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+             X  Y
+        0  0.0  ?
+        1  NaN  b
+        2  3.0  c
+        3  3.0  c
+
+        :class:`SimpleImputer` can also be used in append mode by providing the
+        name of the output_columns that should hold the imputed values.
+
+        >>> preprocessor = SimpleImputer(columns=["X"], output_columns=["X_imputed"], strategy="mean")
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+             X     Y  X_imputed
+        0  0.0  None        0.0
+        1  NaN     b        2.0
+        2  3.0     c        3.0
+        3  3.0     c        3.0
+
+    Args:
+        columns: The columns to apply imputation to.
+        strategy: How imputed values are chosen.
+
+            * ``"mean"``: The mean of non-missing values. This strategy only works with numeric columns.
+            * ``"most_frequent"``: The most common value.
+            * ``"constant"``: The value passed to ``fill_value``.
+
+        fill_value: The value to use when ``strategy`` is ``"constant"``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    Raises:
+        ValueError: if ``strategy`` is not ``"mean"``, ``"most_frequent"``, or
+            ``"constant"``.
+    """  # noqa: E501
+
+    _valid_strategies = ["mean", "most_frequent", "constant"]
+
+    def __init__(
+        self,
+        columns: List[str],
+        strategy: str = "mean",
+        fill_value: Optional[Union[str, Number]] = None,
+        *,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        self.columns = columns
+        self.strategy = strategy
+        self.fill_value = fill_value
+
+        if strategy not in self._valid_strategies:
+            raise ValueError(
+                f"Strategy {strategy} is not supported."
+                f"Supported values are: {self._valid_strategies}"
+            )
+
+        if strategy == "constant":
+            # There is no information to be fitted.
+            self._is_fittable = False
+            if fill_value is None:
+                raise ValueError(
+                    '`fill_value` must be set when using "constant" strategy.'
+                )
+
+        self.output_columns = (
+            SerializablePreprocessorBase._derive_and_validate_output_columns(
+                columns, output_columns
+            )
+        )
+
+    def _fit(self, dataset: "Dataset") -> SerializablePreprocessorBase:
+        if self.strategy == "mean":
+            self.stat_computation_plan.add_aggregator(
+                aggregator_fn=Mean, columns=self.columns
+            )
+        elif self.strategy == "most_frequent":
+            self.stat_computation_plan.add_callable_stat(
+                stat_fn=lambda key_gen: _get_most_frequent_values(
+                    dataset=dataset,
+                    columns=self.columns,
+                    key_gen=key_gen,
+                ),
+                stat_key_fn=lambda col: f"most_frequent({col})",
+                columns=self.columns,
+            )
+
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        for column, output_column in zip(self.columns, self.output_columns):
+            value = self._get_fill_value(column)
+
+            if value is None:
+                raise ValueError(
+                    f"Column {column} has no fill value. "
+                    "Check the data used to fit the SimpleImputer."
+                )
+
+            if column not in df.columns:
+                # Create the column with the fill_value if it doesn't exist
+                df[output_column] = value
+            else:
+                if is_categorical_dtype(df.dtypes[column]):
+                    df[output_column] = df[column].cat.add_categories([value])
+
+                if (
+                    output_column != column
+                    # If the backing array is memory-mapped from shared memory, then the
+                    # array won't be writeable.
+                    or (
+                        isinstance(df[output_column].values, np.ndarray)
+                        and not df[output_column].values.flags.writeable
+                    )
+                ):
+                    df[output_column] = df[column].copy(deep=True)
+
+                df.fillna({output_column: value}, inplace=True)
+
+        return df
+
+    def _get_fill_value(self, column):
+        if self.strategy == "mean":
+            return self.stats_[f"mean({column})"]
+        elif self.strategy == "most_frequent":
+            return self.stats_[f"most_frequent({column})"]
+        elif self.strategy == "constant":
+            return self.fill_value
+        else:
+            raise ValueError(
+                f"Strategy {self.strategy} is not supported. "
+                "Supported values are: {self._valid_strategies}"
+            )
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"strategy={self.strategy!r}, fill_value={self.fill_value!r}, "
+            f"output_columns={self.output_columns!r})"
+        )
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "_fitted": getattr(self, "_fitted", None),
+            "strategy": self.strategy,
+            "fill_value": getattr(self, "fill_value", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        self.strategy = fields["strategy"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+        self.fill_value = fields.get("fill_value")
+
+        if self.strategy == "constant":
+            self._is_fittable = False
+
+
+def _get_most_frequent_values(
+    dataset: "Dataset",
+    columns: List[str],
+    key_gen: Callable[[str], str],
+) -> Dict[str, Union[str, Number]]:
+    def get_pd_value_counts(df: pd.DataFrame) -> Dict[str, List[Counter]]:
+        return {col: [Counter(df[col].value_counts().to_dict())] for col in columns}
+
+    value_counts = dataset.map_batches(get_pd_value_counts, batch_format="pandas")
+    final_counters = {col: Counter() for col in columns}
+    for batch in value_counts.iter_batches(batch_size=None):
+        for col, counters in batch.items():
+            for counter in counters:
+                final_counters[col] += counter
+
+    return {
+        key_gen(column): final_counters[column].most_common(1)[0][0]  # noqa
+        for column in columns
+    }
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/normalizer.py b/lib/python3.12/site-packages/ray/data/preprocessors/normalizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..f94f22b8fcc16a305a4546298e4bbf1ea05773c3
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/normalizer.py
@@ -0,0 +1,131 @@
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+
+from ray.data.preprocessor import Preprocessor
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class Normalizer(Preprocessor):
+    r"""Scales each sample to have unit norm.
+
+    This preprocessor works by dividing each sample (i.e., row) by the sample's norm.
+    The general formula is given by
+
+    .. math::
+
+        s' = \frac{s}{\lVert s \rVert_p}
+
+    where :math:`s` is the sample, :math:`s'` is the transformed sample,
+    :math:\lVert s \rVert`, and :math:`p` is the norm type.
+
+    The following norms are supported:
+
+    * `"l1"` (:math:`L^1`): Sum of the absolute values.
+    * `"l2"` (:math:`L^2`): Square root of the sum of the squared values.
+    * `"max"` (:math:`L^\infty`): Maximum value.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import Normalizer
+        >>>
+        >>> df = pd.DataFrame({"X1": [1, 1], "X2": [1, 0], "X3": [0, 1]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> ds.to_pandas()  # doctest: +SKIP
+           X1  X2  X3
+        0   1   1   0
+        1   1   0   1
+
+        The :math:`L^2`-norm of the first sample is :math:`\sqrt{2}`, and the
+        :math:`L^2`-norm of the second sample is :math:`1`.
+
+        >>> preprocessor = Normalizer(columns=["X1", "X2"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                 X1        X2  X3
+        0  0.707107  0.707107   0
+        1  1.000000  0.000000   1
+
+        The :math:`L^1`-norm of the first sample is :math:`2`, and the
+        :math:`L^1`-norm of the second sample is :math:`1`.
+
+        >>> preprocessor = Normalizer(columns=["X1", "X2"], norm="l1")
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+            X1   X2  X3
+        0  0.5  0.5   0
+        1  1.0  0.0   1
+
+        The :math:`L^\infty`-norm of the both samples is :math:`1`.
+
+        >>> preprocessor = Normalizer(columns=["X1", "X2"], norm="max")
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+            X1   X2  X3
+        0  1.0  1.0   0
+        1  1.0  0.0   1
+
+        :class:`Normalizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the normalized values.
+
+        >>> preprocessor = Normalizer(columns=["X1", "X2"], output_columns=["X1_normalized", "X2_normalized"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2  X3  X1_normalized  X2_normalized
+        0   1   1   0       0.707107       0.707107
+        1   1   0   1       1.000000       0.000000
+
+    Args:
+        columns: The columns to scale. For each row, these colmumns are scaled to
+            unit-norm.
+        norm: The norm to use. The supported values are ``"l1"``, ``"l2"``, or
+            ``"max"``. Defaults to ``"l2"``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    Raises:
+        ValueError: if ``norm`` is not ``"l1"``, ``"l2"``, or ``"max"``.
+    """
+
+    _norm_fns = {
+        "l1": lambda cols: np.abs(cols).sum(axis=1),
+        "l2": lambda cols: np.sqrt(np.power(cols, 2).sum(axis=1)),
+        "max": lambda cols: np.max(abs(cols), axis=1),
+    }
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        norm="l2",
+        *,
+        output_columns: Optional[List[str]] = None,
+    ):
+        self.columns = columns
+        self.norm = norm
+
+        if norm not in self._norm_fns:
+            raise ValueError(
+                f"Norm {norm} is not supported."
+                f"Supported values are: {self._norm_fns.keys()}"
+            )
+
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        columns = df.loc[:, self.columns]
+        column_norms = self._norm_fns[self.norm](columns)
+
+        df[self.output_columns] = columns.div(column_norms, axis=0)
+        return df
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"norm={self.norm!r}, "
+            f"output_columns={self.output_columns!r})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/scaler.py b/lib/python3.12/site-packages/ray/data/preprocessors/scaler.py
new file mode 100644
index 0000000000000000000000000000000000000000..90cb0fffbb8f5b56cd5614d9c87ffcecfa5587fa
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/scaler.py
@@ -0,0 +1,519 @@
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+from ray.data.aggregate import AbsMax, ApproximateQuantile, Max, Mean, Min, Std
+from ray.data.preprocessor import Preprocessor, SerializablePreprocessorBase
+from ray.data.preprocessors.version_support import SerializablePreprocessor
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Dataset
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.standard_scaler")
+class StandardScaler(SerializablePreprocessorBase):
+    r"""Translate and scale each column by its mean and standard deviation,
+    respectively.
+
+    The general formula is given by
+
+    .. math::
+
+        x' = \frac{x - \bar{x}}{s}
+
+    where :math:`x` is the column, :math:`x'` is the transformed column,
+    :math:`\bar{x}` is the column average, and :math:`s` is the column's sample
+    standard deviation. If :math:`s = 0` (i.e., the column is constant-valued),
+    then the transformed column will contain zeros.
+
+    .. warning::
+        :class:`StandardScaler` works best when your data is normal. If your data isn't
+        approximately normal, then the transformed features won't be meaningful.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import StandardScaler
+        >>>
+        >>> df = pd.DataFrame({"X1": [-2, 0, 2], "X2": [-3, -3, 3], "X3": [1, 1, 1]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> ds.to_pandas()  # doctest: +SKIP
+           X1  X2  X3
+        0  -2  -3   1
+        1   0  -3   1
+        2   2   3   1
+
+        Columns are scaled separately.
+
+        >>> preprocessor = StandardScaler(columns=["X1", "X2"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                 X1        X2  X3
+        0 -1.224745 -0.707107   1
+        1  0.000000 -0.707107   1
+        2  1.224745  1.414214   1
+
+        Constant-valued columns get filled with zeros.
+
+        >>> preprocessor = StandardScaler(columns=["X3"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2   X3
+        0  -2  -3  0.0
+        1   0  -3  0.0
+        2   2   3  0.0
+
+        >>> preprocessor = StandardScaler(
+        ...     columns=["X1", "X2"],
+        ...     output_columns=["X1_scaled", "X2_scaled"]
+        ... )
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2  X3  X1_scaled  X2_scaled
+        0  -2  -3   1  -1.224745  -0.707107
+        1   0  -3   1   0.000000  -0.707107
+        2   2   3   1   1.224745   1.414214
+
+    Args:
+        columns: The columns to separately scale.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """
+
+    def __init__(self, columns: List[str], output_columns: Optional[List[str]] = None):
+        super().__init__()
+        self.columns = columns
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        self.stat_computation_plan.add_aggregator(
+            aggregator_fn=Mean,
+            columns=self.columns,
+        )
+        self.stat_computation_plan.add_aggregator(
+            aggregator_fn=lambda col: Std(col, ddof=0),
+            columns=self.columns,
+        )
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_standard_scaler(s: pd.Series):
+            s_mean = self.stats_[f"mean({s.name})"]
+            s_std = self.stats_[f"std({s.name})"]
+
+            if s_std is None or s_mean is None:
+                s[:] = np.nan
+                return s
+
+            # Handle division by zero.
+            # TODO: extend this to handle near-zero values.
+            if s_std == 0:
+                s_std = 1
+
+            return (s - s_mean) / s_std
+
+        df[self.output_columns] = df[self.columns].transform(column_standard_scaler)
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(columns={self.columns!r}, output_columns={self.output_columns!r})"
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.min_max_scaler")
+class MinMaxScaler(SerializablePreprocessorBase):
+    r"""Scale each column by its range.
+
+    The general formula is given by
+
+    .. math::
+
+        x' = \frac{x - \min(x)}{\max{x} - \min{x}}
+
+    where :math:`x` is the column and :math:`x'` is the transformed column. If
+    :math:`\max{x} - \min{x} = 0` (i.e., the column is constant-valued), then the
+    transformed column will get filled with zeros.
+
+    Transformed values are always in the range :math:`[0, 1]`.
+
+    .. tip::
+        This can be used as an alternative to :py:class:`StandardScaler`.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import MinMaxScaler
+        >>>
+        >>> df = pd.DataFrame({"X1": [-2, 0, 2], "X2": [-3, -3, 3], "X3": [1, 1, 1]})   # noqa: E501
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> ds.to_pandas()  # doctest: +SKIP
+           X1  X2  X3
+        0  -2  -3   1
+        1   0  -3   1
+        2   2   3   1
+
+        Columns are scaled separately.
+
+        >>> preprocessor = MinMaxScaler(columns=["X1", "X2"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+            X1   X2  X3
+        0  0.0  0.0   1
+        1  0.5  0.0   1
+        2  1.0  1.0   1
+
+        Constant-valued columns get filled with zeros.
+
+        >>> preprocessor = MinMaxScaler(columns=["X3"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2   X3
+        0  -2  -3  0.0
+        1   0  -3  0.0
+        2   2   3  0.0
+
+        >>> preprocessor = MinMaxScaler(columns=["X1", "X2"], output_columns=["X1_scaled", "X2_scaled"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2  X3  X1_scaled  X2_scaled
+        0  -2  -3   1        0.0        0.0
+        1   0  -3   1        0.5        0.0
+        2   2   3   1        1.0        1.0
+
+    Args:
+        columns: The columns to separately scale.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """
+
+    def __init__(self, columns: List[str], output_columns: Optional[List[str]] = None):
+        super().__init__()
+        self.columns = columns
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        aggregates = [Agg(col) for Agg in [Min, Max] for col in self.columns]
+        self.stats_ = dataset.aggregate(*aggregates)
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_min_max_scaler(s: pd.Series):
+            s_min = self.stats_[f"min({s.name})"]
+            s_max = self.stats_[f"max({s.name})"]
+            diff = s_max - s_min
+
+            # Handle division by zero.
+            # TODO: extend this to handle near-zero values.
+            if diff == 0:
+                diff = 1
+
+            return (s - s_min) / diff
+
+        df[self.output_columns] = df[self.columns].transform(column_min_max_scaler)
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(columns={self.columns!r}, output_columns={self.output_columns!r})"
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.max_abs_scaler")
+class MaxAbsScaler(SerializablePreprocessorBase):
+    r"""Scale each column by its absolute max value.
+
+    The general formula is given by
+
+    .. math::
+
+        x' = \frac{x}{\max{\vert x \vert}}
+
+    where :math:`x` is the column and :math:`x'` is the transformed column. If
+    :math:`\max{\vert x \vert} = 0` (i.e., the column contains all zeros), then the
+    column is unmodified.
+
+    .. tip::
+        This is the recommended way to scale sparse data. If you data isn't sparse,
+        you can use :class:`MinMaxScaler` or :class:`StandardScaler` instead.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import MaxAbsScaler
+        >>>
+        >>> df = pd.DataFrame({"X1": [-6, 3], "X2": [2, -4], "X3": [0, 0]})   # noqa: E501
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> ds.to_pandas()  # doctest: +SKIP
+           X1  X2  X3
+        0  -6   2   0
+        1   3  -4   0
+
+        Columns are scaled separately.
+
+        >>> preprocessor = MaxAbsScaler(columns=["X1", "X2"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+            X1   X2  X3
+        0 -1.0  0.5   0
+        1  0.5 -1.0   0
+
+        Zero-valued columns aren't scaled.
+
+        >>> preprocessor = MaxAbsScaler(columns=["X3"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2   X3
+        0  -6   2  0.0
+        1   3  -4  0.0
+
+        >>> preprocessor = MaxAbsScaler(columns=["X1", "X2"], output_columns=["X1_scaled", "X2_scaled"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2  X3  X1_scaled  X2_scaled
+        0  -2  -3   1       -1.0       -1.0
+        1   0  -3   1        0.0       -1.0
+        2   2   3   1        1.0        1.0
+
+    Args:
+        columns: The columns to separately scale.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """
+
+    def __init__(self, columns: List[str], output_columns: Optional[List[str]] = None):
+        super().__init__()
+        self.columns = columns
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        aggregates = [AbsMax(col) for col in self.columns]
+        self.stats_ = dataset.aggregate(*aggregates)
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_abs_max_scaler(s: pd.Series):
+            s_abs_max = self.stats_[f"abs_max({s.name})"]
+
+            # Handle division by zero.
+            # All values are 0.
+            if s_abs_max == 0:
+                s_abs_max = 1
+
+            return s / s_abs_max
+
+        df[self.output_columns] = df[self.columns].transform(column_abs_max_scaler)
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(columns={self.columns!r}, output_columns={self.output_columns!r})"
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.robust_scaler")
+class RobustScaler(SerializablePreprocessorBase):
+    r"""Scale and translate each column using approximate quantiles.
+
+    The general formula is given by
+
+    .. math::
+        x' = \frac{x - \mu_{1/2}}{\mu_h - \mu_l}
+
+    where :math:`x` is the column, :math:`x'` is the transformed column,
+    :math:`\mu_{1/2}` is the column median. :math:`\mu_{h}` and :math:`\mu_{l}` are the
+    high and low quantiles, respectively. By default, :math:`\mu_{h}` is the third
+    quartile and :math:`\mu_{l}` is the first quartile.
+
+    Internally, the `ApproximateQuantile` aggregator is used to calculate the
+    approximate quantiles.
+
+    .. tip::
+        This scaler works well when your data contains many outliers.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import RobustScaler
+        >>>
+        >>> df = pd.DataFrame({
+        ...     "X1": [1, 2, 3, 4, 5],
+        ...     "X2": [13, 5, 14, 2, 8],
+        ...     "X3": [1, 2, 2, 2, 3],
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>> ds.to_pandas()  # doctest: +SKIP
+           X1  X2  X3
+        0   1  13   1
+        1   2   5   2
+        2   3  14   2
+        3   4   2   2
+        4   5   8   3
+
+        :class:`RobustScaler` separately scales each column.
+
+        >>> preprocessor = RobustScaler(columns=["X1", "X2"])
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+            X1     X2  X3
+        0 -1.0  0.625   1
+        1 -0.5 -0.375   2
+        2  0.0  0.750   2
+        3  0.5 -0.750   2
+        4  1.0  0.000   3
+
+        >>> preprocessor = RobustScaler(
+        ...    columns=["X1", "X2"],
+        ...    output_columns=["X1_scaled", "X2_scaled"]
+        ... )
+        >>> preprocessor.fit_transform(ds).to_pandas()  # doctest: +SKIP
+           X1  X2  X3  X1_scaled  X2_scaled
+        0   1  13   1       -1.0      0.625
+        1   2   5   2       -0.5     -0.375
+        2   3  14   2        0.0      0.750
+        3   4   2   2        0.5     -0.750
+        4   5   8   3        1.0      0.000
+
+    Args:
+        columns: The columns to separately scale.
+        quantile_range: A tuple that defines the lower and upper quantiles. Values
+            must be between 0 and 1. Defaults to the 1st and 3rd quartiles:
+            ``(0.25, 0.75)``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+        quantile_precision: Controls the accuracy and memory footprint of the sketch (K in KLL);
+            higher values yield lower error but use more memory. Defaults to 800. See
+            https://datasketches.apache.org/docs/KLL/KLLAccuracyAndSize.html
+            for details on accuracy and size.
+    """
+
+    DEFAULT_QUANTILE_PRECISION = 800
+
+    def __init__(
+        self,
+        columns: List[str],
+        quantile_range: Tuple[float, float] = (0.25, 0.75),
+        output_columns: Optional[List[str]] = None,
+        quantile_precision: int = DEFAULT_QUANTILE_PRECISION,
+    ):
+        super().__init__()
+        self.columns = columns
+        self.quantile_range = quantile_range
+        self.quantile_precision = quantile_precision
+
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        quantiles = [
+            self.quantile_range[0],
+            0.50,
+            self.quantile_range[1],
+        ]
+        aggregates = [
+            ApproximateQuantile(
+                on=col,
+                quantiles=quantiles,
+                quantile_precision=self.quantile_precision,
+            )
+            for col in self.columns
+        ]
+        aggregated = dataset.aggregate(*aggregates)
+
+        self.stats_ = {}
+        for col in self.columns:
+            low_q, med_q, high_q = aggregated[f"approx_quantile({col})"]
+            self.stats_[f"low_quantile({col})"] = low_q
+            self.stats_[f"median({col})"] = med_q
+            self.stats_[f"high_quantile({col})"] = high_q
+
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_robust_scaler(s: pd.Series):
+            s_low_q = self.stats_[f"low_quantile({s.name})"]
+            s_median = self.stats_[f"median({s.name})"]
+            s_high_q = self.stats_[f"high_quantile({s.name})"]
+            diff = s_high_q - s_low_q
+
+            # Handle division by zero.
+            # Return all zeros.
+            if diff == 0:
+                return np.zeros_like(s)
+
+            return (s - s_median) / diff
+
+        df[self.output_columns] = df[self.columns].transform(column_robust_scaler)
+        return df
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self.columns,
+            "output_columns": self.output_columns,
+            "quantile_range": self.quantile_range,
+            "quantile_precision": self.quantile_precision,
+            "_fitted": getattr(self, "_fitted", None),
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self.columns = fields["columns"]
+        self.output_columns = fields["output_columns"]
+        self.quantile_range = fields["quantile_range"]
+        self.quantile_precision = fields["quantile_precision"]
+        # optional fields
+        self._fitted = fields.get("_fitted")
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"quantile_range={self.quantile_range!r}), "
+            f"output_columns={self.output_columns!r})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/serialization_handlers.py b/lib/python3.12/site-packages/ray/data/preprocessors/serialization_handlers.py
new file mode 100644
index 0000000000000000000000000000000000000000..386dba5c6eb2482ba180aca8707a1e5a582a6c5f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/serialization_handlers.py
@@ -0,0 +1,195 @@
+"""
+Serialization handlers for preprocessor save/load functionality.
+
+This module implements a factory pattern to abstract different serialization formats,
+making it easier to add new formats and maintain existing ones.
+"""
+
+import abc
+import base64
+import pickle
+from enum import Enum
+from typing import Any, Dict, Optional, Union
+
+from ray.cloudpickle import cloudpickle
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class HandlerFormatName(Enum):
+    """Enum for consistent format naming in the factory."""
+
+    CLOUDPICKLE = "cloudpickle"
+    PICKLE = "pickle"
+
+
+@DeveloperAPI
+class SerializationHandler(abc.ABC):
+    """Abstract base class for handling preprocessor serialization formats."""
+
+    @abc.abstractmethod
+    def serialize(
+        self, data: Union["Preprocessor", Dict[str, Any]]  # noqa: F821
+    ) -> Union[str, bytes]:
+        """Serialize preprocessor data to the specific format.
+
+        Args:
+            data: Dictionary containing preprocessor metadata and stats
+
+        Returns:
+            Serialized data in format-specific representation
+        """
+        pass
+
+    @abc.abstractmethod
+    def deserialize(self, serialized: Union[str, bytes]) -> Any:
+        """Deserialize data from the specific format.
+
+        Args:
+            serialized: Serialized data in format-specific representation
+
+        Returns:
+            For structured formats (CloudPickle/JSON/MessagePack): Dictionary containing preprocessor metadata and stats
+            For pickle format: The actual deserialized object
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_magic_bytes(self) -> Union[str, bytes]:
+        """Get the magic bytes/prefix for this format."""
+        pass
+
+    def strip_magic_bytes(self, serialized: Union[str, bytes]) -> Union[str, bytes]:
+        """Remove magic bytes from serialized data."""
+        magic = self.get_magic_bytes()
+        if isinstance(serialized, (str, bytes)) and serialized.startswith(magic):
+            return serialized[len(magic) :]
+        return serialized
+
+
+@DeveloperAPI
+class CloudPickleSerializationHandler(SerializationHandler):
+    """Handler for CloudPickle serialization format."""
+
+    MAGIC_CLOUDPICKLE = b"CPKL:"
+
+    def serialize(
+        self, data: Union["Preprocessor", Dict[str, Any]]  # noqa: F821
+    ) -> bytes:
+        """Serialize to CloudPickle format with magic prefix."""
+        return self.MAGIC_CLOUDPICKLE + cloudpickle.dumps(data)
+
+    def deserialize(self, serialized: bytes) -> Dict[str, Any]:
+        """Deserialize from CloudPickle format."""
+        if not isinstance(serialized, bytes):
+            raise ValueError(
+                f"Expected bytes for CloudPickle deserialization, got {type(serialized)}"
+            )
+
+        if not serialized.startswith(self.MAGIC_CLOUDPICKLE):
+            raise ValueError(f"Invalid CloudPickle magic bytes: {serialized[:10]}")
+
+        cloudpickle_data = self.strip_magic_bytes(serialized)
+        return cloudpickle.loads(cloudpickle_data)
+
+    def get_magic_bytes(self) -> bytes:
+        return self.MAGIC_CLOUDPICKLE
+
+
+@DeveloperAPI
+class PickleSerializationHandler(SerializationHandler):
+    """Handler for legacy Pickle serialization format."""
+
+    def serialize(
+        self, data: Union["Preprocessor", Dict[str, Any]]  # noqa: F821
+    ) -> str:
+        """
+        Serialize using pickle format (for backward compatibility).
+        data is ignored, but kept for consistency
+
+        """
+        return base64.b64encode(pickle.dumps(data)).decode("ascii")
+
+    def deserialize(
+        self, serialized: str
+    ) -> Any:  # Returns the actual object, not metadata
+        """Deserialize from pickle format (legacy support)."""
+        # For pickle, we return the actual deserialized object directly
+        return pickle.loads(base64.b64decode(serialized))
+
+    def get_magic_bytes(self) -> str:
+        return ""  # Pickle format doesn't use magic bytes
+
+
+class SerializationHandlerFactory:
+    """Factory class for creating appropriate serialization handlers."""
+
+    _handlers = {
+        HandlerFormatName.CLOUDPICKLE: CloudPickleSerializationHandler,
+        HandlerFormatName.PICKLE: PickleSerializationHandler,
+    }
+
+    @classmethod
+    def register_handler(cls, format_name: HandlerFormatName, handler_class: type):
+        """Register a new serialization handler."""
+        cls._handlers[format_name] = handler_class
+
+    @classmethod
+    def get_handler(
+        cls,
+        format_identifier: Optional[HandlerFormatName] = None,
+        data: Optional[Union[str, bytes]] = None,
+        **kwargs,
+    ) -> SerializationHandler:
+        """Get the appropriate serialization handler for a format or serialized data.
+
+        Args:
+            format_identifier: The format to use for serialization. If None, will detect from data.
+            data: Serialized data to detect format from (used when format_identifier is None).
+            **kwargs: Additional keyword arguments (currently unused).
+
+        Returns:
+            SerializationHandler instance for the format
+
+        Raises:
+            ValueError: If format is not supported or cannot be detected
+        """
+        # If it's already a format enum, use it directly
+        if not format_identifier:
+            format_identifier = cls.detect_format(data)
+
+        if format_identifier not in cls._handlers:
+            raise ValueError(
+                f"Unsupported serialization format: {format_identifier.value}. "
+                f"Supported formats: {list(cls._handlers.keys())}"
+            )
+
+        handler_class = cls._handlers[format_identifier]
+        return handler_class()
+
+    @classmethod
+    def detect_format(cls, serialized: Union[str, bytes]) -> HandlerFormatName:
+        """Detect the serialization format from the magic bytes.
+
+        Args:
+            serialized: Serialized data
+
+        Returns:
+            Format name enum
+
+        Raises:
+            ValueError: If format cannot be detected
+        """
+        # Check for CloudPickle first (binary format)
+        if isinstance(serialized, bytes) and serialized.startswith(
+            CloudPickleSerializationHandler.MAGIC_CLOUDPICKLE
+        ):
+            return HandlerFormatName.CLOUDPICKLE
+
+        # Check for legacy pickle format (no magic bytes, should be base64 encoded)
+        if isinstance(serialized, str):
+            return HandlerFormatName.PICKLE
+
+        raise ValueError(
+            f"Cannot detect serialization format from: {serialized[:20]}..."
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/tokenizer.py b/lib/python3.12/site-packages/ray/data/preprocessors/tokenizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e44d3934579d01d498e57fb36ec25f1ab9fb17b
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/tokenizer.py
@@ -0,0 +1,91 @@
+from typing import Callable, List, Optional
+
+import pandas as pd
+
+from ray.data.preprocessor import Preprocessor
+from ray.data.preprocessors.utils import simple_split_tokenizer
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class Tokenizer(Preprocessor):
+    """Replace each string with a list of tokens.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> df = pd.DataFrame({"text": ["Hello, world!", "foo bar\\nbaz"]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+
+        The default ``tokenization_fn`` delimits strings using the space character.
+
+        >>> from ray.data.preprocessors import Tokenizer
+        >>> tokenizer = Tokenizer(columns=["text"])
+        >>> tokenizer.transform(ds).to_pandas()  # doctest: +SKIP
+                       text
+        0  [Hello,, world!]
+        1   [foo, bar\\nbaz]
+
+        If the default logic isn't adequate for your use case, you can specify a
+        custom ``tokenization_fn``.
+
+        >>> import string
+        >>> def tokenization_fn(s):
+        ...     for character in string.punctuation:
+        ...         s = s.replace(character, "")
+        ...     return s.split()
+        >>> tokenizer = Tokenizer(columns=["text"], tokenization_fn=tokenization_fn)
+        >>> tokenizer.transform(ds).to_pandas()  # doctest: +SKIP
+                      text
+        0   [Hello, world]
+        1  [foo, bar, baz]
+
+        :class:`Tokenizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the tokenized values.
+
+        >>> tokenizer = Tokenizer(columns=["text"], output_columns=["text_tokenized"])
+        >>> tokenizer.transform(ds).to_pandas()  # doctest: +SKIP
+                    text    text_tokenized
+        0  Hello, world!  [Hello,, world!]
+        1   foo bar\\nbaz   [foo, bar\\nbaz]
+
+    Args:
+        columns: The columns to tokenize.
+        tokenization_fn: The function used to generate tokens. This function
+            should accept a string as input and return a list of tokens as
+            output. If unspecified, the tokenizer uses a function equivalent to
+            ``lambda s: s.split(" ")``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        tokenization_fn: Optional[Callable[[str], List[str]]] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        self.columns = columns
+        # TODO(matt): Add a more robust default tokenizer.
+        self.tokenization_fn = tokenization_fn or simple_split_tokenizer
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_tokenizer(s: pd.Series):
+            return s.map(self.tokenization_fn)
+
+        df[self.output_columns] = df.loc[:, self.columns].transform(column_tokenizer)
+        return df
+
+    def __repr__(self):
+        name = getattr(self.tokenization_fn, "__name__", self.tokenization_fn)
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"tokenization_fn={name}, output_columns={self.output_columns!r})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/torch.py b/lib/python3.12/site-packages/ray/data/preprocessors/torch.py
new file mode 100644
index 0000000000000000000000000000000000000000..4206bc18e214f36d566931ad3cc6c365983c352f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/torch.py
@@ -0,0 +1,149 @@
+from typing import TYPE_CHECKING, Callable, Dict, List, Mapping, Optional, Union
+
+import numpy as np
+
+from ray.air.util.data_batch_conversion import BatchFormat
+from ray.air.util.tensor_extensions.utils import _create_possibly_ragged_ndarray
+from ray.data.preprocessor import Preprocessor
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    import torch
+
+
+@PublicAPI(stability="alpha")
+class TorchVisionPreprocessor(Preprocessor):
+    """Apply a `TorchVision transform <https://pytorch.org/vision/stable/transforms.html>`_
+    to image columns.
+
+    Examples:
+
+        Torch models expect inputs of shape :math:`(B, C, H, W)` in the range
+        :math:`[0.0, 1.0]`. To convert images to this format, add ``ToTensor`` to your
+        preprocessing pipeline.
+
+        .. testcode::
+
+            from torchvision import transforms
+
+            import ray
+            from ray.data.preprocessors import TorchVisionPreprocessor
+
+            transform = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Resize((224, 224)),
+            ])
+            preprocessor = TorchVisionPreprocessor(["image"], transform=transform)
+
+            dataset = ray.data.read_images("s3://anonymous@air-example-data-2/imagenet-sample-images")
+            dataset = preprocessor.transform(dataset)
+
+
+        For better performance, set ``batched`` to ``True`` and replace ``ToTensor``
+        with a batch-supporting ``Lambda``.
+
+        .. testcode::
+
+            import numpy as np
+            import torch
+
+            def to_tensor(batch: np.ndarray) -> torch.Tensor:
+                tensor = torch.as_tensor(batch, dtype=torch.float)
+                # (B, H, W, C) -> (B, C, H, W)
+                tensor = tensor.permute(0, 3, 1, 2).contiguous()
+                # [0., 255.] -> [0., 1.]
+                tensor = tensor.div(255)
+                return tensor
+
+            transform = transforms.Compose([
+                transforms.Lambda(to_tensor),
+                transforms.Resize((224, 224))
+            ])
+            preprocessor = TorchVisionPreprocessor(["image"], transform=transform, batched=True)
+
+            dataset = ray.data.read_images("s3://anonymous@air-example-data-2/imagenet-sample-images")
+            dataset = preprocessor.transform(dataset)
+
+    Args:
+        columns: The columns to apply the TorchVision transform to.
+        transform: The TorchVision transform you want to apply. This transform should
+            accept a ``np.ndarray`` or ``torch.Tensor`` as input and return a
+            ``torch.Tensor`` as output.
+        output_columns: The output name for each input column. If not specified, this
+            defaults to the same set of columns as the columns.
+        batched: If ``True``, apply ``transform`` to batches of shape
+            :math:`(B, H, W, C)`. Otherwise, apply ``transform`` to individual images.
+    """  # noqa: E501
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        transform: Callable[[Union["np.ndarray", "torch.Tensor"]], "torch.Tensor"],
+        output_columns: Optional[List[str]] = None,
+        batched: bool = False,
+    ):
+        if not output_columns:
+            output_columns = columns
+        if len(columns) != len(output_columns):
+            raise ValueError(
+                "The length of columns should match the "
+                f"length of output_columns: {columns} vs {output_columns}."
+            )
+        self._columns = columns
+        self._output_columns = output_columns
+        self._torchvision_transform = transform
+        self._batched = batched
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"columns={self._columns}, "
+            f"output_columns={self._output_columns}, "
+            f"transform={self._torchvision_transform!r})"
+        )
+
+    def _transform_numpy(
+        self, data_batch: Dict[str, "np.ndarray"]
+    ) -> Dict[str, "np.ndarray"]:
+        import torch
+
+        from ray.air._internal.torch_utils import convert_ndarray_to_torch_tensor
+
+        def apply_torchvision_transform(array: np.ndarray) -> np.ndarray:
+            try:
+                tensor = convert_ndarray_to_torch_tensor(array)
+                output = self._torchvision_transform(tensor)
+            except TypeError:
+                # Transforms like `ToTensor` expect a `np.ndarray` as input.
+                output = self._torchvision_transform(array)
+            if isinstance(output, torch.Tensor):
+                output = output.numpy()
+            if not isinstance(output, np.ndarray):
+                raise ValueError(
+                    "`TorchVisionPreprocessor` expected your transform to return a "
+                    "`torch.Tensor` or `np.ndarray`, but your transform returned a "
+                    f"`{type(output).__name__}` instead."
+                )
+            return output
+
+        def transform_batch(batch: np.ndarray) -> np.ndarray:
+            if self._batched:
+                return apply_torchvision_transform(batch)
+            return _create_possibly_ragged_ndarray(
+                [apply_torchvision_transform(array) for array in batch]
+            )
+
+        if isinstance(data_batch, Mapping):
+            for input_col, output_col in zip(self._columns, self._output_columns):
+                data_batch[output_col] = transform_batch(data_batch[input_col])
+        else:
+            # TODO(ekl) deprecate this code path. Unfortunately, predictors are still
+            # sending schemaless arrays to preprocessors.
+            data_batch = transform_batch(data_batch)
+
+        return data_batch
+
+    def preferred_batch_format(cls) -> BatchFormat:
+        return BatchFormat.NUMPY
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/transformer.py b/lib/python3.12/site-packages/ray/data/preprocessors/transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..6618eb31e7ba45b4fbe2dcf55748b74b03126cbc
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/transformer.py
@@ -0,0 +1,101 @@
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+
+from ray.data.preprocessor import Preprocessor
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+class PowerTransformer(Preprocessor):
+    """Apply a `power transform <https://en.wikipedia.org/wiki/Power_transform>`_ to
+    make your data more normally distributed.
+
+    Some models expect data to be normally distributed. By making your data more
+    Gaussian-like, you might be able to improve your model's performance.
+
+    This preprocessor supports the following transformations:
+
+    * `Yeo-Johnson <https://en.wikipedia.org/wiki/Power_transform#Yeo%E2%80%93Johnson_transformation>`_
+    * `Box-Cox <https://en.wikipedia.org/wiki/Power_transform#Box%E2%80%93Cox_transformation>`_
+
+    Box-Cox requires all data to be positive.
+
+    .. warning::
+
+        You need to manually specify the transform's power parameter. If you
+        choose a bad value, the transformation might not work well.
+
+    Args:
+        columns: The columns to separately transform.
+        power: A parameter that determines how your data is transformed. Practioners
+            typically set ``power`` between :math:`-2.5` and :math:`2.5`, although you
+            may need to try different values to find one that works well.
+        method: A string representing which transformation to apply. Supports
+            ``"yeo-johnson"`` and ``"box-cox"``. If you choose ``"box-cox"``, your data
+            needs to be positive. Defaults to ``"yeo-johnson"``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """  # noqa: E501
+
+    _valid_methods = ["yeo-johnson", "box-cox"]
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        power: float,
+        method: str = "yeo-johnson",
+        *,
+        output_columns: Optional[List[str]] = None,
+    ):
+        self.columns = columns
+        self.method = method
+        self.power = power
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+        if method not in self._valid_methods:
+            raise ValueError(
+                f"Method {method} is not supported."
+                f"Supported values are: {self._valid_methods}"
+            )
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_power_transformer(s: pd.Series):
+            if self.method == "yeo-johnson":
+                result = np.zeros_like(s, dtype=np.float64)
+                pos = s >= 0  # binary mask
+
+                if self.power != 0:
+                    result[pos] = (np.power(s[pos] + 1, self.power) - 1) / self.power
+                else:
+                    result[pos] = np.log(s[pos] + 1)
+
+                if self.power != 2:
+                    result[~pos] = -(np.power(-s[~pos] + 1, 2 - self.power) - 1) / (
+                        2 - self.power
+                    )
+                else:
+                    result[~pos] = -np.log(-s[~pos] + 1)
+                return result
+
+            else:  # box-cox
+                if self.power != 0:
+                    return (np.power(s, self.power) - 1) / self.power
+                else:
+                    return np.log(s)
+
+        df[self.output_columns] = df[self.columns].transform(column_power_transformer)
+        return df
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"power={self.power!r}, method={self.method!r}, "
+            f"output_columns={self.output_columns!r})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/utils.py b/lib/python3.12/site-packages/ray/data/preprocessors/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d22120623615a52d2a193010d4d09448a3418cd
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/utils.py
@@ -0,0 +1,227 @@
+import hashlib
+from collections import deque
+from typing import Any, Callable, Deque, Dict, List, Optional, Union
+
+from ray.data import Dataset
+from ray.data.aggregate import AggregateFnV2
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def simple_split_tokenizer(value: str) -> List[str]:
+    """Tokenize a string using a split on spaces."""
+    return value.split(" ")
+
+
+@DeveloperAPI
+def simple_hash(value: object, num_features: int) -> int:
+    """Deterministically hash a value into the integer space."""
+    encoded_value = str(value).encode()
+    hashed_value = hashlib.sha1(encoded_value)
+    hashed_value_int = int(hashed_value.hexdigest(), 16)
+    return hashed_value_int % num_features
+
+
+class BaseStatSpec:
+    """Encapsulates a statistical computation with optional post-processing."""
+
+    def __init__(
+        self,
+        *,
+        stat_fn: Union[AggregateFnV2, Callable],
+        post_process_fn: Callable = lambda x: x,
+        post_key_fn: Callable[[str], str],
+    ):
+        self.stat_fn = stat_fn
+        self.post_process_fn = post_process_fn
+        self.post_key_fn = post_key_fn
+
+
+class AggregateStatSpec(BaseStatSpec):
+    """Represents an AggregateFnV2 spec for a single column."""
+
+    def __init__(
+        self,
+        *,
+        aggregator_fn: Union[AggregateFnV2, Callable[[str], AggregateFnV2]],
+        post_process_fn: Callable = lambda x: x,
+        post_key_fn: Callable[[str], str],
+        column: Optional[str] = None,
+    ):
+        super().__init__(
+            stat_fn=aggregator_fn,
+            post_process_fn=post_process_fn,
+            post_key_fn=post_key_fn,
+        )
+        self.column = column
+
+
+class CallableStatSpec(BaseStatSpec):
+    """Represents a user-defined stat function that operates outside Dataset.aggregate."""
+
+    def __init__(
+        self,
+        *,
+        stat_fn: Callable,
+        stat_key_fn: Optional[Callable[[str], str]],
+        post_key_fn: Optional[Callable[[str], str]],
+        post_process_fn: Callable = lambda x: x,
+        columns: List[str],
+    ):
+        super().__init__(
+            stat_fn=stat_fn, post_process_fn=post_process_fn, post_key_fn=post_key_fn
+        )
+        self.columns = columns
+        self.stat_key_fn = stat_key_fn
+
+
+class StatComputationPlan:
+    """
+    Encapsulates a set of aggregators (AggregateFnV2) and legacy stat functions
+    to compute statistics over a Ray dataset.
+
+    Supports two types of aggregations:
+    1. AggregateFnV2-based aggregators, which are batch-executed using `Dataset.aggregate(...)`.
+    2. Callable-based stat functions, executed sequentially (legacy use case).
+    """
+
+    def __init__(self):
+        self._aggregators: Deque[BaseStatSpec] = deque()
+
+    def reset(self):
+        self._aggregators.clear()
+
+    def add_aggregator(
+        self,
+        *,
+        aggregator_fn: Callable[[str], AggregateFnV2],
+        post_process_fn: Callable = lambda x: x,
+        post_key_fn: Optional[Callable[[str], str]] = None,
+        columns: List[str],
+    ) -> None:
+        """
+        Registers an AggregateFnV2 factory for one or more columns.
+
+        Args:
+            aggregator_fn: A callable (typically a lambda or class) that accepts a column name and returns an instance of AggregateFnV2.
+            post_process_fn: Function to post-process the aggregated result.
+            post_key_fn: Optional key generator to use to save aggregation results after post-processing.
+            columns: List of column names to aggregate.
+        """
+        for column in columns:
+            agg_instance = aggregator_fn(column)
+            self._aggregators.append(
+                AggregateStatSpec(
+                    aggregator_fn=agg_instance,
+                    post_process_fn=post_process_fn,
+                    post_key_fn=post_key_fn,
+                    column=column,
+                )
+            )
+
+    def add_callable_stat(
+        self,
+        *,
+        stat_fn: Callable[[], Any],
+        post_process_fn: Callable = lambda x: x,
+        stat_key_fn: Callable[[str], str],
+        post_key_fn: Optional[Callable[[str], str]] = None,
+        columns: List[str],
+    ) -> None:
+        """
+        Registers a custom stat function to be run sequentially.
+
+        This supports legacy use cases where arbitrary callables are needed
+        and cannot be run via Dataset.aggregate().
+
+        Args:
+            stat_fn: A zero-argument callable that returns the stat.
+            post_process_fn: Function to apply to the result.
+            stat_key_fn:
+            post_key_fn:
+            columns:
+        """
+        self._aggregators.append(
+            CallableStatSpec(
+                stat_fn=stat_fn,
+                post_process_fn=post_process_fn,
+                columns=columns,
+                stat_key_fn=stat_key_fn,
+                post_key_fn=post_key_fn or stat_key_fn,
+            )
+        )
+
+    def compute(self, dataset: Dataset) -> Dict[str, Any]:
+        """
+        Executes all registered aggregators and stat functions.
+
+        AggregateFnV2-based aggregators are batched and executed via Dataset.aggregate().
+        Callable-based stat functions are run sequentially.
+
+        Args:
+            dataset: The Ray Dataset to compute statistics on.
+
+        Returns:
+            A dictionary of computed statistics.
+        """
+        stats = {}
+        # Run batched aggregators (AggregateFnV2)
+        aggregators = self._get_aggregate_fn_list()
+        if aggregators:
+            raw_result = dataset.aggregate(*aggregators)
+            for spec in self._get_aggregate_specs():
+                stat_key = spec.stat_fn.name
+                post_key = (
+                    spec.post_key_fn(spec.column)
+                    if spec.post_key_fn is not None
+                    else stat_key
+                )
+                stats[post_key] = spec.post_process_fn(raw_result[stat_key])
+
+        # Run sequential stat functions
+        for spec in self._get_custom_stat_fn_specs():
+            result = spec.stat_fn(spec.stat_key_fn)
+            for col in spec.columns:
+                stat_key = spec.stat_key_fn(col)
+                post_key = spec.post_key_fn(col)
+                stats[post_key] = spec.post_process_fn(result[stat_key])
+
+        return stats
+
+    def _get_aggregate_fn_list(self) -> List[AggregateFnV2]:
+        return [
+            spec.stat_fn
+            for spec in self._aggregators
+            if isinstance(spec, AggregateStatSpec)
+        ]
+
+    def _get_aggregate_specs(self) -> List[AggregateStatSpec]:
+        return [
+            spec for spec in self._aggregators if isinstance(spec, AggregateStatSpec)
+        ]
+
+    def _get_custom_stat_fn_specs(self) -> List[CallableStatSpec]:
+        return [
+            spec for spec in self._aggregators if isinstance(spec, CallableStatSpec)
+        ]
+
+    def __iter__(self):
+        """
+        Iterates over all AggregatorSpecs.
+        """
+        return iter(self._get_aggregate_specs())
+
+
+def make_post_processor(base_fn, callbacks: List[Callable]):
+    """
+    Wraps a base post-processing function with a sequence of callback functions.
+    Useful when multiple post-processing steps need to be applied in order.
+    """
+
+    def wrapper(result):
+        processed = base_fn(result)
+        for cb in callbacks:
+            processed = cb(processed)
+        return processed
+
+    return wrapper
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/vectorizer.py b/lib/python3.12/site-packages/ray/data/preprocessors/vectorizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b5bc21ec03cc370f765fa85a31850b754c2aeaa
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/vectorizer.py
@@ -0,0 +1,328 @@
+from collections import Counter
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+import pandas as pd
+
+from ray.data.preprocessor import Preprocessor
+from ray.data.preprocessors.utils import simple_hash, simple_split_tokenizer
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Dataset
+
+
+@PublicAPI(stability="alpha")
+class HashingVectorizer(Preprocessor):
+    """Count the frequency of tokens using the
+    `hashing trick <https://en.wikipedia.org/wiki/Feature_hashing>`_.
+
+    This preprocessors creates a list column for each input column. For each row,
+    the list contains the frequency counts of tokens (for CountVectorizer) or hash values
+    (for HashingVectorizer). For HashingVectorizer, the list will have length
+    ``num_features``. If ``num_features`` is large enough relative to the size of your
+    vocabulary, then each index approximately corresponds to the frequency of a unique
+    token.
+
+    :class:`HashingVectorizer` is memory efficient and quick to pickle. However, given a
+    transformed column, you can't know which tokens correspond to it. This might make it
+    hard to determine which tokens are important to your model.
+
+    .. note::
+
+        This preprocessor transforms each input column to a
+        `document-term matrix <https://en.wikipedia.org/wiki/Document-term_matrix>`_.
+
+        A document-term matrix is a table that describes the frequency of tokens in a
+        collection of documents. For example, the strings `"I like Python"` and `"I
+        dislike Python"` might have the document-term matrix below:
+
+        .. code-block::
+
+                    corpus_I  corpus_Python  corpus_dislike  corpus_like
+                0         1              1               1            0
+                1         1              1               0            1
+
+        To generate the matrix, you typically map each token to a unique index. For
+        example:
+
+        .. code-block::
+
+                        token  index
+                0        I      0
+                1   Python      1
+                2  dislike      2
+                3     like      3
+
+        The problem with this approach is that memory use scales linearly with the size
+        of your vocabulary. :class:`HashingVectorizer` circumvents this problem by
+        computing indices with a hash function:
+        :math:`\\texttt{index} = hash(\\texttt{token})`.
+
+    .. warning::
+        Sparse matrices aren't currently supported. If you use a large ``num_features``,
+        this preprocessor might behave poorly.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import HashingVectorizer
+        >>>
+        >>> df = pd.DataFrame({
+        ...     "corpus": [
+        ...         "Jimmy likes volleyball",
+        ...         "Bob likes volleyball too",
+        ...         "Bob also likes fruit jerky"
+        ...     ]
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>>
+        >>> vectorizer = HashingVectorizer(["corpus"], num_features=8)
+        >>> vectorizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                             corpus
+        0  [1, 0, 1, 0, 0, 0, 0, 1]
+        1  [1, 0, 1, 0, 0, 0, 1, 1]
+        2  [0, 0, 1, 1, 0, 2, 1, 0]
+
+        :class:`HashingVectorizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> vectorizer = HashingVectorizer(["corpus"], num_features=8, output_columns=["corpus_hashed"])
+        >>> vectorizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                               corpus             corpus_hashed
+        0      Jimmy likes volleyball  [1, 0, 1, 0, 0, 0, 0, 1]
+        1    Bob likes volleyball too  [1, 0, 1, 0, 0, 0, 1, 1]
+        2  Bob also likes fruit jerky  [0, 0, 1, 1, 0, 2, 1, 0]
+
+    Args:
+        columns: The columns to separately tokenize and count.
+        num_features: The number of features used to represent the vocabulary. You
+            should choose a value large enough to prevent hash collisions between
+            distinct tokens.
+        tokenization_fn: The function used to generate tokens. This function
+            should accept a string as input and return a list of tokens as
+            output. If unspecified, the tokenizer uses a function equivalent to
+            ``lambda s: s.split(" ")``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+
+    .. seealso::
+
+        :class:`CountVectorizer`
+            Another method for counting token frequencies. Unlike :class:`HashingVectorizer`,
+            :class:`CountVectorizer` creates a feature for each unique token. This
+            enables you to compute the inverse transformation.
+
+        :class:`FeatureHasher`
+            This preprocessor is similar to :class:`HashingVectorizer`, except it expects
+            a table describing token frequencies. In contrast,
+            :class:`FeatureHasher` expects a column containing documents.
+    """  # noqa: E501
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        num_features: int,
+        tokenization_fn: Optional[Callable[[str], List[str]]] = None,
+        *,
+        output_columns: Optional[List[str]] = None,
+    ):
+        self.columns = columns
+        self.num_features = num_features
+        self.tokenization_fn = tokenization_fn or simple_split_tokenizer
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def hash_count(tokens: List[str]) -> Counter:
+            hashed_tokens = [simple_hash(token, self.num_features) for token in tokens]
+            return Counter(hashed_tokens)
+
+        for col, output_col in zip(self.columns, self.output_columns):
+            tokenized = df[col].map(self.tokenization_fn)
+            hashed = tokenized.map(hash_count)
+            # Create a list to store the hash columns
+            hash_columns = []
+            for i in range(self.num_features):
+                series = hashed.map(lambda counts: counts[i])
+                series.name = f"hash_{i}"
+                hash_columns.append(series)
+            # Concatenate all hash columns into a single list column
+            df[output_col] = pd.concat(hash_columns, axis=1).values.tolist()
+
+        return df
+
+    def __repr__(self):
+        fn_name = getattr(self.tokenization_fn, "__name__", self.tokenization_fn)
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"num_features={self.num_features!r}, tokenization_fn={fn_name}, "
+            f"output_columns={self.output_columns!r})"
+        )
+
+
+@PublicAPI(stability="alpha")
+class CountVectorizer(Preprocessor):
+    """Count the frequency of tokens in a column of strings.
+
+    :class:`CountVectorizer` operates on columns that contain strings. For example:
+
+    .. code-block::
+
+                        corpus
+        0    I dislike Python
+        1       I like Python
+
+    This preprocessor creates a list column for each input column. Each list contains
+    the frequency counts of tokens in order of their first appearance. For example:
+
+    .. code-block::
+
+                    corpus
+        0    [1, 1, 1, 0]  # Counts for [I, dislike, Python, like]
+        1    [1, 0, 1, 1]  # Counts for [I, dislike, Python, like]
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> from ray.data.preprocessors import CountVectorizer
+        >>>
+        >>> df = pd.DataFrame({
+        ...     "corpus": [
+        ...         "Jimmy likes volleyball",
+        ...         "Bob likes volleyball too",
+        ...         "Bob also likes fruit jerky"
+        ...     ]
+        ... })
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+        >>>
+        >>> vectorizer = CountVectorizer(["corpus"])
+        >>> vectorizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                             corpus
+        0  [1, 0, 1, 1, 0, 0, 0, 0]
+        1  [1, 1, 1, 0, 0, 0, 0, 1]
+        2  [1, 1, 0, 0, 1, 1, 1, 0]
+
+        You can limit the number of tokens in the vocabulary with ``max_features``.
+
+        >>> vectorizer = CountVectorizer(["corpus"], max_features=3)
+        >>> vectorizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+              corpus
+        0  [1, 0, 1]
+        1  [1, 1, 1]
+        2  [1, 1, 0]
+
+        :class:`CountVectorizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the encoded values.
+
+        >>> vectorizer = CountVectorizer(["corpus"], output_columns=["corpus_counts"])
+        >>> vectorizer.fit_transform(ds).to_pandas()  # doctest: +SKIP
+                               corpus             corpus_counts
+        0      Jimmy likes volleyball  [1, 0, 1, 1, 0, 0, 0, 0]
+        1    Bob likes volleyball too  [1, 1, 1, 0, 0, 0, 0, 1]
+        2  Bob also likes fruit jerky  [1, 1, 0, 0, 1, 1, 1, 0]
+
+    Args:
+        columns: The columns to separately tokenize and count.
+        tokenization_fn: The function used to generate tokens. This function
+            should accept a string as input and return a list of tokens as
+            output. If unspecified, the tokenizer uses a function equivalent to
+            ``lambda s: s.split(" ")``.
+        max_features: The maximum number of tokens to encode in the transformed
+            dataset. If specified, only the most frequent tokens are encoded.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        columns: List[str],
+        tokenization_fn: Optional[Callable[[str], List[str]]] = None,
+        max_features: Optional[int] = None,
+        *,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        self.columns = columns
+        self.tokenization_fn = tokenization_fn or simple_split_tokenizer
+        self.max_features = max_features
+        self.output_columns = Preprocessor._derive_and_validate_output_columns(
+            columns, output_columns
+        )
+
+    def _fit(self, dataset: "Dataset") -> Preprocessor:
+        def stat_fn(key_gen):
+            def get_pd_value_counts(df: pd.DataFrame) -> List[Counter]:
+                def get_token_counts(col):
+                    token_series = df[col].apply(self.tokenization_fn)
+                    tokens = token_series.sum()
+                    return Counter(tokens)
+
+                return {col: [get_token_counts(col)] for col in self.columns}
+
+            value_counts = dataset.map_batches(
+                get_pd_value_counts, batch_format="pandas"
+            )
+            total_counts = {col: Counter() for col in self.columns}
+            for batch in value_counts.iter_batches(batch_size=None):
+                for col, counters in batch.items():
+                    for counter in counters:
+                        total_counts[col].update(counter)
+
+            def most_common(counter: Counter, n: int):
+                return Counter(dict(counter.most_common(n)))
+
+            top_counts = [
+                most_common(counter, self.max_features)
+                for counter in total_counts.values()
+            ]
+
+            return {
+                key_gen(col): counts  # noqa
+                for (col, counts) in zip(self.columns, top_counts)
+            }
+
+        self.stat_computation_plan.add_callable_stat(
+            stat_fn=lambda key_gen: stat_fn(key_gen),
+            stat_key_fn=lambda col: f"token_counts({col})",
+            columns=self.columns,
+        )
+
+        return self
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        result_columns = []
+        for col, output_col in zip(self.columns, self.output_columns):
+            token_counts = self.stats_[f"token_counts({col})"]
+            sorted_tokens = [token for (token, count) in token_counts.most_common()]
+            tokenized = df[col].map(self.tokenization_fn).map(Counter)
+
+            # Create a list to store token frequencies
+            token_columns = []
+            for token in sorted_tokens:
+                series = tokenized.map(lambda val: val[token])
+                series.name = token
+                token_columns.append(series)
+
+            # Concatenate all token columns into a single list column
+            if token_columns:
+                df[output_col] = pd.concat(token_columns, axis=1).values.tolist()
+            else:
+                df[output_col] = [[]] * len(df)
+            result_columns.append(output_col)
+
+        return df
+
+    def __repr__(self):
+        fn_name = getattr(self.tokenization_fn, "__name__", self.tokenization_fn)
+        return (
+            f"{self.__class__.__name__}(columns={self.columns!r}, "
+            f"tokenization_fn={fn_name}, max_features={self.max_features!r}, "
+            f"output_columns={self.output_columns!r})"
+        )
diff --git a/lib/python3.12/site-packages/ray/data/preprocessors/version_support.py b/lib/python3.12/site-packages/ray/data/preprocessors/version_support.py
new file mode 100644
index 0000000000000000000000000000000000000000..df6d0a7acd588b833e1084d55e2cd38ff87b2ea1
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/data/preprocessors/version_support.py
@@ -0,0 +1,86 @@
+class UnknownPreprocessorError(ValueError):
+    """Raised when attempting to deserialize an unknown/unregistered preprocessor type."""
+
+    def __init__(self, preprocessor_type: str):
+        self.preprocessor_type = preprocessor_type
+        super().__init__(f"Unknown preprocessor type: {preprocessor_type}")
+
+
+_PREPROCESSOR_REGISTRY = {}
+
+
+def SerializablePreprocessor(version: int, identifier: str):
+    """Register a preprocessor class for serialization.
+
+    This decorator registers a preprocessor class in the serialization registry,
+    enabling it to be serialized and deserialized. The decorated class MUST inherit
+    from SerializablePreprocessor.
+
+    Args:
+        version: Version number for this preprocessor's serialization format
+        identifier: Stable identifier for serialization. This identifier will be used
+            in serialized data. Using an explicit identifier allows classes to be
+            renamed without breaking compatibility with existing serialized data.
+
+    Returns:
+        A decorator function that registers the class and returns it unchanged.
+
+    Raises:
+        TypeError: If the decorated class does not inherit from SerializablePreprocessor
+
+    Note:
+        If a class with the same identifier is already registered, logs an info message
+        and overwrites the previous registration.
+
+    Examples:
+        @SerializablePreprocessor(version=1, identifier="my_preprocessor_v1")
+        class MyPreprocessor(SerializablePreprocessor):
+            pass
+    """
+
+    def decorator(cls):
+        import logging
+
+        from ray.data.preprocessor import SerializablePreprocessorBase
+
+        # Verify that the class inherits from SerializablePreprocessor
+        if not issubclass(cls, SerializablePreprocessorBase):
+            raise TypeError(
+                f"Class {cls.__module__}.{cls.__qualname__} must inherit from "
+                f"SerializablePreprocessor to use @SerializablePreprocessor decorator."
+            )
+
+        cls.set_version(version)
+        cls.set_preprocessor_class_id(identifier)
+
+        # Check for collisions and log info message
+        if identifier in _PREPROCESSOR_REGISTRY:
+            existing = _PREPROCESSOR_REGISTRY[identifier]
+            if existing != cls:
+                logging.info(
+                    f"Preprocessor id collision: '{identifier}' was already registered "
+                    f"by {existing.__module__}.{existing.__qualname__}. "
+                    f"Overwriting with {cls.__module__}.{cls.__qualname__}."
+                )
+
+        _PREPROCESSOR_REGISTRY[identifier] = cls
+        return cls
+
+    return decorator
+
+
+def _lookup_class(serialization_id: str):
+    """Look up a preprocessor class by its serialization ID.
+
+    Args:
+        serialization_id: The serialization ID of the preprocessor (either explicit or class name)
+
+    Returns:
+        The registered preprocessor class
+
+    Raises:
+        UnknownPreprocessorError: If the serialization ID is not registered
+    """
+    if serialization_id not in _PREPROCESSOR_REGISTRY:
+        raise UnknownPreprocessorError(serialization_id)
+    return _PREPROCESSOR_REGISTRY[serialization_id]
diff --git a/lib/python3.12/site-packages/ray/serve/autoscaling_policy.py b/lib/python3.12/site-packages/ray/serve/autoscaling_policy.py
new file mode 100644
index 0000000000000000000000000000000000000000..bd7b426416b12e856afbb56c7d9f1431cef7b92e
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/serve/autoscaling_policy.py
@@ -0,0 +1,176 @@
+import logging
+import math
+from typing import Any, Dict, Optional, Tuple
+
+from ray.serve._private.constants import CONTROL_LOOP_INTERVAL_S, SERVE_LOGGER_NAME
+from ray.serve.config import AutoscalingConfig, AutoscalingContext
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(SERVE_LOGGER_NAME)
+
+
+def _calculate_desired_num_replicas(
+    autoscaling_config: AutoscalingConfig,
+    total_num_requests: int,
+    num_running_replicas: int,
+    override_min_replicas: Optional[float] = None,
+    override_max_replicas: Optional[float] = None,
+) -> int:
+    """Returns the number of replicas to scale to based on the given metrics.
+
+    Args:
+        autoscaling_config: The autoscaling parameters to use for this
+            calculation.
+        current_num_ongoing_requests (List[float]): A list of the number of
+            ongoing requests for each replica.  Assumes each entry has already
+            been time-averaged over the desired lookback window.
+        override_min_replicas: Overrides min_replicas from the config
+            when calculating the final number of replicas.
+        override_max_replicas: Overrides max_replicas from the config
+            when calculating the final number of replicas.
+
+    Returns:
+        desired_num_replicas: The desired number of replicas to scale to, based
+            on the input metrics and the current number of replicas.
+
+    """
+    if num_running_replicas == 0:
+        raise ValueError("Number of replicas cannot be zero")
+
+    # Example: if error_ratio == 2.0, we have two times too many ongoing
+    # requests per replica, so we desire twice as many replicas.
+    target_num_requests = (
+        autoscaling_config.get_target_ongoing_requests() * num_running_replicas
+    )
+    error_ratio: float = total_num_requests / target_num_requests
+
+    # If error ratio >= 1, then the number of ongoing requests per
+    # replica exceeds the target and we will make an upscale decision,
+    # so we apply the upscale smoothing factor. Otherwise, the number of
+    # ongoing requests per replica is lower than the target and we will
+    # make a downscale decision, so we apply the downscale smoothing
+    # factor.
+    if error_ratio >= 1:
+        scaling_factor = autoscaling_config.get_upscaling_factor()
+    else:
+        scaling_factor = autoscaling_config.get_downscaling_factor()
+
+    # Multiply the distance to 1 by the smoothing ("gain") factor (default=1).
+    smoothed_error_ratio = 1 + ((error_ratio - 1) * scaling_factor)
+    desired_num_replicas = math.ceil(num_running_replicas * smoothed_error_ratio)
+
+    # If desired num replicas is "stuck" because of the smoothing factor
+    # (meaning the traffic is low enough for the replicas to downscale
+    # without the smoothing factor), decrease desired_num_replicas by 1.
+    if (
+        math.ceil(num_running_replicas * error_ratio) < num_running_replicas
+        and desired_num_replicas == num_running_replicas
+    ):
+        desired_num_replicas -= 1
+
+    min_replicas = autoscaling_config.min_replicas
+    max_replicas = autoscaling_config.max_replicas
+    if override_min_replicas is not None:
+        min_replicas = override_min_replicas
+    if override_max_replicas is not None:
+        max_replicas = override_max_replicas
+
+    # Ensure scaled_min_replicas <= desired_num_replicas <= scaled_max_replicas.
+    desired_num_replicas = max(min_replicas, min(max_replicas, desired_num_replicas))
+
+    return desired_num_replicas
+
+
+@PublicAPI(stability="alpha")
+def replica_queue_length_autoscaling_policy(
+    ctx: AutoscalingContext,
+) -> Tuple[int, Dict[str, Any]]:
+    """The default autoscaling policy based on basic thresholds for scaling.
+    There is a minimum threshold for the average queue length in the cluster
+    to scale up and a maximum threshold to scale down. Each period, a 'scale
+    up' or 'scale down' decision is made. This decision must be made for a
+    specified number of periods in a row before the number of replicas is
+    actually scaled. See config options for more details.  Assumes
+    `get_decision_num_replicas` is called once every CONTROL_LOOP_PERIOD_S
+    seconds.
+    """
+
+    curr_target_num_replicas: int = ctx.target_num_replicas
+    total_num_requests: int = ctx.total_num_requests
+    num_running_replicas: int = ctx.current_num_replicas
+    config: Optional[AutoscalingConfig] = ctx.config
+    capacity_adjusted_min_replicas: int = ctx.capacity_adjusted_min_replicas
+    capacity_adjusted_max_replicas: int = ctx.capacity_adjusted_max_replicas
+    policy_state: Dict[str, Any] = ctx.policy_state
+    decision_counter = policy_state.get("decision_counter", 0)
+    if num_running_replicas == 0:
+        # When 0 replicas and queries are queued, scale up the replicas
+        if total_num_requests > 0:
+            return (
+                max(
+                    math.ceil(1 * config.get_upscaling_factor()),
+                    curr_target_num_replicas,
+                ),
+                policy_state,
+            )
+        return curr_target_num_replicas, policy_state
+
+    decision_num_replicas = curr_target_num_replicas
+
+    desired_num_replicas = _calculate_desired_num_replicas(
+        config,
+        total_num_requests,
+        num_running_replicas=num_running_replicas,
+        override_min_replicas=capacity_adjusted_min_replicas,
+        override_max_replicas=capacity_adjusted_max_replicas,
+    )
+    # Scale up.
+    if desired_num_replicas > curr_target_num_replicas:
+        # If the previous decision was to scale down (the counter was
+        # negative), we reset it and then increment it (set to 1).
+        # Otherwise, just increment.
+        if decision_counter < 0:
+            decision_counter = 0
+        decision_counter += 1
+
+        # Only actually scale the replicas if we've made this decision for
+        # 'scale_up_consecutive_periods' in a row.
+        if decision_counter > int(config.upscale_delay_s / CONTROL_LOOP_INTERVAL_S):
+            decision_counter = 0
+            decision_num_replicas = desired_num_replicas
+
+    # Scale down.
+    elif desired_num_replicas < curr_target_num_replicas:
+        # If the previous decision was to scale up (the counter was
+        # positive), reset it to zero before decrementing.
+
+        if decision_counter > 0:
+            decision_counter = 0
+        decision_counter -= 1
+        # Downscaling to zero is only allowed from 1 -> 0
+        is_scaling_to_zero = curr_target_num_replicas == 1
+        # Determine the delay to use
+        if is_scaling_to_zero:
+            # Check if the downscale_to_zero_delay_s is set
+            if config.downscale_to_zero_delay_s is not None:
+                delay_s = config.downscale_to_zero_delay_s
+            else:
+                delay_s = config.downscale_delay_s
+        else:
+            delay_s = config.downscale_delay_s
+            # The desired_num_replicas>0 for downscaling cases other than 1->0
+            desired_num_replicas = max(1, desired_num_replicas)
+        # Only actually scale the replicas if we've made this decision for
+        # 'scale_down_consecutive_periods' in a row.
+        if decision_counter < -int(delay_s / CONTROL_LOOP_INTERVAL_S):
+            decision_counter = 0
+            decision_num_replicas = desired_num_replicas
+    # Do nothing.
+    else:
+        decision_counter = 0
+
+    policy_state["decision_counter"] = decision_counter
+    return decision_num_replicas, policy_state
+
+
+default_autoscaling_policy = replica_queue_length_autoscaling_policy
diff --git a/lib/python3.12/site-packages/ray/serve/grpc_util.py b/lib/python3.12/site-packages/ray/serve/grpc_util.py
new file mode 100644
index 0000000000000000000000000000000000000000..2b258c5a261437041f98afcd2089b51b2e4a48fb
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/serve/grpc_util.py
@@ -0,0 +1,169 @@
+from typing import Any, Dict, List, Optional, Tuple
+
+import grpc
+
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="beta")
+class RayServegRPCContext:
+    """Context manager to set and get gRPC context.
+
+    This class implements most of the methods from ServicerContext
+    (see: https://grpc.github.io/grpc/python/grpc.html#grpc.ServicerContext). It's
+    serializable and can be passed with the request to be used on the deployment.
+    """
+
+    def __init__(self, grpc_context: grpc._cython.cygrpc._ServicerContext):
+        self._auth_context = grpc_context.auth_context()
+        self._code = grpc_context.code()
+        self._details = grpc_context.details()
+        self._invocation_metadata = [  # noqa: C416
+            (key, value) for key, value in grpc_context.invocation_metadata()
+        ]
+        self._peer = grpc_context.peer()
+        self._peer_identities = grpc_context.peer_identities()
+        self._peer_identity_key = grpc_context.peer_identity_key()
+        self._trailing_metadata = [  # noqa: C416
+            (key, value) for key, value in grpc_context.trailing_metadata()
+        ]
+        self._compression = None
+
+    def auth_context(self) -> Dict[str, Any]:
+        """Gets the auth context for the call.
+
+        Returns:
+          A map of strings to an iterable of bytes for each auth property.
+        """
+        return self._auth_context
+
+    def code(self) -> grpc.StatusCode:
+        """Accesses the value to be used as status code upon RPC completion.
+
+        Returns:
+          The StatusCode value for the RPC.
+        """
+        return self._code
+
+    def details(self) -> str:
+        """Accesses the value to be used as detail string upon RPC completion.
+
+        Returns:
+          The details string of the RPC.
+        """
+        return self._details
+
+    def invocation_metadata(self) -> List[Tuple[str, str]]:
+        """Accesses the metadata sent by the client.
+
+        Returns:
+          The invocation :term:`metadata`.
+        """
+        return self._invocation_metadata
+
+    def peer(self) -> str:
+        """Identifies the peer that invoked the RPC being serviced.
+
+        Returns:
+          A string identifying the peer that invoked the RPC being serviced.
+          The string format is determined by gRPC runtime.
+        """
+        return self._peer
+
+    def peer_identities(self) -> Optional[bytes]:
+        """Gets one or more peer identity(s).
+
+        Equivalent to
+        servicer_context.auth_context().get(servicer_context.peer_identity_key())
+
+        Returns:
+          An iterable of the identities, or None if the call is not
+          authenticated. Each identity is returned as a raw bytes type.
+        """
+        return self._peer_identities
+
+    def peer_identity_key(self) -> Optional[str]:
+        """The auth property used to identify the peer.
+
+        For example, "x509_common_name" or "x509_subject_alternative_name" are
+        used to identify an SSL peer.
+
+        Returns:
+          The auth property (string) that indicates the
+          peer identity, or None if the call is not authenticated.
+        """
+        return self._peer_identity_key
+
+    def trailing_metadata(self) -> List[Tuple[str, str]]:
+        return self._trailing_metadata
+
+    def set_code(self, code: grpc.StatusCode):
+        """Sets the value to be used as status code upon RPC completion.
+
+        This method need not be called by method implementations if they wish
+        the gRPC runtime to determine the status code of the RPC.
+
+        Args:
+          code: A StatusCode object to be sent to the client.
+        """
+        self._code = code
+
+    def set_compression(self, compression: grpc.Compression):
+        """Set the compression algorithm to be used for the entire call.
+
+        Args:
+          compression: An element of grpc.compression, e.g.
+            grpc.compression.Gzip.
+        """
+        self._compression = compression
+
+    def set_details(self, details: str):
+        """Sets the value to be used as detail string upon RPC completion.
+
+        Calling this method is only needed if method implementations have
+        details to transmit.
+
+        Args:
+          details: A UTF-8-encodable string to be sent to the client upon
+                   termination of the RPC.
+        """
+        self._details = details
+
+    def _request_id_metadata(self) -> List[Tuple[str, str]]:
+        # Request id metadata should be carried over to the trailing metadata and passed
+        # back to the request client. This function helps pick it out if it exists.
+        for key, value in self._trailing_metadata:
+            if key == "request_id":
+                return [(key, value)]
+        return []
+
+    def set_trailing_metadata(self, trailing_metadata: List[Tuple[str, str]]):
+        """Sets the trailing metadata for the RPC.
+
+        Sets the trailing metadata to be sent upon completion of the RPC.
+
+        If this method is invoked multiple times throughout the lifetime of an
+        RPC, the value supplied in the final invocation + request id will be the value
+        sent over the wire.
+
+        This method need not be called by implementations if they have no
+        metadata to add to what the gRPC runtime will transmit.
+
+        Args:
+          trailing_metadata: The trailing :term:`metadata`.
+        """
+        self._trailing_metadata = self._request_id_metadata() + trailing_metadata
+
+    def _set_on_grpc_context(self, grpc_context: grpc._cython.cygrpc._ServicerContext):
+        """Serve's internal method to set attributes on the gRPC context."""
+        if self._code:
+            grpc_context.set_code(self._code)
+
+        if self._compression:
+            grpc_context.set_compression(self._compression)
+
+        if self._details:
+            grpc_context.set_details(self._details)
+
+        if self._trailing_metadata:
+            grpc_context.set_trailing_metadata(self._trailing_metadata)
diff --git a/lib/python3.12/site-packages/ray/serve/request_router.py b/lib/python3.12/site-packages/ray/serve/request_router.py
new file mode 100644
index 0000000000000000000000000000000000000000..78a42829517519ad9f72a3e7f8939269ef9534ed
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/serve/request_router.py
@@ -0,0 +1,14 @@
+from ray.serve._private.common import ReplicaID  # noqa: F401
+from ray.serve._private.replica_result import ReplicaResult  # noqa: F401
+from ray.serve._private.request_router.common import (  # noqa: F401
+    PendingRequest,
+)
+from ray.serve._private.request_router.replica_wrapper import (  # noqa: F401
+    RunningReplica,
+)
+from ray.serve._private.request_router.request_router import (  # noqa: F401
+    FIFOMixin,
+    LocalityMixin,
+    MultiplexMixin,
+    RequestRouter,
+)
diff --git a/lib/python3.12/site-packages/ray/serve/task_processor.py b/lib/python3.12/site-packages/ray/serve/task_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..92e5f68de90c7b6bc35639ea7343e4ab50d8265f
--- /dev/null
+++ b/lib/python3.12/site-packages/ray/serve/task_processor.py
@@ -0,0 +1,356 @@
+import logging
+import threading
+import time
+from typing import Any, Dict, List, Optional
+
+from celery import Celery
+from celery.signals import task_failure, task_unknown
+
+from ray.serve import get_replica_context
+from ray.serve._private.constants import (
+    DEFAULT_CONSUMER_CONCURRENCY,
+    SERVE_LOGGER_NAME,
+)
+from ray.serve.schema import (
+    CeleryAdapterConfig,
+    TaskProcessorAdapter,
+    TaskProcessorConfig,
+    TaskResult,
+)
+from ray.util.annotations import PublicAPI
+
+logger = logging.getLogger(SERVE_LOGGER_NAME)
+
+
+CELERY_WORKER_POOL = "worker_pool"
+CELERY_WORKER_CONCURRENCY = "worker_concurrency"
+CELERY_TASK_IGNORE_RESULT = "task_ignore_result"
+CELERY_TASK_ACKS_LATE = "task_acks_late"
+CELERY_TASK_REJECT_ON_WORKER_LOST = "task_reject_on_worker_lost"
+
+CELERY_DEFAULT_APP_CONFIG = [
+    CELERY_WORKER_POOL,
+    CELERY_WORKER_CONCURRENCY,
+    CELERY_TASK_IGNORE_RESULT,
+    CELERY_TASK_ACKS_LATE,
+    CELERY_TASK_REJECT_ON_WORKER_LOST,
+]
+
+
+@PublicAPI(stability="alpha")
+class CeleryTaskProcessorAdapter(TaskProcessorAdapter):
+    """
+    Celery-based task processor adapter.
+    This adapter does NOT support any async operations.
+    All operations must be performed synchronously.
+    """
+
+    _app: Celery
+    _config: TaskProcessorConfig
+    _worker_thread: Optional[threading.Thread] = None
+    _worker_hostname: Optional[str] = None
+    _worker_concurrency: int = DEFAULT_CONSUMER_CONCURRENCY
+
+    def __init__(self, config: TaskProcessorConfig, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if not isinstance(config.adapter_config, CeleryAdapterConfig):
+            raise TypeError(
+                "TaskProcessorConfig.adapter_config must be an instance of CeleryAdapterConfig"
+            )
+
+        # Check if any app_custom_config keys conflict with default Celery app config
+        if config.adapter_config.app_custom_config:
+            conflicting_keys = set(
+                config.adapter_config.app_custom_config.keys()
+            ) & set(CELERY_DEFAULT_APP_CONFIG)
+            if conflicting_keys:
+                raise ValueError(
+                    f"The following configuration keys cannot be changed via app_custom_config: {sorted(conflicting_keys)}. "
+                    f"These are managed internally by the CeleryTaskProcessorAdapter."
+                )
+
+        self._config = config
+
+        # Celery adapter does not support any async capabilities
+        # self._async_capabilities is already an empty set from parent class
+
+    def initialize(self, consumer_concurrency: int = DEFAULT_CONSUMER_CONCURRENCY):
+        self._app = Celery(
+            self._config.queue_name,
+            backend=self._config.adapter_config.backend_url,
+            broker=self._config.adapter_config.broker_url,
+        )
+
+        app_configuration = {
+            CELERY_WORKER_POOL: "threads",
+            CELERY_WORKER_CONCURRENCY: consumer_concurrency,
+            CELERY_TASK_IGNORE_RESULT: False,  # Store task results so they can be retrieved after completion
+            CELERY_TASK_ACKS_LATE: True,  # Acknowledge tasks only after completion (not when received) for better reliability
+            CELERY_TASK_REJECT_ON_WORKER_LOST: True,  # Reject and requeue tasks when worker is lost to prevent data loss
+        }
+
+        if self._config.adapter_config.app_custom_config:
+            app_configuration.update(self._config.adapter_config.app_custom_config)
+
+        self._app.conf.update(app_configuration)
+
+        queue_config = {
+            self._config.queue_name: {
+                "exchange": self._config.queue_name,
+                "exchange_type": "direct",
+                "routing_key": self._config.queue_name,
+            },
+        }
+
+        if self._config.failed_task_queue_name:
+            queue_config[self._config.failed_task_queue_name] = {
+                "exchange": self._config.failed_task_queue_name,
+                "exchange_type": "direct",
+                "routing_key": self._config.failed_task_queue_name,
+            }
+
+        if self._config.unprocessable_task_queue_name:
+            queue_config[self._config.unprocessable_task_queue_name] = {
+                "exchange": self._config.unprocessable_task_queue_name,
+                "exchange_type": "direct",
+                "routing_key": self._config.unprocessable_task_queue_name,
+            }
+
+        self._app.conf.update(
+            task_queues=queue_config,
+            task_routes={
+                # Default tasks go to main queue
+                "*": {"queue": self._config.queue_name},
+            },
+        )
+
+        if self._config.adapter_config.broker_transport_options is not None:
+            self._app.conf.update(
+                broker_transport_options=self._config.adapter_config.broker_transport_options,
+            )
+
+        if self._config.failed_task_queue_name:
+            task_failure.connect(self._handle_task_failure)
+
+        if self._config.unprocessable_task_queue_name:
+            task_unknown.connect(self._handle_unknown_task)
+
+    def register_task_handle(self, func, name=None):
+        task_options = {
+            "autoretry_for": (Exception,),
+            "retry_kwargs": {"max_retries": self._config.max_retries},
+            "retry_backoff": True,
+            "retry_backoff_max": 60,  # Max backoff of 60 seconds
+            "retry_jitter": False,  # Disable jitter for predictable testing
+        }
+        if self._config.adapter_config.task_custom_config:
+            task_options.update(self._config.adapter_config.task_custom_config)
+
+        if name:
+            self._app.task(name=name, **task_options)(func)
+        else:
+            self._app.task(**task_options)(func)
+
+    def enqueue_task_sync(
+        self, task_name, args=None, kwargs=None, **options
+    ) -> TaskResult:
+        task_response = self._app.send_task(
+            task_name,
+            args=args,
+            kwargs=kwargs,
+            queue=self._config.queue_name,
+            **options,
+        )
+
+        return TaskResult(
+            id=task_response.id,
+            status=task_response.status,
+            created_at=time.time(),
+            result=task_response.result,
+        )
+
+    def get_task_status_sync(self, task_id) -> TaskResult:
+        task_details = self._app.AsyncResult(task_id)
+        return TaskResult(
+            id=task_details.id,
+            result=task_details.result,
+            status=task_details.status,
+        )
+
+    def start_consumer(self, **kwargs):
+        """Starts the Celery worker thread."""
+        if self._worker_thread is not None and self._worker_thread.is_alive():
+            logger.info("Celery worker thread is already running.")
+            return
+
+        unique_id = get_replica_context().replica_tag
+        self._worker_hostname = f"{self._app.main}_{unique_id}"
+
+        worker_args = [
+            "worker",
+            f"--hostname={self._worker_hostname}",
+            "-Q",
+            self._config.queue_name,
+        ]
+
+        self._worker_thread = threading.Thread(
+            target=self._app.worker_main,
+            args=(worker_args,),
+        )
+        self._worker_thread.start()
+
+        logger.info(
+            f"Celery worker thread started with hostname: {self._worker_hostname}"
+        )
+
+    def stop_consumer(self, timeout: float = 10.0):
+        """Signals the Celery worker to shut down and waits for it to terminate."""
+        if self._worker_thread is None or not self._worker_thread.is_alive():
+            logger.info("Celery worker thread is not running.")
+            return
+
+        logger.info("Sending shutdown signal to Celery worker...")
+
+        # Use the worker's hostname for targeted shutdown
+        self._app.control.broadcast(
+            "shutdown", destination=[f"celery@{self._worker_hostname}"]
+        )
+        self._worker_thread.join(timeout=timeout)
+
+        if self._worker_thread.is_alive():
+            logger.warning(f"Worker thread did not terminate after {timeout} seconds.")
+        else:
+            logger.info("Celery worker thread has stopped.")
+
+        self._worker_thread = None
+
+    def shutdown(self):
+        logger.info("Shutting down Celery worker...")
+        self._app.control.shutdown()
+        logger.info("Celery worker shutdown complete...")
+
+    def cancel_task_sync(self, task_id):
+        """
+        Cancels a task synchronously. Only supported for Redis and RabbitMQ brokers by Celery.
+        More details can be found here: https://docs.celeryq.dev/en/stable/userguide/workers.html#revoke-revoking-tasks
+        """
+        self._app.control.revoke(task_id)
+
+    def get_metrics_sync(self) -> Dict[str, Any]:
+        """
+        Returns the metrics of the Celery worker synchronously.
+        More details can be found here: https://docs.celeryq.dev/en/stable/reference/celery.app.control.html#celery.app.control.Inspect.stats
+        """
+        return self._app.control.inspect().stats()
+
+    def health_check_sync(self) -> List[Dict]:
+        """
+        Checks the health of the Celery worker synchronously.
+        Returns a list of dictionaries, each containing the worker name and a dictionary with the health status.
+        Example: [{'celery@192.168.1.100': {'ok': 'pong'}}]
+        More details can be found here: https://docs.celeryq.dev/en/stable/reference/celery.app.control.html#celery.app.control.Control.ping
+        """
+        return self._app.control.ping()
+
+    def _handle_task_failure(
+        self,
+        sender: Any = None,
+        task_id: str = None,
+        args: Any = None,
+        kwargs: Any = None,
+        einfo: Any = None,
+        **kw,
+    ):
+        """Handle task failures and route them to appropriate dead letter queues.
+
+        This method is called when a task fails after all retry attempts have been
+        exhausted. It logs the failure and moves the task to failed_task_queue
+
+        Args:
+            sender: The task object that failed
+            task_id: Unique identifier of the failed task
+            args: Positional arguments passed to the task
+            kwargs: Keyword arguments passed to the task
+            einfo: Exception info object containing exception details and traceback
+            **kw: Additional keyword arguments passed by Celery
+        """
+        logger.info(
+            f"Task failure detected for task_id: {task_id}, einfo: {str(einfo)}"
+        )
+
+        dlq_args = [
+            task_id,
+            str(einfo.exception),
+            str(args),
+            str(kwargs),
+            str(einfo),
+        ]
+
+        if self._config.failed_task_queue_name:
+            self._move_task_to_queue(
+                self._config.failed_task_queue_name,
+                sender.name,
+                dlq_args,
+            )
+
+            logger.error(
+                f"Task {task_id} failed after max retries. Exception: {einfo}. Moved it to the {self._config.failed_task_queue_name} queue."
+            )
+
+    def _handle_unknown_task(
+        self,
+        sender: Any = None,
+        name: str = None,
+        id: str = None,
+        message: Any = None,
+        exc: Any = None,
+        **kwargs,
+    ):
+        """Handle unknown or unregistered tasks received by Celery.
+
+        This method is called when Celery receives a task that it doesn't recognize
+        (i.e., a task that hasn't been registered with the Celery app). These tasks
+        are moved to the unprocessable task queue if configured.
+
+        Args:
+            sender: The Celery app or worker that detected the unknown task
+            name: Name of the unknown task
+            id: Task ID of the unknown task
+            message: The raw message received for the unknown task
+            exc: The exception raised when trying to process the unknown task
+            **kwargs: Additional context information from Celery
+        """
+        logger.info(
+            f"Unknown task detected by Celery. Name: {name}, ID: {id}, Exc: {str(exc)}"
+        )
+
+        if self._config.unprocessable_task_queue_name:
+            self._move_task_to_queue(
+                self._config.unprocessable_task_queue_name,
+                name,
+                [
+                    name,
+                    id,
+                    str(message),
+                    str(exc),
+                    str(kwargs),
+                ],
+            )
+
+    def _move_task_to_queue(self, queue_name: str, task_name: str, args: list):
+        """Helper function to move a task to a specified queue."""
+        try:
+            logger.info(
+                f"Moving task: {task_name} to queue: {queue_name}, args: {args}"
+            )
+            self._app.send_task(
+                name=task_name,
+                queue=queue_name,
+                args=args,
+            )
+        except Exception as e:
+            logger.error(
+                f"Failed to move task: {task_name} to queue: {queue_name}, error: {e}"
+            )
+            raise e