diff --git a/python/ray/data/_internal/arrow_block.py b/python/ray/data/_internal/arrow_block.py index 9370398afe07..2bdb15e702db 100644 --- a/python/ray/data/_internal/arrow_block.py +++ b/python/ray/data/_internal/arrow_block.py @@ -539,9 +539,13 @@ def to_pylist(self) -> List[Any]: return self._column.to_pylist() def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray: - # NOTE: Pyarrow < 13.0.0 does not support ``zero_copy_only`` if get_pyarrow_version() < _MIN_PYARROW_VERSION_TO_NUMPY_ZERO_COPY_ONLY: - return self._column.to_numpy() + if isinstance( + self._column, pyarrow.ChunkedArray + ): # NOTE: ChunkedArray in Pyarrow < 13.0.0 does not support ``zero_copy_only`` + return self._column.to_numpy() + else: + return self._column.to_numpy(zero_copy_only=zero_copy_only) return self._column.to_numpy(zero_copy_only=zero_copy_only) diff --git a/python/ray/data/tests/test_block.py b/python/ray/data/tests/test_block.py index f897a8bc4529..5c4b9cfea8de 100644 --- a/python/ray/data/tests/test_block.py +++ b/python/ray/data/tests/test_block.py @@ -156,3 +156,9 @@ def test_find_partitions_duplicates(): assert partitions[1].to_pydict() == {"value": []} # [1,2) assert partitions[2].to_pydict() == {"value": [2, 2, 2, 2, 2]} # [2,3) assert partitions[3].to_pydict() == {"value": []} # >=3 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", __file__])) diff --git a/semgrep.yml b/semgrep.yml index 088afbda5d35..db4eeb392c95 100644 --- a/semgrep.yml +++ b/semgrep.yml @@ -35,7 +35,6 @@ rules: exclude: # FIXME: These tests weren't run in CI, and now they're failing. - "python/ray/data/tests/test_arrow_serialization.py" - - "python/ray/data/tests/test_block.py" - "python/ray/data/tests/test_hash_shuffle.py" - "python/ray/data/tests/test_operator_fusion.py" languages: