/python/ray/data/tests/test_tensor_extension.py

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__]))