Revert "[Datasets] Add support for slicing Arrow blocks that contain tensor columns."

python/ray/data/impl/arrow_block.py

@@ -169,15 +169,9 @@ def __next__(self):
     def slice(self, start: int, end: int, copy: bool) -> "pyarrow.Table":
         view = self._table.slice(start, end - start)
         if copy:
-            # TODO(Clark): Find a better method to ensure a copy. This Pandas
-            # roundtrip isn't ideal because:
-            #  1. It may break for nested schemas.
-            #  2. Future optimizations of zero-copy roundtrip with Pandas could
-            #     break this copy guarantee.
-            #  3. It creates two full memory copies,
-            #     Arrow --> Pandas --> Arrow, instead of just one.
-            return pyarrow.Table.from_pandas(
-                view.to_pandas(), schema=view.schema)
+            # TODO(ekl) there must be a cleaner way to force a copy of a table.
+            copy = [c.to_pandas() for c in view.itercolumns()]
+            return pyarrow.Table.from_arrays(copy, schema=self._table.schema)
         else:
             return view
 

python/ray/data/tests/test_dataset.py

@@ -295,40 +295,6 @@ def test_batch_tensors(ray_start_regular_shared):
     assert df.to_dict().keys() == {0, 1}
 
 
-def test_arrow_block_slice_copy():
-    # Test that ArrowBlock slicing properly copies the underlying Arrow
-    # table.
-    n = 20
-    df = pd.DataFrame({"one": list(range(n))})
-    table = pa.Table.from_pandas(df)
-    og_chunk = table.column(0).chunk(0)
-    og_bufs = og_chunk.buffers()
-    a, b = 5, 10
-    expected_slice = table.slice(a, b - a)
-    block_accessor = BlockAccessor.for_block(table)
-
-    # Test with copy.
-    table2 = block_accessor.slice(a, b, True)
-    assert table2.equals(expected_slice)
-    assert table2.schema == table.schema
-    chunk = table2.column(0).chunk(0)
-    assert chunk.offset == 0
-    assert len(chunk) == b - a
-    bufs = chunk.buffers()
-    print(len(bufs))
-    assert bufs[1].address != og_bufs[1].address
-
-    # Test without copy.
-    table2 = block_accessor.slice(a, b, False)
-    assert table2.equals(expected_slice)
-    assert table2.schema == table.schema
-    chunk = table2.column(0).chunk(0)
-    assert chunk.offset == a
-    assert len(chunk) == b - a
-    bufs = chunk.buffers()
-    assert bufs[1].address == og_bufs[1].address
-
-
 def test_tensors(ray_start_regular_shared):
     # Create directly.
     ds = ray.data.range_tensor(5, shape=(3, 5))
@@ -393,37 +359,6 @@ def test_tensor_array_reductions(ray_start_regular_shared):
                                 reducer(arr, axis=0, **np_kwargs))
 
 
-def test_tensor_array_block_slice():
-    # Test that ArrowBlock slicing workers with tensor column extension type.
-    n = 20
-    df = pd.DataFrame({"one": TensorArray(np.array(list(range(n))))})
-    table = pa.Table.from_pandas(df)
-    og_chunk = table.column(0).chunk(0)
-    og_bufs = og_chunk.buffers()
-    a, b = 5, 10
-    expected_slice = table.slice(a, b - a)
-    block_accessor = BlockAccessor.for_block(table)
-
-    # Test with copy.
-    table2 = block_accessor.slice(a, b, True)
-    assert table2.equals(expected_slice)
-    chunk = table2.column(0).chunk(0)
-    assert chunk.offset == 0
-    assert len(chunk) == b - a
-    bufs = chunk.buffers()
-    print(len(bufs))
-    assert bufs[1].address != og_bufs[1].address
-
-    # Test without copy.
-    table2 = block_accessor.slice(a, b, False)
-    assert table2.equals(expected_slice)
-    chunk = table2.column(0).chunk(0)
-    assert chunk.offset == a
-    assert len(chunk) == b - a
-    bufs = chunk.buffers()
-    assert bufs[1].address == og_bufs[1].address
-
-
 def test_arrow_tensor_array_getitem(ray_start_regular_shared):
     outer_dim = 3
     inner_shape = (2, 2, 2)