[Data][doc] Update on ray.data.Dataset.map() type hints

[chuang0221]

Why are these changes needed?

Currently, the type hints for Dataset.map() suggest it supports both direct returns and generators, but in practice generators are not supported and raise errors. This is misleading for users.

This PR updates the type hints to accurately reflect supported types.

Related issue number

Closes #52279

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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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 :(

[chuang0221]

During investigation, I found that _CallableClassProtocol with Iterator return type also doesn't work with map(), despite what the internal type hint suggests:

# In block.py
class _CallableClassProtocol(Protocol[T, U]):
    def __call__(self, __arg: T) -> Union[U, Iterator[U]]  # Iterator not actually supported for map()

Since this is an internal implementation detail not exposed in the public API, I'm keeping this PR focused on fixing the user-facing documentation. The internal type hint inconsistency could be addressed separately if needed.

[richardliaw]

hmm, I don't think we should expose _CallableClassProtocol as a public typehint since this is an internal class.

Can we use instead just the Callable[[Dict[str, Any]], Dict[str, Any]]?

[chuang0221]

This makes sense to me. I'll change it.

[chuang0221]

I'm curious about the original design here - since UserDefinedFunction is used in other type hints, was there a specific reason for including the internal class _CallableClassProtocol rather than just using Callable in the UserDefinedFunction?

[richardliaw]

hmm, i don't have the context there unfortunately.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[github-actions]

This pull request has been automatically closed because there has been no more activity in the 14 days
since being marked stale.

Please feel free to reopen or open a new pull request if you'd still like this to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for your contribution!