[Data] Break "UDF time" down into input prep, function body and output build

[marwan116]

Why are these changes needed?

ds.stats() reports one figure per operator for a map stage's transform and calls it "UDF time". It isn't the user's function: it spans turning input blocks into batches or rows, calling the function, and assembling the output back into blocks. So a slow operator tells you nothing about whether to optimise your code or your data layout.

That bites hardest when rows carry Python objects, because both ends of the window get expensive. A UDF that never reads an object column is still charged for unpickling it, and nothing in the output says so.

What this change does

Renames the figure to say what it measures, and breaks it down. ds.stats() prints * Block transform time: where it printed * UDF time:, and under verbose_stats_logs splits it into three lines that sum to it:

* Block transform time: 4.02s min, 4.11s max, 4.06s mean, 16.2s total
	* Input prep: ...           forming batches or rows, where object unpickling lands
	* Function body: ...        the stage bodies, yours and Ray Data's alike
	* Output block build: ...   assembling output back into blocks

The total and all three phases are metric_fields on OpRuntimeMetrics, so they reach Prometheus per operator tagged dataset and operator. The three phases stay None rather than zero when a chain is timed only as a whole, so a dashboard can tell "not measured" from "took no time".

Default output is unchanged: the breakdown renders only under verbose_stats_logs, the same treatment extra_metrics gets, and the figures are always on get_stats_summary() regardless. Row transforms (map, filter) report only a total by default, since a timer per row costs roughly 0.6µs across a three-stage chain; DataContext.accurate_map_phase_timing opts them in.

The docs gain a worked example that runs a fused ReadRange->Project->MapBatches(map1)->MapBatches(map2) and reads the three figures off its output, because "which stages land in which bucket?" was the first thing review asked.

Under the hood

map_transformer.py is rewritten by #65855, merged in here: a task's transform is now a flat list of TimedSteps you can print and assert on, each timing itself with no coordination. Subtracting neighbouring totals once at drain recovers each step's own work, at 580 ns/item against the 1016 the per-item cursor cost.

Timing is no longer gated on the chain containing a UDF, so reads and writes are measured too and MapTransformFn.is_udf is gone, along with the separate Built-in stages bucket. Removing the gate surfaced an ordering bug in the checkpoint 2PC chain, fixed in c89db16.

What needs review

1. The name. ray_data_block_transform_time_s is introduced here, so it is free to change now and expensive after a release. @iamjustinhsu suggested total_task_per_block_s.
2. The checkpoint fix wants a checkpoint owner. It belongs in this PR, since the bug is created by this PR's deferred apply and doesn't exist on master, but nobody else has looked at it.
3. Folding Built-in stages into Function body costs the "my code or Ray Data's?" split until per-stage attribution lands in [Data] Report UDF time per fused stage, behind a flag #65810. @iamjustinhsu asked for it; worth confirming the interim is fine.
4. Reads and writes now report a figure where they reported nothing, which changes ds.stats() output for every read operator.

Related issue number

None. Closed #55052 raised adjacent complaints about operator-level metric granularity and was addressed by ds.explain(); it did not cover this.

Checks

• I've signed off every commit (git commit -s).
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've added any new APIs to the API Reference. (DataContext.accurate_map_phase_timing is documented in the DataContext docstring; no new public API)
• I've made sure the tests are passing.
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

Testing

Eight tests added:

test	asserts
test_block_transform_time_phases_sum_to_the_total	the components add up, which is what makes this a decomposition rather than three loose numbers
test_block_transform_time_phases_separate_object_serde	input prep dominates the body for a UDF that only counts rows beside an object column
test_row_transform_phases_are_opt_in	row transforms report only a total by default, and the flag adds detail without moving it
test_eagerly_consuming_stage_body_is_timed	a datasink's I/O reaches Function body instead of being reported as microseconds
test_operator_without_a_udf_is_still_timed	a standalone read reports its phases, and they stay inside the task's wall time
test_map_transformer.py::test_every_output_block_is_timed	every output block of a task is measured, not just the first
test_op_runtime_metrics.py::test_phase_metrics_stay_none_when_unmeasured	a chain timed only as a whole exports no phases, rather than zeros beside a non-zero total
test_op_runtime_metrics.py::test_phase_metrics_accumulate_when_measured	the None default doesn't swallow a measured phase

Existing tests that assert on the rendered ds.stats() string keep their non-verbose expectations unchanged, which is what gating the breakdown buys; the verbose ones gained three lines. The new OpRuntimeMetrics keys are canonicalized to a placeholder in those expectations, as obj_store_mem_used already is, because a trivial UDF's phases round to zero or not depending on the run.

Run locally against a source build, with PYTHONPATH=python/ray/data/tests:

file	result
test_stats.py	105 passed, 2 skipped, 1 failed
test_map.py	252 passed, 1 failed
test_checkpoint.py	60 passed, 13 skipped, 4 errors
test_consumption.py	28 passed, 2 skipped, 3 failed
test_map_transformer.py, unit/test_auto_batch_size.py, test_op_runtime_metrics.py, test_context.py	all passed
datasource/{test_datasink,test_file_datasink,test_parquet}.py	failures identical to the merge base

Every failure above was A/B'd against master and traces to something absent from my machine rather than to this change: the dashboard API server isn't running, pandas 3.0 breaks an untouched conversion helper, and the rest are s3 fixtures. pre-commit run --all-files passes.

AI assistance

AI assistance (Claude Code) was used to investigate the attribution gap, write the change and the tests, and draft this description. Per AGENTS.md, the author has reviewed every changed line and the tests reported above were run locally.

Duplicate check

gh pr list --repo ray-project/ray --state open was searched for udf time phase breakdown input prep and map_transformer phase timing; neither returns anything. The open PRs touching nearby code — #61379 (caching the serialized map transformer), #64090 (per-actor placement groups), #65663 (exposing backpressure policy in diagnostics) — do not touch this timing.

[gemini-code-assist]

Code Review

This pull request decomposes the overall 'UDF time' metric in Ray Data into four distinct phases: input preparation, UDF body execution, output block building, and other fused stages. It introduces these granular metrics across execution stats, operator summaries, and block metadata, with an opt-in configuration (accurate_map_phase_timing) for row-based transforms to avoid overhead. The review feedback highlights a critical missing update to BlockExecStatsBuilder.build that would cause a runtime TypeError when constructing block stats, and suggests a formatting improvement to ensure core UDF phases are consistently displayed in verbose logs even when their sum is zero.

[iamjustinhsu]

Nice, left some questions and some higher-level feedback since I think the code is getting a bit complicated

[iamjustinhsu]

So just from the description of the breakdown(not looking at code yet), I'm still a bit uncertain how the timing phases work. So if there is something like this:

Read -> select_columns -> map1 -> map2

does the function body only include the sum of map1, map2. Does the built-in stage include the sum of Read and select_columns? And then we in between each phase there is input and output batching?

[marwan116]

Built-in stages is now gone: Function body holds all four bodies: ReadRange, Project, map1 and map2. Input prep and Output block build span all four too, since every stage forms its own batches and builds its own blocks.

So the breakdown tells you which phase the time went to, not which stage. Your exact pipeline is now a worked example in the docs.

[iamjustinhsu]

Is UDF time here the same as total task time? Or is UDF time per block?

[marwan116]

Per output block. The scope drains after each output block, so min/max/mean are per block and sum is the operator total. I've put "per output block" into the opening sentence of the bullet.

[marwan116]

Updating this after the rename: the figure is Block transform time now. Still per output block — min, max and mean are across a task's output blocks, and the total is the operator's.

[iamjustinhsu]

I think UDF time could be misleading because it sounds like the function-body time. I think we should rename this to total_task_s or total_task_per_block_s depending on how its broken down

[marwan116]

I went with block_transform_time_s:

1. BlockExecStats.udf_time_s, OperatorStatsSummary.udf_time and OpRuntimeMetrics.udf_time_s are all renamed, so the Prometheus metric is ray_data_block_transform_time_s.
2. ds.stats() now prints Block transform time: instead of UDF time.

[iamjustinhsu]

Higher level feedback for this map transformer (lmk if this makes sense, and whether you want to tackle this in ur PR).

So to give context, during planning operators that launch tasks (read, write, map) use the MapTransformer. The map transformer takes in a MapTransformFn, which also contains pre- and post-processing. I think with all the lazy iterators it's getting harder to interpret what is getting timed, and what is being wrapped, since iterators don't show upstream iterators well. To make this clear, I'm wondering if you could take some time experimenting an approach that is bit more explicit on what's being timed. My suggestion would be:

1. _map_task is the main entrypoint for running a task. It contains referrences to the input blocks.
2. Remove apply_transform in _map_task. Instead, create another function that returns an iterable of MapTransformFn. Suppose we call that get_transform_fns(). It should contain the function bodies as well as well as pre and post processing.
3. Then, for each input block, we should call get_transform_fns(), and apply those functions. So something like

transform_fns = map_transformer.get_transform_fns()
for input_block in blocks:
    for transform_fn in transform_fns:
        # here you can time each of these very explicity
        out = transform_fn.preprocess(input_block)
        out = transform_fn(out)
        out = transform_fn.postprocess(out)

This ^ will require some work because the transform_fns right now accept an input iterable, since everything is passed by reference. If you go this route, we'll need to change the API signature

Also, 2 side notes

• we should treat built-in stages the same as UDFs. I don't think we need a separate breakdown, since at the end of the day it's just a function (whether it be read or write)
• Operators who eagerly compute (like write) should don't need special casing since we are running the function-body per block

[iamjustinhsu]

Also note that there is room for optimization here. We might be able to skip consecutive pre and postprocess in the same fused task, but we can worry about that later

[marwan116]

agreed with you that the lazy chain is hard to follow.

One concrete obstacle with the sketch as written: batches span input blocks.

E.g. with 4 input blocks of 1 row each and batch_size=4, a for input_block in blocks: outer loop can't express that correctly - i.e. it would silently turn batch_size=4 into four batches of 1.

Output block shaping buffers across inputs the same way. So the refactor needs an answer for cross-block buffering.

[marwan116]

@iamjustinhsu I iterated on this with Claude — let me know if the code now reads as sufficiently simpler.

Both side notes are addressed:

1. Built-in stages is gone as a bucket; a read or write body lands in Function body like any UDF's.
2. Deferring each step's apply to its first pull removed the write special case — the upload now happens inside the same window as the pulls. It did surface a real 2PC ordering bug in the checkpoint chain, fixed in c89db16.

[cursor]

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

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit c3e49a8. Configure here.}

[marwan116]

@iamjustinhsu this has moved a lot since your review — worth re-reading the description rather than the old threads.

• [Data] RFC: build a map task's timing from a declared list of steps #65855 is merged in, so map_transformer.py builds on the prototype you suggested.
• Both your requests are done:
  • udf_time is now block_transform_time_s
  • Built-in stages is now folded into Function body.
• Removing the UDF gate surfaced a 2PC ordering bug in the checkpoint write chain. Fixed here, and it's the part I'd most like another pair of eyes on.

Four things need your call; they're listed under What needs review in the description.

[dstrodtman]

Docs-team style pass from Douglas Strodtman (Anyscale docs). Claude Code assisted; I read every comment below and stand behind each one.

Scope: prose style, grammar, and docs conventions only. I'm deliberately staying out of the naming (UDF time → Block transform time) and the phase-model questions in @iamjustinhsu's open threads — those are the Data team's call, and they look like they're still in motion. The new prose reads clearly; the notes below are all mechanical.

Four small fixes inline, all serial-comma and filler-word nits against the Google developer style guide (Ray's stated fallback for anything the Ray writing-style guide doesn't cover). Nothing here blocks — take or leave each individually. No technical review implied, and no approval: a Data maintainer is already reviewing.

[dstrodtman]

Serial comma: with three items, put a comma before the final and (Google developer style guide).

Suggested change

	covers scheduling and writing blocks to the object store. Read, write and map operators all report it, because a read and a
	covers scheduling and writing blocks to the object store. Read, write, and map operators all report it, because a read and a

[dstrodtman]

Note that is filler — the sentence is just as clear without it.

Suggested change

	objects into Arrow. Note that this is separate from the object store write, which Ray Data reports as the
	objects into Arrow. This is separate from the object store write, which Ray Data reports as the

[dstrodtman]

Serial comma before the final and in the four-item list.

Suggested change

	* **Function body** holds all four stage bodies: ``ReadRange``, ``Project``, ``map1`` and ``map2``. Two tasks each slept 0.1
	* **Function body** holds all four stage bodies: ``ReadRange``, ``Project``, ``map1``, and ``map2``. Two tasks each slept 0.1

[dstrodtman]

Serial comma before the final and.

Suggested change

	block, into the input prep, function body and output block build phases, so you can see which part of the transform the
	block, into the input prep, function body, and output block build phases, so you can see which part of the transform the