[Datasets] Can't return ragged array in UDF

[bveeramani]

What happened + What you expected to happen

I tried to apply a function to a batch of ragged images, but I got ValueError: could not broadcast input array from shape (3,300,400) into shape (3,).

Versions / Dependencies

Ray: b0efdb8
Dependencies: https://gist.github.com/bveeramani/2b8678c289ae2faf19d00b7371ddfeba
OS: Linux
Python: 3.7

Reproduction script

import numpy as np
import ray
from torchvision.transforms import Compose, ToTensor, Normalize

dataset = ray.data.read_images("s3://air-example-data-2/imagenet-sample-images/")

transform = Compose([
    ToTensor(), 
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 
])

def map_fn(batch):
    return {"image": np.array([transform(image).numpy() for image in batch], dtype=object)}

dataset.map_batches(map_fn, batch_format="numpy")

RayTaskError(ValueError): ray::_map_block_nosplit() (pid=297, ip=172.28.0.2)
  File "/usr/local/lib/python3.7/dist-packages/ray/data/_internal/compute.py", line 484, in _map_block_nosplit
    for new_block in block_fn(blocks, *fn_args, **fn_kwargs):
  File "/usr/local/lib/python3.7/dist-packages/ray/data/dataset.py", line 597, in transform
    yield from process_next_batch(batch)
  File "/usr/local/lib/python3.7/dist-packages/ray/data/dataset.py", line 579, in process_next_batch
    batch = batch_fn(batch, *fn_args, **fn_kwargs)
  File "<ipython-input-7-2658550a83d0>", line 14, in map_fn
ValueError: could not broadcast input array from shape (3,300,400) into shape (3,)

Issue Severity

Medium: It is a significant difficulty but I can work around it.