import sys from typing import Union import numpy as np import pyarrow as pa import pytest from ray.data._internal.arrow_block import ( ArrowBlockAccessor, ArrowBlockBuilder, ArrowBlockColumnAccessor, _get_max_chunk_size, ) from ray.data._internal.arrow_ops.transform_pyarrow import combine_chunked_array, concat from ray.data._internal.tensor_extensions.arrow import ( ArrowTensorArray, ) def simple_array(): return pa.array([1, 2, None, 6], type=pa.int64()) def simple_chunked_array(): return pa.chunked_array([pa.array([1, 2]), pa.array([None, 6])]) def _wrap_as_pa_scalar(v, dtype: pa.DataType): return pa.scalar(v, type=dtype) @pytest.mark.parametrize("arr", [simple_array(), simple_chunked_array()]) @pytest.mark.parametrize("as_py", [True, False]) class TestArrowBlockColumnAccessor: @pytest.mark.parametrize( "ignore_nulls, expected", [ (True, 3), (False, 4), ], ) def test_count(self, arr, ignore_nulls, as_py, expected): accessor = ArrowBlockColumnAccessor(arr) result = accessor.count(ignore_nulls=ignore_nulls, as_py=as_py) if not as_py: expected = _wrap_as_pa_scalar(expected, dtype=pa.int64()) assert result == expected @pytest.mark.parametrize( "ignore_nulls, expected", [ (True, 9), (False, None), ], ) def test_sum(self, arr, ignore_nulls, as_py, expected): accessor = ArrowBlockColumnAccessor(arr) result = accessor.sum(ignore_nulls=ignore_nulls, as_py=as_py) if not as_py: expected = _wrap_as_pa_scalar(expected, dtype=pa.int64()) assert result == expected @pytest.mark.parametrize( "ignore_nulls, expected", [ (True, 1), (False, None), ], ) def test_min(self, arr, ignore_nulls, as_py, expected): accessor = ArrowBlockColumnAccessor(arr) result = accessor.min(ignore_nulls=ignore_nulls, as_py=as_py) if not as_py: expected = _wrap_as_pa_scalar(expected, dtype=pa.int64()) assert result == expected @pytest.mark.parametrize( "ignore_nulls, expected", [ (True, 6), (False, None), ], ) def test_max(self, arr, ignore_nulls, as_py, expected): accessor = ArrowBlockColumnAccessor(arr) result = accessor.max(ignore_nulls=ignore_nulls, as_py=as_py) if not as_py: expected = _wrap_as_pa_scalar(expected, dtype=pa.int64()) assert result == expected @pytest.mark.parametrize( "ignore_nulls, expected", [ (True, 3), (False, None), ], ) def test_mean(self, arr, ignore_nulls, as_py, expected): accessor = ArrowBlockColumnAccessor(arr) result = accessor.mean(ignore_nulls=ignore_nulls, as_py=as_py) if not as_py: expected = _wrap_as_pa_scalar(expected, dtype=pa.float64()) assert result == expected @pytest.mark.parametrize( "provided_mean, expected", [ (3.0, 14.0), (None, 14.0), ], ) def test_sum_of_squared_diffs_from_mean(self, arr, provided_mean, as_py, expected): accessor = ArrowBlockColumnAccessor(arr) result = accessor.sum_of_squared_diffs_from_mean( ignore_nulls=True, mean=provided_mean, as_py=as_py ) if not as_py: expected = _wrap_as_pa_scalar(expected, dtype=pa.float64()) assert result == expected def test_to_pylist(self, arr, as_py): accessor = ArrowBlockColumnAccessor(arr) assert accessor.to_pylist() == arr.to_pylist() @pytest.mark.parametrize( "input_,expected_output", [ # Empty chunked array (pa.chunked_array([], type=pa.int8()), pa.array([], type=pa.int8())), # Fixed-shape tensors ( pa.chunked_array( [ ArrowTensorArray.from_numpy(np.arange(3).reshape(3, 1)), ArrowTensorArray.from_numpy(np.arange(3).reshape(3, 1)), ] ), ArrowTensorArray.from_numpy( np.concatenate( [ np.arange(3).reshape(3, 1), np.arange(3).reshape(3, 1), ] ) ), ), # Ragged (variable-shaped) tensors ( pa.chunked_array( [ ArrowTensorArray.from_numpy(np.arange(3).reshape(3, 1)), ArrowTensorArray.from_numpy(np.arange(5).reshape(5, 1)), ] ), ArrowTensorArray.from_numpy( np.concatenate( [ np.arange(3).reshape(3, 1), np.arange(5).reshape(5, 1), ] ) ), ), # Small (< 2 GiB) arrays ( pa.chunked_array( [ pa.array([1, 2, 3], type=pa.int16()), pa.array([4, 5, 6], type=pa.int16()), ] ), pa.array([1, 2, 3, 4, 5, 6], type=pa.int16()), ), ], ) def test_combine_chunked_array_small( input_, expected_output: Union[pa.Array, pa.ChunkedArray] ): result = combine_chunked_array(input_) assert expected_output.equals(result) @pytest.mark.parametrize( "input_block, fill_column_name, fill_value, expected_output_block", [ ( pa.Table.from_pydict({"a": [0, 1]}), "b", 2, pa.Table.from_pydict({"a": [0, 1], "b": [2, 2]}), ), ( pa.Table.from_pydict({"a": [0, 1]}), "b", pa.scalar(2), pa.Table.from_pydict({"a": [0, 1], "b": [2, 2]}), ), ], ) def test_fill_column(input_block, fill_column_name, fill_value, expected_output_block): block_accessor = ArrowBlockAccessor.for_block(input_block) actual_output_block = block_accessor.fill_column(fill_column_name, fill_value) assert actual_output_block.equals(expected_output_block) def test_add_blocks_with_different_column_names(): builder = ArrowBlockBuilder() builder.add_block(pa.Table.from_pydict({"col1": ["spam"]})) builder.add_block(pa.Table.from_pydict({"col2": ["foo"]})) block = builder.build() expected_table = pa.Table.from_pydict( {"col1": ["spam", None], "col2": [None, "foo"]} ) assert block.equals(expected_table) @pytest.mark.parametrize( "table_data,max_chunk_size_bytes,expected", [ ({"a": []}, 100, None), ({"a": list(range(100))}, 7, 1), ({"a": list(range(100))}, 10, 1), ({"a": list(range(100))}, 25, 3), ({"a": list(range(100))}, 50, 6), ({"a": list(range(100))}, 100, 12), ], ) def test_arrow_block_max_chunk_size(table_data, max_chunk_size_bytes, expected): table = pa.table(table_data) assert _get_max_chunk_size(table, max_chunk_size_bytes) == expected def test_arrow_block_concat(): table1 = pa.table( { "a": [1, 2, 3], "s": [{"x": 1} for _ in range(3)], } ) table2 = pa.table( { "b": [4, 5, 6], } ) concatenated = concat([table1, table2]) assert set(concatenated.column_names) == {"a", "s", "b"} expected = pa.table( { "a": [1, 2, 3, None, None, None], "s": [{"x": 1} for _ in range(3)] + [None] * 3, "b": [None, None, None, 4, 5, 6], } ) assert concatenated.select(["a", "s", "b"]) == expected if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))