[data] More informative exceptions in block impl (#24665)

python/ray/data/impl/simple_block.py

@@ -36,7 +36,13 @@ class SimpleBlockBuilder(BlockBuilder[T]):
         self._size_estimator.add(item)
 
     def add_block(self, block: List[T]) -> None:
-        assert isinstance(block, list), block
+        if not isinstance(block, list):
+            raise TypeError(
+                f"Got a block of type {type(block)}, expected list. "
+                "If you are mapping a function, ensure it returns an "
+                "object with the expected type. Block:\n"
+                f"{block}"
+            )
         self._items.extend(block)
         for item in block:
             self._size_estimator.add(item)

python/ray/data/impl/table_block.py

@@ -45,7 +45,13 @@ class TableBlockBuilder(BlockBuilder[T]):
         self._uncompacted_size.add(item)
 
     def add_block(self, block: Any) -> None:
-        assert isinstance(block, self._block_type), block
+        if not isinstance(block, self._block_type):
+            raise TypeError(
+                f"Got a block of type {type(block)}, expected {self._block_type}."
+                "If you are mapping a function, ensure it returns an "
+                "object with the expected type. Block:\n"
+                f"{block}"
+            )
         accessor = BlockAccessor.for_block(block)
         self._tables.append(block)
         self._tables_size_bytes += accessor.size_bytes()

python/ray/data/tests/test_dataset.py

@@ -2960,11 +2960,11 @@ def test_groupby_map_groups_merging_invalid_result(ray_start_regular_shared):
     grouped = ds.groupby(lambda x: x)
 
     # The UDF returns None, which is invalid.
-    with pytest.raises(AssertionError):
+    with pytest.raises(TypeError):
         grouped.map_groups(lambda x: None if x == [1] else x)
 
     # The UDF returns a type that's different than the input type, which is invalid.
-    with pytest.raises(AssertionError):
+    with pytest.raises(TypeError):
         grouped.map_groups(lambda x: pd.DataFrame([1]) if x == [1] else x)