import itertools import threading from unittest.mock import patch import numpy as np import pandas as pd import pyarrow as pa import pytest from packaging.version import parse as parse_version from ray.data._internal.arrow_ops.transform_pyarrow import combine_chunked_array from ray.data._internal.tensor_extensions.arrow import ( ArrowConversionError, ArrowTensorArray, ArrowTensorType, ArrowTensorTypeV2, ArrowVariableShapedTensorArray, ArrowVariableShapedTensorType, FixedShapeTensorFormat, FixedShapeTensorType, _are_contiguous_1d_views, _concat_ndarrays, _extension_array_concat_supported, concat_tensor_arrays, create_arrow_fixed_shape_tensor_type, fixed_shape_extension_scalar_to_ndarray, unify_tensor_arrays, ) from ray.data._internal.tensor_extensions.pandas import TensorArray, TensorDtype from ray.data._internal.tensor_extensions.utils import ( create_ragged_ndarray, ) from ray.data._internal.utils.arrow_utils import get_pyarrow_version @pytest.mark.parametrize( "values", [ [np.zeros((3, 1)), np.zeros((3, 2))], [np.zeros((3,))], ], ) def test_create_ragged_ndarray(values, tensor_format_context): ragged_array = create_ragged_ndarray(values) assert len(ragged_array) == len(values) for actual_array, expected_array in zip(ragged_array, values): np.testing.assert_array_equal(actual_array, expected_array) def test_tensor_array_validation(): # Test unknown input type raises TypeError. with pytest.raises(TypeError): TensorArray(object()) # Test non-primitive element raises TypeError. with pytest.raises(TypeError): TensorArray(np.array([object(), object()])) with pytest.raises(TypeError): TensorArray([object(), object()]) def test_pandas_to_arrow_fixed_shape_tensor_conversion(tensor_format_context): # First, convert Pandas serise w/ nulls to numpy array = pd.Series([1, 2, 3, None], dtype=pd.Int64Dtype).to_numpy().reshape((2, 2)) # First, check on singular tensor of shape (2, 2, 2) input_tensor = np.stack([array, array]) pa_tensor = ArrowTensorArray.from_numpy(input_tensor) res_tensor = pa_tensor.to_numpy_ndarray() np.testing.assert_array_equal(res_tensor, np.stack([array.astype(np.float64)] * 2)) # Next, check "ragged" tensor # - Outermost ndarray is of shape (2,) (dtype='O') # - Internal ndarrays are homogeneously shaped (2, 2) (dtype='O') input_tensor = create_ragged_ndarray([array, array]) pa_tensor = ArrowTensorArray.from_numpy(input_tensor) res_tensor = pa_tensor.to_numpy_ndarray() np.testing.assert_array_equal(res_tensor, np.stack([array.astype(np.float64)] * 2)) def test_pandas_to_arrow_var_shape_tensor_conversion(): # First, convert Pandas series w/ nulls to numpy array = pd.Series([1, 2, 3, None], dtype=pd.Int64Dtype).to_numpy() input_tensor = create_ragged_ndarray([array.reshape(1, 4), array.reshape((2, 2))]) # For ragged arrays, we need to convert each element individually expected_np_tensor = create_ragged_ndarray( [t.astype(np.float64) for t in input_tensor] ) pa_tensor = ArrowVariableShapedTensorArray.from_numpy(input_tensor) res_tensor = pa_tensor.to_numpy() assert len(res_tensor) == len(expected_np_tensor) for actual, expected in zip(res_tensor, expected_np_tensor): np.testing.assert_array_equal(actual, expected) def test_arrow_scalar_tensor_array_roundtrip(tensor_format_context): tensor_format = tensor_format_context arr = np.arange(1000).reshape((10, 1, 100)) ata = ArrowTensorArray.from_numpy(arr) assert isinstance(ata.type, tensor_format.to_type()) assert len(ata) == len(arr) if tensor_format == FixedShapeTensorFormat.ARROW_NATIVE: out = ata.to_numpy_ndarray() else: out = ata.to_numpy() np.testing.assert_array_equal(out, arr) def test_arrow_scalar_tensor_array_roundtrip_boolean(tensor_format_context): arr = np.array([True, False, False, True]) ata = ArrowTensorArray.from_numpy(arr) assert isinstance(ata.type, pa.DataType) assert len(ata) == len(arr) # Zero-copy is not possible since Arrow bitpacks boolean arrays while NumPy does # not. out = ata.to_numpy(zero_copy_only=False) np.testing.assert_array_equal(out, arr) def test_scalar_tensor_array_roundtrip(tensor_format_context): tensor_format = tensor_format_context arr = np.arange(1000).reshape(10, 1, 100) ta = TensorArray(arr) assert isinstance(ta.dtype, TensorDtype) assert len(ta) == len(arr) out = ta.to_numpy() np.testing.assert_array_equal(out, arr) # Check Arrow conversion. ata = ta.__arrow_array__() assert isinstance(ata.type, pa.DataType) assert len(ata) == len(arr) if tensor_format == FixedShapeTensorFormat.ARROW_NATIVE: out = ata.to_numpy_ndarray() else: out = ata.to_numpy() np.testing.assert_array_equal(out, arr) def test_arrow_variable_shaped_tensor_array_validation(tensor_format_context): # Test tensor elements with differing dimensions raises ValueError. with pytest.raises(ValueError): ArrowVariableShapedTensorArray.from_numpy([np.ones((2, 2)), np.ones((3, 3, 3))]) # Test arbitrary object raises ValueError. with pytest.raises(ValueError): ArrowVariableShapedTensorArray.from_numpy(object()) # Test empty array raises ValueError. with pytest.raises(ValueError): ArrowVariableShapedTensorArray.from_numpy(np.array([])) # Test deeply ragged tensor raises ValueError. with pytest.raises(ValueError): ArrowVariableShapedTensorArray.from_numpy( np.array( [ np.array( [ np.array([1, 2]), np.array([3, 4, 5]), ], dtype=object, ), np.array( [ np.array([5, 6, 7, 8]), ], dtype=object, ), np.array( [ np.array([5, 6, 7, 8]), np.array([5, 6, 7, 8]), np.array([5, 6, 7, 8]), ], dtype=object, ), ], dtype=object, ) ) def test_arrow_variable_shaped_tensor_array_roundtrip(restore_data_context): shapes = [(2, 2), (3, 3), (4, 4)] cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]]) arrs = [ np.arange(offset, offset + np.prod(shape)).reshape(shape) for offset, shape in zip(cumsum_sizes, shapes) ] arr = create_ragged_ndarray(arrs) ata = ArrowVariableShapedTensorArray.from_numpy(arr) assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) out = ata.to_numpy() for o, a in zip(out, arr): np.testing.assert_array_equal(o, a) def test_arrow_variable_shaped_tensor_array_roundtrip_boolean(restore_data_context): arr = np.array( [[True, False], [False, False, True], [False], [True, True, False, True]], dtype=object, ) ata = ArrowVariableShapedTensorArray.from_numpy(arr) assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) out = ata.to_numpy() for o, a in zip(out, arr): np.testing.assert_array_equal(o, a) def test_arrow_variable_shaped_tensor_array_roundtrip_contiguous_optimization( restore_data_context, ): # Test that a roundtrip on slices of an already-contiguous 1D base array does not # create any unnecessary copies. base = np.arange(6) base_address = base.__array_interface__["data"][0] arr = np.array([base[:2], base[2:]], dtype=object) ata = ArrowVariableShapedTensorArray.from_numpy(arr) assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) assert ata.storage.field("data").buffers()[3].address == base_address out = ata.to_numpy() for o, a in zip(out, arr): assert o.base.address == base_address np.testing.assert_array_equal(o, a) def test_arrow_variable_shaped_tensor_array_slice(restore_data_context): shapes = [(2, 2), (3, 3), (4, 4)] cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]]) arrs = [ np.arange(offset, offset + np.prod(shape)).reshape(shape) for offset, shape in zip(cumsum_sizes, shapes) ] arr = np.array(arrs, dtype=object) ata = ArrowVariableShapedTensorArray.from_numpy(arr) assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) indices = [0, 1, 2] for i in indices: np.testing.assert_array_equal(ata[i], arr[i]) slices = [ slice(0, 1), slice(1, 2), slice(2, 3), slice(0, 2), slice(1, 3), slice(0, 3), ] for slice_ in slices: ata_slice = ata[slice_] ata_slice_np = ata_slice.to_numpy() arr_slice = arr[slice_] # Check for equivalent dtypes and shapes. assert ata_slice_np.dtype == arr_slice.dtype assert ata_slice_np.shape == arr_slice.shape # Iteration over tensor array slices triggers NumPy conversion. for o, e in zip(ata_slice, arr_slice): np.testing.assert_array_equal(o, e) def test_arrow_variable_shaped_bool_tensor_array_slice(restore_data_context): arr = np.array( [ [True], [True, False], [False, True, False], ], dtype=object, ) ata = ArrowVariableShapedTensorArray.from_numpy(arr) assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) indices = [0, 1, 2] for i in indices: np.testing.assert_array_equal(ata[i], arr[i]) slices = [ slice(0, 1), slice(1, 2), slice(2, 3), slice(0, 2), slice(1, 3), slice(0, 3), ] for slice_ in slices: ata_slice = ata[slice_] ata_slice_np = ata_slice.to_numpy() arr_slice = arr[slice_] # Check for equivalent dtypes and shapes. assert ata_slice_np.dtype == arr_slice.dtype assert ata_slice_np.shape == arr_slice.shape # Iteration over tensor array slices triggers NumPy conversion. for o, e in zip(ata_slice, arr_slice): np.testing.assert_array_equal(o, e) def test_arrow_variable_shaped_string_tensor_array_slice(restore_data_context): arr = np.array( [ ["Philip", "J", "Fry"], ["Leela", "Turanga"], ["Professor", "Hubert", "J", "Farnsworth"], ["Lrrr"], ], dtype=object, ) ata = ArrowVariableShapedTensorArray.from_numpy(arr) assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) indices = [0, 1, 2, 3] for i in indices: np.testing.assert_array_equal(ata[i], arr[i]) slices = [ slice(0, 1), slice(1, 2), slice(2, 3), slice(3, 4), slice(0, 2), slice(1, 3), slice(2, 4), slice(0, 3), slice(1, 4), slice(0, 4), ] for slice_ in slices: ata_slice = ata[slice_] ata_slice_np = ata_slice.to_numpy() arr_slice = arr[slice_] # Check for equivalent dtypes and shapes. assert ata_slice_np.dtype == arr_slice.dtype assert ata_slice_np.shape == arr_slice.shape # Iteration over tensor array slices triggers NumPy conversion. for o, e in zip(ata_slice, arr_slice): np.testing.assert_array_equal(o, e) def test_variable_shaped_tensor_array_roundtrip(restore_data_context): shapes = [(2, 2), (3, 3), (4, 4)] cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]]) arrs = [ np.arange(offset, offset + np.prod(shape)).reshape(shape) for offset, shape in zip(cumsum_sizes, shapes) ] arr = np.array(arrs, dtype=object) ta = TensorArray(arr) assert isinstance(ta.dtype, TensorDtype) assert len(ta) == len(arr) out = ta.to_numpy() for o, a in zip(out, arr): np.testing.assert_array_equal(o, a) # Check Arrow conversion. ata = ta.__arrow_array__() assert isinstance(ata.type, ArrowVariableShapedTensorType) assert len(ata) == len(arr) out = ata.to_numpy() for o, a in zip(out, arr): np.testing.assert_array_equal(o, a) def test_variable_shaped_tensor_array_slice(restore_data_context): shapes = [(2, 2), (3, 3), (4, 4)] cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]]) arrs = [ np.arange(offset, offset + np.prod(shape)).reshape(shape) for offset, shape in zip(cumsum_sizes, shapes) ] arr = np.array(arrs, dtype=object) ta = TensorArray(arr) assert isinstance(ta.dtype, TensorDtype) assert len(ta) == len(arr) indices = [0, 1, 2] for i in indices: np.testing.assert_array_equal(ta[i], arr[i]) slices = [ slice(0, 1), slice(1, 2), slice(2, 3), slice(0, 2), slice(1, 3), slice(0, 3), ] for slice_ in slices: for o, e in zip(ta[slice_], arr[slice_]): np.testing.assert_array_equal(o, e) def test_tensor_array_ops(tensor_format_context): outer_dim = 3 inner_shape = (2, 2, 2) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items).reshape(shape) df = pd.DataFrame({"one": [1, 2, 3], "two": TensorArray(arr)}) def apply_arithmetic_ops(arr): return 2 * (arr + 1) / 3 def apply_comparison_ops(arr): return arr % 2 == 0 def apply_logical_ops(arr): return arr & (3 * arr) | (5 * arr) # Op tests, using NumPy as the groundtruth. np.testing.assert_equal(apply_arithmetic_ops(arr), apply_arithmetic_ops(df["two"])) np.testing.assert_equal(apply_comparison_ops(arr), apply_comparison_ops(df["two"])) np.testing.assert_equal(apply_logical_ops(arr), apply_logical_ops(df["two"])) def test_tensor_array_array_protocol(tensor_format_context): outer_dim = 3 inner_shape = (2, 2, 2) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items).reshape(shape) t_arr = TensorArray(arr) np.testing.assert_array_equal( np.asarray(t_arr, dtype=np.float32), arr.astype(np.float32) ) t_arr_elem = t_arr[0] np.testing.assert_array_equal( np.asarray(t_arr_elem, dtype=np.float32), arr[0].astype(np.float32) ) def test_tensor_array_dataframe_repr(tensor_format_context): outer_dim = 3 inner_shape = (2, 2) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items).reshape(shape) t_arr = TensorArray(arr) df = pd.DataFrame({"a": t_arr}) expected_repr = """ a 0 [[ 0, 1], [ 2, 3]] 1 [[ 4, 5], [ 6, 7]] 2 [[ 8, 9], [10, 11]]""" assert repr(df) == expected_repr def test_tensor_array_scalar_cast(tensor_format_context): outer_dim = 3 inner_shape = (1,) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items).reshape(shape) t_arr = TensorArray(arr) for t_arr_elem, arr_elem in zip(t_arr, arr): assert float(t_arr_elem) == float(arr_elem) arr = np.arange(1).reshape((1, 1, 1)) t_arr = TensorArray(arr) assert float(t_arr) == float(arr) def test_tensor_array_reductions(tensor_format_context): outer_dim = 3 inner_shape = (2, 2, 2) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items).reshape(shape) df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arr)}) # Reduction tests, using NumPy as the groundtruth. for name, reducer in TensorArray.SUPPORTED_REDUCERS.items(): np_kwargs = {} if name in ("std", "var"): # Pandas uses a ddof default of 1 while NumPy uses 0. # Give NumPy a ddof kwarg of 1 in order to ensure equivalent # standard deviation calculations. np_kwargs["ddof"] = 1 np.testing.assert_equal(df["two"].agg(name), reducer(arr, axis=0, **np_kwargs)) @pytest.mark.parametrize("shape", [(2, 0), (2, 5, 0), (0, 5), (0, 0)]) def test_zero_length_arrow_tensor_array_roundtrip(tensor_format_context, shape): arr = np.empty(shape, dtype=np.int8) t_arr = ArrowTensorArray.from_numpy(arr) assert len(t_arr) == len(arr) out = t_arr.to_numpy_ndarray() np.testing.assert_array_equal(out, arr) @pytest.mark.parametrize("chunked", [False, True]) def test_arrow_tensor_array_getitem(chunked, tensor_format_context): tensor_format = tensor_format_context outer_dim = 3 inner_shape = (2, 2, 2) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items).reshape(shape) t_arr = ArrowTensorArray.from_numpy(arr) if chunked: t_arr = pa.chunked_array(t_arr) pyarrow_version = get_pyarrow_version() if ( chunked and pyarrow_version >= parse_version("8.0.0") and pyarrow_version < parse_version("9.0.0") ): for idx in range(outer_dim): item = t_arr[idx] assert isinstance(item, pa.ExtensionScalar) item = item.type._extension_scalar_to_ndarray(item) np.testing.assert_array_equal(item, arr[idx]) else: for idx in range(outer_dim): item = t_arr[idx] if pyarrow_version >= parse_version("16.0.0"): # Returns native FixedShapeTensorScalar np.testing.assert_array_equal(item.to_numpy(), arr[idx]) else: # Returns an ExtensionScalar, item.type: FixedShapeTensorType np.testing.assert_array_equal( fixed_shape_extension_scalar_to_ndarray(item), arr[idx] ) # NOTE: In addition we verify that for existing ``ArrowTensorScalar`` # implements `__array__` method therefore implementing Numpy # array protocol if tensor_format != FixedShapeTensorFormat.ARROW_NATIVE: np.testing.assert_array_equal(item, arr[idx]) # Test __iter__. for t_subarr, subarr in zip(t_arr, arr): if pyarrow_version >= parse_version("16.0.0"): # Returns native FixedShapeTensorScalar np.testing.assert_array_equal(t_subarr.to_numpy(), subarr) else: # Returns an ExtensionScalar np.testing.assert_array_equal( fixed_shape_extension_scalar_to_ndarray(t_subarr), subarr ) # Test to_pylist. # Note: With tensor_format_context fixture, ARROW_NATIVE is only tested when # FixedShapeTensorType is available (PyArrow >= 12), so no fallback needed. if tensor_format == FixedShapeTensorFormat.ARROW_NATIVE: np.testing.assert_array_equal(t_arr.to_pylist(), arr.reshape(outer_dim, -1)) else: np.testing.assert_array_equal(t_arr.to_pylist(), list(arr)) # Test slicing and indexing. t_arr2 = t_arr[1:] if chunked: # For extension arrays, ChunkedArray.to_numpy() concatenates chunk storage # arrays and calls to_numpy() on the resulting array, which returns the wrong # ndarray. # TODO(Clark): Fix this in Arrow by (1) providing an ExtensionArray hook for # concatenation, and (2) using that + a to_numpy() call on the resulting # ExtensionArray. t_arr2_npy = t_arr2.chunk(0).to_numpy_ndarray() else: t_arr2_npy = t_arr2.to_numpy_ndarray() np.testing.assert_array_equal(t_arr2_npy, arr[1:]) if ( chunked and pyarrow_version >= parse_version("8.0.0") and pyarrow_version < parse_version("9.0.0") and tensor_format != FixedShapeTensorFormat.ARROW_NATIVE ): for idx in range(1, outer_dim): item = t_arr2[idx - 1] assert isinstance(item, pa.ExtensionScalar) item = item.type._extension_scalar_to_ndarray(item) np.testing.assert_array_equal(item, arr[idx]) else: for idx in range(1, outer_dim): item = t_arr2[idx - 1] if pyarrow_version >= parse_version("16.0.0"): # Returns native FixedShapeTensorScalar np.testing.assert_array_equal(item.to_numpy(), arr[idx]) else: # Returns an ExtensionScalar np.testing.assert_array_equal( fixed_shape_extension_scalar_to_ndarray(item), arr[idx] ) @pytest.mark.parametrize("chunked", [False, True]) def test_arrow_variable_shaped_tensor_array_getitem(chunked, tensor_format_context): shapes = [(2, 2), (3, 3), (4, 4)] outer_dim = len(shapes) cumsum_sizes = np.cumsum([0] + [np.prod(shape) for shape in shapes[:-1]]) arrs = [ np.arange(offset, offset + np.prod(shape)).reshape(shape) for offset, shape in zip(cumsum_sizes, shapes) ] arr = np.array(arrs, dtype=object) t_arr = ArrowVariableShapedTensorArray.from_numpy(arr) if chunked: t_arr = pa.chunked_array(t_arr) pyarrow_version = get_pyarrow_version() if ( chunked and pyarrow_version >= parse_version("8.0.0") and pyarrow_version < parse_version("9.0.0") ): for idx in range(outer_dim): item = t_arr[idx] assert isinstance(item, pa.ExtensionScalar) item = item.type._extension_scalar_to_ndarray(item) np.testing.assert_array_equal(item, arr[idx]) else: for idx in range(outer_dim): np.testing.assert_array_equal(t_arr[idx], arr[idx]) # Test __iter__. for t_subarr, subarr in zip(t_arr, arr): np.testing.assert_array_equal(t_subarr, subarr) # Test to_pylist. for t_subarr, subarr in zip(t_arr.to_pylist(), list(arr)): np.testing.assert_array_equal(t_subarr, subarr) # Test slicing and indexing. t_arr2 = t_arr[1:] if chunked: # For extension arrays, ChunkedArray.to_numpy() concatenates chunk storage # arrays and calls to_numpy() on the resulting array, which returns the wrong # ndarray. # TODO(Clark): Fix this in Arrow by (1) providing an ExtensionArray hook for # concatenation, and (2) using that + a to_numpy() call on the resulting # ExtensionArray. t_arr2_npy = t_arr2.chunk(0).to_numpy() else: t_arr2_npy = t_arr2.to_numpy() for t_subarr, subarr in zip(t_arr2_npy, arr[1:]): np.testing.assert_array_equal(t_subarr, subarr) if ( chunked and pyarrow_version >= parse_version("8.0.0") and pyarrow_version < parse_version("9.0.0") ): for idx in range(1, outer_dim): item = t_arr2[idx - 1] assert isinstance(item, pa.ExtensionScalar) item = item.type._extension_scalar_to_ndarray(item) np.testing.assert_array_equal(item, arr[idx]) else: for idx in range(1, outer_dim): np.testing.assert_array_equal(t_arr2[idx - 1], arr[idx]) @pytest.mark.parametrize( "test_arr,dtype", [ ([[1, 2], [3, 4], [5, 6], [7, 8]], None), ([[1, 2], [3, 4], [5, 6], [7, 8]], np.int32), ([[1, 2], [3, 4], [5, 6], [7, 8]], np.int16), ([[1, 2], [3, 4], [5, 6], [7, 8]], np.longlong), ([[1.5, 2.5], [3.3, 4.2], [5.2, 6.9], [7.6, 8.1]], None), ([[1.5, 2.5], [3.3, 4.2], [5.2, 6.9], [7.6, 8.1]], np.float32), ([[1.5, 2.5], [3.3, 4.2], [5.2, 6.9], [7.6, 8.1]], np.float16), ([["B", "A"], ["A", "B"], ["A", "A"], ["B", "B"]], None), ([[False, True], [True, False], [True, True], [False, False]], None), ], ) def test_arrow_tensor_array_slice(test_arr, dtype, tensor_format_context): # Test that ArrowTensorArray slicing works as expected. arr = np.array(test_arr, dtype=dtype) ata = ArrowTensorArray.from_numpy(arr) np.testing.assert_array_equal(ata.to_numpy_ndarray(), arr) slice1 = ata.slice(0, 2) np.testing.assert_array_equal(slice1.to_numpy_ndarray(), arr[0:2]) np.testing.assert_array_equal(slice1[1].as_py(), arr[1]) slice2 = ata.slice(2, 2) np.testing.assert_array_equal(slice2.to_numpy_ndarray(), arr[2:4]) np.testing.assert_array_equal(slice2[1].as_py(), arr[3]) pytest_tensor_array_concat_shapes = [(1, 2, 2), (3, 2, 2), (2, 3, 3)] pytest_tensor_array_concat_arrs = [ np.arange(np.prod(shape)).reshape(shape) for shape in pytest_tensor_array_concat_shapes ] pytest_tensor_array_concat_arrs += [ create_ragged_ndarray( [np.arange(4).reshape((2, 2)), np.arange(4, 13).reshape((3, 3))] ) ] pytest_tensor_array_concat_arr_combinations = list( itertools.combinations(pytest_tensor_array_concat_arrs, 2) ) @pytest.mark.parametrize("a1,a2", pytest_tensor_array_concat_arr_combinations) def test_tensor_array_concat(a1, a2, tensor_format_context): ta1 = TensorArray(a1) ta2 = TensorArray(a2) ta = TensorArray._concat_same_type([ta1, ta2]) assert len(ta) == a1.shape[0] + a2.shape[0] assert ta.dtype.element_dtype == ta1.dtype.element_dtype if a1.shape[1:] == a2.shape[1:]: assert ta.dtype.element_shape == a1.shape[1:] np.testing.assert_array_equal(ta.to_numpy(), np.concatenate([a1, a2])) else: assert ta.dtype.element_shape == (None,) * (len(a1.shape) - 1) for arr, expected in zip( ta.to_numpy(), np.array([e for a in [a1, a2] for e in a], dtype=object) ): np.testing.assert_array_equal(arr, expected) @pytest.mark.parametrize("a1,a2", pytest_tensor_array_concat_arr_combinations) def test_arrow_tensor_array_concat(a1, a2, tensor_format_context): tensor_format = tensor_format_context ta1 = ArrowTensorArray.from_numpy(a1) ta2 = ArrowTensorArray.from_numpy(a2) ta = concat_tensor_arrays([ta1, ta2]) assert len(ta) == a1.shape[0] + a2.shape[0] if a1.shape[1:] == a2.shape[1:]: # With tensor_format_context, ARROW_NATIVE is only tested when # FixedShapeTensorType is available, so to_type() is safe to use assert isinstance(ta.type, tensor_format.to_type()) assert ta.type.storage_type == ta1.type.storage_type assert ta.type.storage_type == ta2.type.storage_type assert tuple(ta.type.shape) == a1.shape[1:] np.testing.assert_array_equal(ta.to_numpy_ndarray(), np.concatenate([a1, a2])) else: assert isinstance(ta.type, ArrowVariableShapedTensorType) assert pa.types.is_struct(ta.type.storage_type) for arr, expected in zip( ta.to_numpy(), np.array([e for a in [a1, a2] for e in a], dtype=object) ): np.testing.assert_array_equal(arr, expected) def test_variable_shaped_tensor_array_chunked_concat(tensor_format_context): # Test that chunking a tensor column and concatenating its chunks preserves typing # and underlying data. shape1 = (2, 2, 2) shape2 = (3, 4, 4) a1 = np.arange(np.prod(shape1)).reshape(shape1) a2 = np.arange(np.prod(shape2)).reshape(shape2) ta1 = ArrowTensorArray.from_numpy(a1) ta2 = ArrowTensorArray.from_numpy(a2) unified_arrs = unify_tensor_arrays([ta1, ta2]) ta = concat_tensor_arrays(unified_arrs) assert len(ta) == shape1[0] + shape2[0] assert isinstance(ta.type, ArrowVariableShapedTensorType) assert pa.types.is_struct(ta.type.storage_type) for arr, expected in zip( ta.to_numpy(), np.array([e for a in [a1, a2] for e in a], dtype=object) ): np.testing.assert_array_equal(arr, expected) def test_variable_shaped_tensor_array_uniform_dim(tensor_format_context): shape1 = (3, 2, 2) shape2 = (3, 4, 4) a1 = np.arange(np.prod(shape1)).reshape(shape1) a2 = np.arange(np.prod(shape2)).reshape(shape2) ta = TensorArray([a1, a2]) assert len(ta) == 2 assert ta.is_variable_shaped for a, expected in zip(ta.to_numpy(), [a1, a2]): np.testing.assert_array_equal(a, expected) def test_large_arrow_tensor_array(tensor_format_context): tensor_format = tensor_format_context test_arr = np.ones((1000, 550), dtype=np.uint8) if tensor_format == FixedShapeTensorFormat.V1: with pytest.raises(ArrowConversionError) as exc_info: ta = ArrowTensorArray.from_numpy([test_arr] * 4000) assert ( repr(exc_info.value.__cause__) == "ArrowInvalid('Negative offsets in list array')" ) else: ta = ArrowTensorArray.from_numpy([test_arr] * 4000) assert len(ta) == 4000 ta = ta.to_numpy_ndarray() for arr in ta: assert arr.shape == (1000, 550) def test_tensor_array_string_tensors_simple(tensor_format_context): """Simple test for fixed-shape string tensor arrays with pandas/arrow roundtrip.""" # Create fixed-shape string tensor string_tensors = np.array( [["hello", "world"], ["arrow", "pandas"], ["tensor", "string"]] ) # Create pandas DataFrame with TensorArray df_pandas = pd.DataFrame({"id": [1, 2, 3], "strings": TensorArray(string_tensors)}) # Convert to Arrow table arrow_table = pa.Table.from_pandas(df_pandas) # Verify the roundtrip preserves the data original_strings = df_pandas["strings"].to_numpy() roundtrip_strings = combine_chunked_array(arrow_table["strings"]).to_numpy_ndarray() np.testing.assert_array_equal(original_strings, roundtrip_strings) np.testing.assert_array_equal(roundtrip_strings, string_tensors) def test_tensor_type_equality_checks(): # Test that different types are not equal fs_tensor_type_v1 = ArrowTensorType((2, 3), pa.int64()) fs_tensor_type_v2 = ArrowTensorTypeV2((2, 3), pa.int64()) assert fs_tensor_type_v1 != fs_tensor_type_v2 # Test different shapes/dtypes aren't equal assert fs_tensor_type_v1 != ArrowTensorType((3, 3), pa.int64()) assert fs_tensor_type_v1 != ArrowTensorType((2, 3), pa.float64()) assert fs_tensor_type_v2 != ArrowTensorTypeV2((3, 3), pa.int64()) assert fs_tensor_type_v2 != ArrowTensorTypeV2((2, 3), pa.float64()) # Test var-shaped tensor type vs_tensor_type = ArrowVariableShapedTensorType(pa.int64(), 2) # Test that different types are not equal assert vs_tensor_type == ArrowVariableShapedTensorType(pa.int64(), 3) assert vs_tensor_type != ArrowVariableShapedTensorType(pa.float64(), 2) assert vs_tensor_type != fs_tensor_type_v1 assert vs_tensor_type != fs_tensor_type_v2 class TestCreateFixedShapeTensorType: """Tests for the create_arrow_fixed_shape_tensor_type factory function.""" @pytest.mark.parametrize( "tensor_format,expected_type_if_native_available,expected_type_fallback", [ # V1 (FixedShapeTensorFormat.V1, ArrowTensorType, ArrowTensorType), # V2 (default) (FixedShapeTensorFormat.V2, ArrowTensorTypeV2, ArrowTensorTypeV2), # NATIVE with V2 fallback when NATIVE unavailable ( FixedShapeTensorFormat.ARROW_NATIVE, FixedShapeTensorType, ArrowTensorTypeV2, ), ], ) def test_context_defaults( self, tensor_format_context, expected_type_if_native_available, expected_type_fallback, ): """Test default tensor type based on context settings with fallback behavior.""" tensor_type = create_arrow_fixed_shape_tensor_type( shape=(2, 3), dtype=pa.int64() ) if FixedShapeTensorType is not None: assert isinstance(tensor_type, expected_type_if_native_available) else: assert isinstance(tensor_type, expected_type_fallback) @pytest.mark.parametrize( "dtype", [pa.int8(), pa.int16(), pa.int32(), pa.int64(), pa.float32(), pa.float64()], ) def test_various_dtypes(self, tensor_format_context, dtype): """Test factory works with various PyArrow dtypes across all formats.""" tensor_type = create_arrow_fixed_shape_tensor_type(shape=(2, 2), dtype=dtype) assert tensor_type.value_type == dtype @pytest.mark.skipif( not _extension_array_concat_supported(), reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0", ) def test_arrow_fixed_shape_tensor_format_eq_with_concat(tensor_format_context): """Test that ArrowTensorType, ArrowTensorTypeV2, and native tensor type __eq__ methods work correctly when concatenating Arrow arrays with the same tensor type.""" tensor_format = tensor_format_context # Create the appropriate tensor type based on format if tensor_format == FixedShapeTensorFormat.V1: tensor_type = ArrowTensorType((2, 3), pa.int64()) elif tensor_format == FixedShapeTensorFormat.V2: tensor_type = ArrowTensorTypeV2((2, 3), pa.int64()) else: # ARROW_NATIVE tensor_type = pa.fixed_shape_tensor(pa.int64(), (2, 3)) first = ArrowTensorArray.from_numpy(np.ones((2, 2, 3), dtype=np.int64)) second = ArrowTensorArray.from_numpy(np.zeros((3, 2, 3), dtype=np.int64)) assert first.type == second.type # Assert commutation assert tensor_type == first.type assert first.type == tensor_type # Test concatenation works appropriately concatenated = pa.concat_arrays([first, second]) assert len(concatenated) == 5 assert concatenated.type == tensor_type expected = np.vstack([first.to_numpy_ndarray(), second.to_numpy_ndarray()]) np.testing.assert_array_equal(concatenated.to_numpy_ndarray(), expected) @pytest.mark.skipif( not _extension_array_concat_supported(), reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0", ) def test_arrow_variable_shaped_tensor_type_eq_with_concat(): """Test that ArrowVariableShapedTensorType __eq__ method works correctly when concatenating Arrow arrays with variable shaped tensors.""" from ray.data.extensions.tensor_extension import ( ArrowVariableShapedTensorArray, ) # # Case 1: Tensors are variable-shaped but same ``ndim`` # # Create arrays with variable-shaped tensors (but same ndim) first_tensors = [ # (2, 2) np.array([[1, 2], [3, 4]]), # (2, 3) np.array([[5, 6, 7], [8, 9, 10]]), ] second_tensors = [ # (1, 4) np.array([[11, 12, 13, 14]]), # (3, 1) np.array([[15], [16], [17]]), ] first_arr = ArrowVariableShapedTensorArray.from_numpy(first_tensors) second_arr = ArrowVariableShapedTensorArray.from_numpy(second_tensors) # Assert commutation assert first_arr.type == second_arr.type assert second_arr.type == first_arr.type # Assert hashing is correct assert hash(first_arr.type) == hash(second_arr.type) assert first_arr.type.ndim == 2 assert second_arr.type.ndim == 2 # Test concatenation works appropriately concatenated = pa.concat_arrays([first_arr, second_arr]) assert len(concatenated) == 4 assert concatenated.type == first_arr.type result_ndarray = concatenated.to_numpy() for i, expected_ndarray in enumerate( itertools.chain.from_iterable([first_tensors, second_tensors]) ): assert result_ndarray[i].shape == expected_ndarray.shape np.testing.assert_array_equal(result_ndarray[i], expected_ndarray) # # Case 2: Tensors are variable-shaped, with diverging ``ndim``s # # Create arrays with variable-shaped tensors (but different ndim) first_tensors = [ # (1, 2, 1) np.array([[[1], [2]], [[3], [4]]]), # (2, 3, 1) np.array([[[5], [6], [7]], [[8], [9], [10]]]), ] second_tensors = [ # (1, 4) np.array([[11, 12, 13, 14]]), # (3, 1) np.array([[15], [16], [17]]), ] first_arr = ArrowVariableShapedTensorArray.from_numpy(first_tensors) second_arr = ArrowVariableShapedTensorArray.from_numpy(second_tensors) # Assert commutation assert first_arr.type == second_arr.type assert second_arr.type == first_arr.type # Assert hashing is correct assert hash(first_arr.type) == hash(second_arr.type) assert first_arr.type.ndim == 3 assert second_arr.type.ndim == 2 # Test concatenation works appropriately concatenated = pa.concat_arrays([first_arr, second_arr]) assert len(concatenated) == 4 assert concatenated.type == first_arr.type result_ndarray = concatenated.to_numpy() for i, expected_ndarray in enumerate( itertools.chain.from_iterable([first_tensors, second_tensors]) ): assert result_ndarray[i].shape == expected_ndarray.shape np.testing.assert_array_equal(result_ndarray[i], expected_ndarray) def test_reverse_order(): """Test views in reverse order.""" base = np.arange(100, dtype=np.float64) raveled = np.empty(3, dtype=np.object_) raveled[0] = base[50:60].ravel() raveled[1] = base[30:50].ravel() raveled[2] = base[0:30].ravel() # Reverse order views should NOT be contiguous assert not _are_contiguous_1d_views(raveled) def test_concat_ndarrays_zero_copy(): """Test that _concat_ndarrays performs zero-copy concatenation when possible.""" # Case 1: Create a base array and contiguous views base = np.arange(100, dtype=np.int64) arrs = [base[0:20], base[20:50], base[50:100]] result = _concat_ndarrays(arrs) np.testing.assert_array_equal(result, base) # Verify it's a zero-copy view (shares memory with base) assert np.shares_memory(result, base) # Case 2: Verify empty views are skipped arrs = [base[0:10], base[10:10], base[10:20]] # Empty array result = _concat_ndarrays(arrs) expected = np.concatenate([base[0:10], base[10:20]]) np.testing.assert_array_equal(result, expected) # Verify it's a zero-copy view (shares memory with base) assert np.shares_memory(result, base) # Case 3: Singleton ndarray is returned as is result = _concat_ndarrays([base]) # Should return the same array or equivalent assert result is base def test_concat_ndarrays_non_contiguous_fallback(): """Test that _concat_ndarrays falls back to np.concatenate when arrays aren't contiguous.""" # Case 1: Non-contiguous arrays arr1 = np.arange(10, dtype=np.float32) _ = np.arange(1000) # Create gap to prevent contiguity arr2 = np.arange(10, 20, dtype=np.float32) _ = np.arange(1000) # Create gap to prevent contiguity arr3 = np.arange(20, 30, dtype=np.float32) arrs = [arr1, arr2, arr3] result = _concat_ndarrays(arrs) expected = np.concatenate(arrs) np.testing.assert_array_equal(result, expected) assert all(not np.shares_memory(result, a) for a in arrs) # Case 2: Non-contiguous arrays (take 2) base = np.arange(100, dtype=np.float64) arrs = [base[0:10], base[20:30], base[30:40]] # Gap from 10-20 result = _concat_ndarrays(arrs) expected = np.concatenate(arrs) np.testing.assert_array_equal(result, expected) # Should have created a copy since there's a gap assert not np.shares_memory(result, base) def test_concat_ndarrays_diff_dtypes_fallback(): """Different dtypes""" base_int16 = np.arange(50, dtype=np.int16) base_int32 = np.arange(50, dtype=np.int32) # Different dtypes should use fallback arrs = [base_int16, base_int32] # This should use np.concatenate with type promotion result = _concat_ndarrays(arrs) expected = np.concatenate(arrs) np.testing.assert_array_equal(result, expected) assert result.dtype == expected.dtype def test_are_contiguous_1d_views_non_raveled(): """Test that _are_contiguous_1d_views rejects non-1D arrays.""" base = np.arange(100, dtype=np.int64).reshape(10, 10) arrs = [ base[0:2].ravel(), # 1D view base[2:4], # 2D array ] # Should reject because second array is not 1D assert not _are_contiguous_1d_views(arrs) def test_are_contiguous_1d_views_non_c_contiguous(): """Test _are_contiguous_1d_views with non-C-contiguous arrays.""" base = np.arange(100, dtype=np.int64).reshape(10, 10) # Column slices are not C-contiguous arrs = [base[:, 0], base[:, 1]] assert not _are_contiguous_1d_views(arrs) def test_are_contiguous_1d_views_different_bases(): """Test _are_contiguous_1d_views with views from different base arrays.""" base1 = np.arange(50, dtype=np.int64) _ = np.arange(1000, dtype=np.int64) # Create gap to prevent contiguity base2 = np.arange(50, 100, dtype=np.int64) arrs = [base1, base2] # Different base arrays assert not _are_contiguous_1d_views(arrs) def test_are_contiguous_1d_views_overlapping(): """Test _are_contiguous_1d_views with overlapping views.""" base = np.arange(100, dtype=np.float64) arrs = [base[0:20], base[10:30]] # Overlaps with first # Overlapping views are not contiguous assert not _are_contiguous_1d_views(arrs) def test_concat_ndarrays_complex_views(): """Test _concat_ndarrays with complex view scenarios.""" # Create a 2D array and take contiguous row views base_2d = np.arange(100, dtype=np.int64).reshape(10, 10) base = base_2d.ravel() # Get 1D view # Take contiguous slices of the 1D view arrs = [base[0:30], base[30:60], base[60:100]] result = _concat_ndarrays(arrs) np.testing.assert_array_equal(result, base) assert np.shares_memory( result, base_2d ) # Should share memory with original 2D array def test_concat_ndarrays_strided_views(): """Test _concat_ndarrays with strided (non-contiguous) views.""" base = np.arange(100, dtype=np.float64) # Every other element - these are strided views arrs = [base[::2], base[1::2]] # Even indices # Odd indices # Strided views are not C-contiguous result = _concat_ndarrays(arrs) expected = np.concatenate(arrs) np.testing.assert_array_equal(result, expected) # Should have created a copy assert not np.shares_memory(result, base) def test_arrow_extension_serialize_deserialize_cache(): """Test caching behavior of ArrowExtensionSerializeDeserializeCache.""" # Test 1: Serialization cache is instance-level # Create a fresh test instance tensor_type = ArrowTensorType(shape=(2, 3), dtype=pa.int64()) # Clear the instance's serialization cache to ensure fresh test tensor_type._serialize_cache = None # Track calls to _arrow_ext_serialize_compute to verify caching with patch.object( tensor_type, "_arrow_ext_serialize_compute", wraps=tensor_type._arrow_ext_serialize_compute, ) as mock_serialize: # First serialization should call compute function serialized1 = tensor_type.__arrow_ext_serialize__() assert mock_serialize.call_count == 1 assert serialized1 is not None assert isinstance(serialized1, bytes) # Second serialization should use cache (no additional call) serialized2 = tensor_type.__arrow_ext_serialize__() assert mock_serialize.call_count == 1 # Still 1, proving cache hit assert serialized1 == serialized2 # Test 2: Deserialization cache is class-level (shared across instances) # Clear the lru_cache to ensure fresh test ArrowTensorType._arrow_ext_deserialize_cache.cache_clear() storage_type = pa.list_(pa.int64()) # Track calls to _arrow_ext_deserialize_compute to verify caching with patch.object( ArrowTensorType, "_arrow_ext_deserialize_compute", wraps=ArrowTensorType._arrow_ext_deserialize_compute, ) as mock_deserialize: # First deserialization should call compute function deserialized1 = ArrowTensorType.__arrow_ext_deserialize__( storage_type, serialized1 ) assert mock_deserialize.call_count == 1 assert deserialized1.shape == (2, 3) assert deserialized1.value_type == pa.int64() # Second deserialization with same key should use cache (no additional call) deserialized2 = ArrowTensorType.__arrow_ext_deserialize__( storage_type, serialized1 ) assert mock_deserialize.call_count == 1 # Still 1, proving cache hit assert deserialized1.shape == deserialized2.shape assert deserialized1.value_type == deserialized2.value_type assert deserialized1.extension_name == deserialized2.extension_name # Test 3: Different serialized data produces different cache entries tensor_type2 = ArrowTensorType(shape=(3, 4), dtype=pa.int32()) tensor_type2._serialize_cache = None different_serialized = tensor_type2.__arrow_ext_serialize__() storage_type2 = pa.list_(pa.int32()) deserialized3 = ArrowTensorType.__arrow_ext_deserialize__( storage_type2, different_serialized ) # Should be different from previous deserialization assert deserialized3.shape == (3, 4) assert deserialized3.value_type == pa.int32() assert deserialized3.shape != deserialized1.shape # Test 4: Cache parameter generation works correctly param1 = ArrowTensorType._get_deserialize_parameter(storage_type, serialized1) param2 = ArrowTensorType._get_deserialize_parameter(storage_type, serialized1) assert param1 == param2 # Same inputs should produce same parameters param3 = ArrowTensorType._get_deserialize_parameter( storage_type2, different_serialized ) assert param1 != param3 # Different inputs should produce different parameters # Test 5: Multiple instances have separate serialization caches tensor_type_a = ArrowTensorType(shape=(2, 3), dtype=pa.int64()) tensor_type_b = ArrowTensorType(shape=(2, 3), dtype=pa.int64()) # Clear caches tensor_type_a._serialize_cache = None tensor_type_b._serialize_cache = None # Track calls to verify separate caches with patch.object( tensor_type_a, "_arrow_ext_serialize_compute", wraps=tensor_type_a._arrow_ext_serialize_compute, ) as mock_a, patch.object( tensor_type_b, "_arrow_ext_serialize_compute", wraps=tensor_type_b._arrow_ext_serialize_compute, ) as mock_b: # Serialize both instances serialized_a = tensor_type_a.__arrow_ext_serialize__() serialized_b = tensor_type_b.__arrow_ext_serialize__() # Each should have been called once (separate caches) assert mock_a.call_count == 1 assert mock_b.call_count == 1 # Both should produce the same serialized data (same shape and dtype) assert serialized_a == serialized_b # Second calls should use respective caches (no additional calls) assert tensor_type_a.__arrow_ext_serialize__() == serialized_a assert tensor_type_b.__arrow_ext_serialize__() == serialized_b assert mock_a.call_count == 1 # Cache hit assert mock_b.call_count == 1 # Cache hit # Test 6: Deserialization cache is shared (class-level) # The cache is class-level, so all instances share it # Note: deserialized1 and deserialized2 were already created in Test 2, # so the cache should already have this entry. Let's verify it's reused. with patch.object( ArrowTensorType, "_arrow_ext_deserialize_compute", wraps=ArrowTensorType._arrow_ext_deserialize_compute, ) as mock_deserialize_shared: # These should use the cache from Test 2 (no new compute calls) deserialized_a = ArrowTensorType.__arrow_ext_deserialize__( storage_type, serialized1 ) deserialized_b = ArrowTensorType.__arrow_ext_deserialize__( storage_type, serialized1 ) # Should not call compute again (cache hit from Test 2) assert mock_deserialize_shared.call_count == 0 # Both should be equal (cache hit) assert deserialized_a.shape == deserialized_b.shape assert deserialized_a.value_type == deserialized_b.value_type assert deserialized_a.extension_name == deserialized_b.extension_name def test_arrow_extension_deserialize_cache_per_class(): """Test that different classes have separate deserialization caches.""" # Create instances of different classes with the same shape and dtype tensor_type_v1 = ArrowTensorType(shape=(2, 3), dtype=pa.int64()) tensor_type_v2 = ArrowTensorTypeV2(shape=(2, 3), dtype=pa.int64()) # Serialize both (they should produce the same serialized data since shape is the same) serialized_v1 = tensor_type_v1.__arrow_ext_serialize__() serialized_v2 = tensor_type_v2.__arrow_ext_serialize__() # They should have the same serialized representation (same shape) assert serialized_v1 == serialized_v2 # Clear both caches to ensure fresh test ArrowTensorType._arrow_ext_deserialize_cache.cache_clear() ArrowTensorTypeV2._arrow_ext_deserialize_cache.cache_clear() # Get storage types for each class storage_type_v1 = pa.list_(pa.int64()) # ArrowTensorType uses list_ storage_type_v2 = pa.large_list(pa.int64()) # ArrowTensorTypeV2 uses large_list # Track calls to verify each class has its own cache with patch.object( ArrowTensorType, "_arrow_ext_deserialize_compute", wraps=ArrowTensorType._arrow_ext_deserialize_compute, ) as mock_v1, patch.object( ArrowTensorTypeV2, "_arrow_ext_deserialize_compute", wraps=ArrowTensorTypeV2._arrow_ext_deserialize_compute, ) as mock_v2: # Deserialize using ArrowTensorType deserialized_v1_1 = ArrowTensorType.__arrow_ext_deserialize__( storage_type_v1, serialized_v1 ) assert mock_v1.call_count == 1 assert mock_v2.call_count == 0 # V2 cache not affected # Deserialize using ArrowTensorTypeV2 with compatible parameters # Note: We use the same serialized data but different storage type deserialized_v2_1 = ArrowTensorTypeV2.__arrow_ext_deserialize__( storage_type_v2, serialized_v2 ) assert mock_v1.call_count == 1 # V1 cache not affected assert mock_v2.call_count == 1 # Verify they are different instances (different classes) assert type(deserialized_v1_1) is not type(deserialized_v2_1) assert isinstance(deserialized_v1_1, ArrowTensorType) assert isinstance(deserialized_v2_1, ArrowTensorTypeV2) assert not isinstance(deserialized_v1_1, ArrowTensorTypeV2) assert not isinstance(deserialized_v2_1, ArrowTensorType) # Verify they have the same shape and dtype (same logical content) assert deserialized_v1_1.shape == deserialized_v2_1.shape assert deserialized_v1_1.value_type == deserialized_v2_1.value_type # But different extension names (different classes) assert deserialized_v1_1.extension_name != deserialized_v2_1.extension_name assert deserialized_v1_1.extension_name == "ray.data.arrow_tensor" assert deserialized_v2_1.extension_name == "ray.data.arrow_tensor_v2" # Verify each class uses its own cache (second calls should hit cache) deserialized_v1_2 = ArrowTensorType.__arrow_ext_deserialize__( storage_type_v1, serialized_v1 ) deserialized_v2_2 = ArrowTensorTypeV2.__arrow_ext_deserialize__( storage_type_v2, serialized_v2 ) # Both should use cache (no additional compute calls) assert mock_v1.call_count == 1 # Cache hit for V1 assert mock_v2.call_count == 1 # Cache hit for V2 # Verify cache returns same instances for same class assert deserialized_v1_1 is deserialized_v1_2 # Same instance from V1 cache assert deserialized_v2_1 is deserialized_v2_2 # Same instance from V2 cache # But instances from different classes are different assert deserialized_v1_1 is not deserialized_v2_1 assert deserialized_v1_2 is not deserialized_v2_2 def test_arrow_extension_serialize_deserialize_cache_thread_safety(): """Test that ArrowExtensionSerializeDeserializeCache is thread-safe.""" tensor_type = ArrowTensorType(shape=(2, 3), dtype=pa.int64()) storage_type = pa.list_(pa.int64()) serialized = tensor_type.__arrow_ext_serialize__() results = [] errors = [] def deserialize_worker(): try: result = ArrowTensorType.__arrow_ext_deserialize__(storage_type, serialized) results.append(result) except Exception as e: errors.append(e) # Create multiple threads that deserialize simultaneously threads = [threading.Thread(target=deserialize_worker) for _ in range(10)] for thread in threads: thread.start() for thread in threads: thread.join() # Should have no errors assert len(errors) == 0, f"Errors occurred: {errors}" # All results should be equal (same deserialized type) assert len(results) == 10 for result in results[1:]: assert result.shape == results[0].shape assert result.value_type == results[0].value_type if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", "-x", __file__]))