[Dataset] Improve str/repr of Dataset to include execution plan

[c21]

Signed-off-by: Cheng Su scnju13@gmail.com

Why are these changes needed?

This is a followup of #31286, we want to improve the Dataset.__repr__() to provide more useful information to users, given lazy execution is default behavior.

The change is to include execution plan (stages as a tree) into Dataset.__repr__(). Currently each stage only has stage name printed out. We shall add more information per stage/operator in the future, which is orthogonal to this PR. This PR is just to print out the existing information we have.

Example:

>>> import ray
>>> ds = ray.data.range(10)
>>> ds = ds.map_batches(lambda x:x)
>>> ds = ds.filter(lambda x: x > 0)
>>> ds = ds.random_shuffle()
>>> ds
RandomShuffle
+- Filter
   +- MapBatches
      +- Dataset(num_blocks=10, num_rows=10, schema=<class 'int'>)
>>> ds.fully_executed()
>>> ds
Dataset(num_blocks=10, num_rows=9, schema=<class 'int'>)

The code change includes:

• Introduce ExecutionPlan.get_plan_as_string() to get the string representation above for the plan.
• Refactor two private methods inside ExecutionPlan - _get_unified_blocks_schema() and _get_num_rows_from_blocks_metadata()
• Change Dataset.__repr__ to call ExecutionPlan.get_plan_as_string() directly.

Related issue number

Closes #31417

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• 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 made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[ericl]

Could we also update the rst docs?

[c21]

Could we also update the rst docs?

@ericl - yeah plan to do that in the same PR here, if reviewers has no objections on the string representation.

[stephanie-wang]

Just a suggestion, but I think it would be nicer to keep the plan history even after fully_executed is called, more like a "cache" call.

[c21]

Just a suggestion, but I think it would be nicer to keep the plan history even after fully_executed is called, more like a "cache" call.

@stephanie-wang - I thought it before, the only thing I am worried about, is the plan gets super long after multiple calls, assuming users only care about the latest Dataset. WDYT? @ericl, @clarkzinzow and @jianoaix.

Alternative is to add Dataset.plan() or Dataset.explain() prints out all plan history.

[jianoaix]

Just a suggestion, but I think it would be nicer to keep the plan history even after fully_executed is called, more like a "cache" call.

@stephanie-wang - I thought it before, the only thing I am worried about, is the plan gets super long after multiple calls, assuming users only care about the latest Dataset. WDYT? @ericl, @clarkzinzow and @jianoaix.

Alternative is to add Dataset.plan() or Dataset.explain() prints out all plan history.

It sounds good to me to have a separate API to display the plan. The repr is used quite often and I think it's too much details for a simple print(ds). It seems not bad idea to just leave plan out of repr as well.

[c21]

Could we also update the rst docs?

@ericl - acutally given @jianoaix is doing change to make from_item() being lazy in parallel. Should we update all Ray Data rst docs in one pass after both PRs are merged? This should save us time to do only one pass for running all code snippets of documentation.

[ericl]

Hmm, for the caching thing I think we should hide the plan if the Dataset is fully independent of the previous stages. if it still has a hidden reference, we should show those previous stages. This might matter since the serialization behavior of the two cases could be different.

[ericl]

I'm going to just merge this, since I think it's a reasonable first step. We can discuss further refinements on a longer timescale.