Handle non-contiguous Tensors based GPU transfer

[srinathk10]

Why are these changes needed?

Handle non-contiguous Tensors based GPU transfer. This allows removing the overhead of combining Arrow chunked arrays during Arrow -> Numpy -> Tensor conversion.

Related issue number

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 :(

[raulchen]

The default should be arrow for better performance.
And we don't need the default numpy / pandas.

[srinathk10]

Yes, fixed now. Changed to Numpy to chase down the issue of null->nan conversion.

[raulchen]

let's add a warning that raw collate_fn will be deprecated. and suggest using ArrowBatchCollateFn for the best performance.

[justinvyu]

i think we should still keep raw collate fn? (and just default to the numpy version).

[raulchen]

yeah, we should keep it. but I'd like to emit a warning to ask users to migrate to the new API.
WDYT?

[raulchen]

I'd like to rewrite this doc-string and make it more descriptive and easy to read.
E.g.

collate_fn: A function that collates data batches before feeding them to the model.  Potential use cases include ...
    If not specified, `iter_torch_batches` converts the data to torch.Tensors and transfer them to the given device. 
    If specified, you can customize the collation logic. The input batch type can be one of .... Arrow is the most recommended for best perf.  And the output can be any type. If the output type is one of ..., the data will be automatically transferred to the given device. Otherwise, you need to transfer the batches in your training loop. 
	Note, this collate_fn will be called in a multi-threaded manner. 

cc @justinvyu for more suggestions

[raulchen]

and the above example needs to be updated as well

[raulchen]

also need to mention how to choose the batch type - subclassing one of ...

[justinvyu]

"""
            collate_fn: [Alpha] A function to customize how data batches are collated before
                being passed to the model. This is useful for last-mile data formatting
                such as padding, masking, or packaging tensors into custom data
                structures. If not provided, `iter_torch_batches` automatically converts
                batches to `torch.Tensor`s and moves them to the device assigned to the
                current worker. The input to `collate_fn` may be:
                    (1) dict of np.ndarray, where you should provide a function
                        that takes in a dict of Numpy arrays
                    (2) pd.DataFrame, where you should provide a callable class
                        that subclasses `PandasCollateFn`
                    (3) pyarrow.Table, where you should provide a callable class
                        that subclasses `ArrowCollateFn` (recommended for best performance)
                The output can be any type. If the output is a `torch.Tensor`,
                `dict[str, torch.Tensor]`, or `list/tuple[torch.Tensor]`, it will be
                automatically moved to the current worker's device. For other types,
                you must handle device transfer manually in your training loop.
                Note: This function is called in a multi-threaded context; avoid using
                thread-unsafe code.
"""

[justinvyu]

"""
            collate_fn: [Alpha] A function to customize how data batches are collated before
                being passed to the model. This is useful for last-mile data formatting
                such as padding, masking, or packaging tensors into custom data
                structures. If not provided, `iter_torch_batches` automatically converts
                batches to `torch.Tensor`s and moves them to the device assigned to the
                current worker. The input to `collate_fn` may be:
                    (1) dict of np.ndarray, where you should provide a function
                        that takes in a dict of Numpy arrays
                    (2) pd.DataFrame, where you should provide a callable class
                        that subclasses `PandasCollateFn`
                    (3) pyarrow.Table, where you should provide a callable class
                        that subclasses `ArrowCollateFn` (recommended for best performance)
                The output can be any type. If the output is a `torch.Tensor`,
                `dict[str, torch.Tensor]`, or `list/tuple[torch.Tensor]`, it will be
                automatically moved to the current worker's device. For other types,
                you must handle device transfer manually in your training loop.
                Note: This function is called in a multi-threaded context; avoid using
                thread-unsafe code.
"""

[justinvyu]

i think we should still keep raw collate fn? (and just default to the numpy version).