[Datasets] Refactor iter_rows() for DatasetIterator classes

[c21]

Why are these changes needed?

This is to fix CI test failure introduced in #33180 (example of buildkite failure). We need to add iter_rows() method for StreamSplitDatasetIterator as well.

Related issue number

Closes #33263

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]

This is basically the same as the other impl. Could we make this not an abstract method, and delete the impl in this file?

[ericl]

(move the common impl to dataset_iterator.py, and then delete the methods in stream_split_dataset_iterator.py and dataset_iterator_impl.py)

[c21]

Ok, I put it in this way because PipelinedDatasetIterator.iter_rows has its own impl. So technically these 3 subclasses not sharing same iter_rows impl. But I can put it inside DatasetIterator.iter_rows.

[ericl]

I think that one can also use the same impl...

[ericl]

(We could make get_batch_format() an abstract method instead, like schema)

[c21]

Actually I found PipelinedDatasetIterator.iter_batches() iterate next dataset internally, I don't want to further complicate PipelinedDatasetIterator.iter_batches interface. How about we leave the implenmentation as it is for now? Or @ericl - you have other suggestion to fix the code?

[ericl]

To be clear, this is my proposal, and I believe we can remove most of the other code after this.

class DatasetIterator:
 def iter_rows(...):
   target_format = self._best_batch_format()
   for batch in self.iter_batches(target_format):
        ...
           yield row
           
 def _best_batch_format():
    return "default"   # can override to customize in DatasetIteratorImpl

[c21]

@ericl - ah got it. Updated.

[jianoaix]

What does it mean "best"? Matching the dataset format doesn't necessarily has the best perf, e.g. pyarrow-to-pandas conversion can be faster than pandas-to-pandas conversion.

[c21]

@jianoaix - I am open to a better name, do you have a suggestion for naming?

[jianoaix]

_default_batch_format to align with Dataset (the idea of "best" can be in comment)?
Also since here is a best determined batch format, should it also be used by iter_batches when the batch_format arg is not specified by users?

[c21]

I guess I won't touch the batch format behavior for iter_batches, that would trigger more change and dsicussion across Dataset and DatasetIterator. ideally I want this PR should be simple one to unblock CI test failure. I can do the renaming part.

[jianoaix]

It should be simple but this invites a question why do we have a class member which determines the best format in the context of iterator, but is not actually used for iter_batches().

[c21]

The current behavior for iter_batches is always using the "default" if user does not provide. The iter_rows behavior is different - in bulk execution, it checks Dataset.dataset_format() and uses the returned one, while in streaming execution, it uses "default".

The difference is already happening before this PR, and here we just introduce a private method in DatasetIterator, so it should be easy to change/remove later. Can we defer the behavior change for iter_batches to a future PR?

[jianoaix]

Yeah, SG

[jianoaix]

If this can be a util, why cannot it be the default impl for DatasetIterator?

[c21]

So DatasetIterator has no definition of ds: "Dataset. PipelinedDatasetIterator just uses base_dataset_pipeline. And DatasetIteratorImpl and StreamSplitDatasetIterator uses _base_dataset.