[data] [strict mode] Allow returning lists instead of arrays for nump…

…y batches (ray-project#34734)

Allow map_batches UDFs to return {"foo": [1, 2, 3]} in addition to {"foo": np.array([1, 2, 3])} by implicitly casting lists to arrays.

python/ray/data/_internal/planner/map_batches.py

@@ -51,17 +51,29 @@ def validate_batch(batch: Block) -> None:
                 )
 
             if isinstance(batch, collections.abc.Mapping):
-                for key, value in batch.items():
-                    if not isinstance(value, np.ndarray):
+                for key, value in list(batch.items()):
+                    if not isinstance(value, (np.ndarray, list)):
                         raise ValueError(
                             f"Error validating {_truncated_repr(batch)}: "
                             "The `fn` you passed to `map_batches` returned a "
                             f"`dict`. `map_batches` expects all `dict` values "
-                            f"to be of type `numpy.ndarray`, but the value "
+                            f"to be `list` or `np.ndarray` type, but the value "
                             f"corresponding to key {key!r} is of type "
                             f"{type(value)}. To fix this issue, convert "
-                            f"the {type(value)} to a `numpy.ndarray`."
+                            f"the {type(value)} to a `np.ndarray`."
                         )
+                    if isinstance(value, list):
+                        # Try to convert list values into an numpy array via
+                        # np.array(), so users don't need to manually cast.
+                        # NOTE: we don't cast generic iterables, since types like
+                        # `str` are also Iterable.
+                        try:
+                            batch[key] = np.array(value)
+                        except Exception:
+                            raise ValueError(
+                                "Failed to convert column values to numpy array: "
+                                f"({_truncated_repr(value)})."
+                            )
 
         def process_next_batch(batch: DataBatch) -> Iterator[Block]:
             # Apply UDF.

python/ray/data/block.py

@@ -439,7 +439,7 @@ def batch_to_block(batch: DataBatch) -> Block:
                 return ArrowBlockAccessor.numpy_to_block(
                     batch, passthrough_arrow_not_implemented_errors=True
                 )
-            except pa.ArrowNotImplementedError:
+            except (pa.ArrowNotImplementedError, pa.ArrowInvalid):
                 import pandas as pd
 
                 # TODO(ekl) once we support Python objects within Arrow blocks, we

python/ray/data/tests/test_strict_mode.py

@@ -80,6 +80,26 @@ def test_strict_map_output(ray_start_regular_shared, enable_strict_mode):
     ds.map(lambda x: UserDict({"x": object()})).materialize()
 
 
+def test_strict_convert_map_output(ray_start_regular_shared, enable_strict_mode):
+    ds = ray.data.range(1).map_batches(lambda x: {"id": [0, 1, 2, 3]}).materialize()
+    assert ds.take_batch()["id"].tolist() == [0, 1, 2, 3]
+
+    with pytest.raises(ValueError):
+        # Strings not converted into array.
+        ray.data.range(1).map_batches(lambda x: {"id": "string"}).materialize()
+
+    class UserObj:
+        def __eq__(self, other):
+            return isinstance(other, UserObj)
+
+    ds = (
+        ray.data.range(1)
+        .map_batches(lambda x: {"id": [0, 1, 2, UserObj()]})
+        .materialize()
+    )
+    assert ds.take_batch()["id"].tolist() == [0, 1, 2, UserObj()]
+
+
 def test_strict_default_batch_format(ray_start_regular_shared, enable_strict_mode):
     ds = ray.data.range(1)