[data] Fix pyarrow numpy element issue (#34215)

python/ray/data/_internal/arrow_block.py

@@ -210,6 +210,7 @@ def _build_tensor_row(
                 element = element.type._extension_scalar_to_ndarray(element)
         # For Arrow < 8.0.0, accessing an element in a chunked tensor array produces an
         # ndarray, which we return directly.
+        assert isinstance(element, np.ndarray), type(element)
         return element
 
     def slice(self, start: int, end: int, copy: bool = False) -> "pyarrow.Table":

python/ray/data/_internal/pandas_block.py

@@ -58,10 +58,16 @@ class PandasRow(TableRow):
     """
 
     def __getitem__(self, key: str) -> Any:
+        from ray.data.extensions import TensorArrayElement
+
         col = self._row[key]
         if len(col) == 0:
             return None
         item = col.iloc[0]
+        if isinstance(item, TensorArrayElement):
+            # Getting an item in a Pandas tensor column may return a TensorArrayElement,
+            # which we have to convert to an ndarray.
+            item = item.to_numpy()
         try:
             # Try to interpret this as a numpy-type value.
             # See https://stackoverflow.com/questions/9452775/converting-numpy-dtypes-to-native-python-types.  # noqa: E501

python/ray/data/_internal/table_block.py

@@ -148,11 +148,10 @@ def __init__(self, table: Any):
         self._table = table
 
     def _get_row(self, index: int, copy: bool = False) -> Union[TableRow, np.ndarray]:
-        row = self.slice(index, index + 1, copy=copy)
+        base_row = self.slice(index, index + 1, copy=copy)
+        row = self.ROW_TYPE(base_row)
         if self.is_tensor_wrapper():
-            row = self._build_tensor_row(row)
-        else:
-            row = self.ROW_TYPE(row)
+            row = row[TENSOR_COLUMN_NAME]
         return row
 
     @staticmethod

python/ray/data/tests/test_arrow_serialization.py

@@ -515,6 +515,17 @@ def test_custom_arrow_data_serializer_fallback(
             assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
 
 
+def test_arrow_scalar_conversion(ray_start_regular_shared):
+    ds = ray.data.from_items([1])
+
+    def fn(batch: list):
+        return np.array([1])
+
+    ds = ds.map_batches(fn)
+    res = ds.take()
+    assert res == [1], res
+
+
 def test_custom_arrow_data_serializer_parquet_roundtrip(
     ray_start_regular_shared, tmp_path
 ):