ray-project
diff --git a/‎ci/lint/pydoclint-baseline.txt‎
Lines changed: 0 additions & 1 deletion b/‎ci/lint/pydoclint-baseline.txt‎
Lines changed: 0 additions & 1 deletion
diff --git a/‎doc/source/conf.py‎
Lines changed: 2 additions & 0 deletions b/‎doc/source/conf.py‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎python/ray/_private/arrow_utils.py‎
Lines changed: 43 additions & 0 deletions b/‎python/ray/_private/arrow_utils.py‎
Lines changed: 43 additions & 0 deletions
diff --git a/‎python/ray/air/tests/test_tensor_extension.py‎
Lines changed: 295 additions & 3 deletions b/‎python/ray/air/tests/test_tensor_extension.py‎
Lines changed: 295 additions & 3 deletions
@@ -415,7 +415,6 @@ python/ray/air/util/tensor_extensions/arrow.py
     DOC101: Function `pyarrow_table_from_pydict`: Docstring contains fewer arguments than in function signature.
     DOC103: Function `pyarrow_table_from_pydict`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [pydict: Dict[str, Union[List[Any], pa.Array]]].
     DOC201: Function `pyarrow_table_from_pydict` does not have a return section in docstring
-    DOC201: Method `ArrowTensorArray._concat_same_type` does not have a return section in docstring
 --------------------
 python/ray/air/util/tensor_extensions/pandas.py
     DOC101: Method `TensorDtype.__init__`: Docstring contains fewer arguments than in function signature.
 
@@ -672,6 +672,8 @@ def filter(self, record):
     "ray._raylet",
     "ray.core.generated",
     "ray.serve.generated",
+    "ray.air.util.tensor_extensions",
+    "ray.data._internal.arrow_ops",
 ]
 
 for mock_target in autodoc_mock_imports:
 
@@ -1,13 +1,56 @@
 import json
+import logging
+import os
 from typing import Dict, Optional
 from urllib.parse import parse_qsl, unquote, urlencode, urlparse, urlunparse
 
 from packaging.version import Version, parse as parse_version
 
+_RAY_DISABLE_PYARROW_VERSION_CHECK = "RAY_DISABLE_PYARROW_VERSION_CHECK"
+
+
 _PYARROW_INSTALLED: Optional[bool] = None
 _PYARROW_VERSION: Optional[Version] = None
 
 
+# NOTE: Make sure that these lower and upper bounds stay in sync with version
+# constraints given in python/setup.py.
+# Inclusive minimum pyarrow version.
+_PYARROW_SUPPORTED_VERSION_MIN = "9.0.0"
+_PYARROW_VERSION_VALIDATED = False
+
+
+logger = logging.getLogger(__name__)
+
+
+def _check_pyarrow_version():
+    """Checks that Pyarrow's version is within the supported bounds."""
+    global _PYARROW_VERSION_VALIDATED
+
+    if not _PYARROW_VERSION_VALIDATED:
+        if os.environ.get(_RAY_DISABLE_PYARROW_VERSION_CHECK, "0") == "1":
+            _PYARROW_VERSION_VALIDATED = True
+            return
+
+        version = get_pyarrow_version()
+        if version is not None:
+            if version < parse_version(_PYARROW_SUPPORTED_VERSION_MIN):
+                raise ImportError(
+                    f"Dataset requires pyarrow >= {_PYARROW_SUPPORTED_VERSION_MIN}, but "
+                    f"{version} is installed. Reinstall with "
+                    f'`pip install -U "pyarrow"`. '
+                )
+        else:
+            logger.warning(
+                "You are using the 'pyarrow' module, but the exact version is unknown "
+                "(possibly carried as an internal component by another module). Please "
+                f"make sure you are using pyarrow >= {_PYARROW_SUPPORTED_VERSION_MIN} to ensure "
+                "compatibility with Ray Dataset. "
+            )
+
+        _PYARROW_VERSION_VALIDATED = True
+
+
 def get_pyarrow_version() -> Optional[Version]:
     """Get the version of the pyarrow package or None if not installed."""
     global _PYARROW_INSTALLED, _PYARROW_VERSION
 
@@ -14,6 +14,11 @@
     ArrowTensorTypeV2,
     ArrowVariableShapedTensorArray,
     ArrowVariableShapedTensorType,
+    _are_contiguous_1d_views,
+    _concat_ndarrays,
+    _extension_array_concat_supported,
+    concat_tensor_arrays,
+    unify_tensor_arrays,
 )
 from ray.air.util.tensor_extensions.pandas import TensorArray, TensorDtype
 from ray.air.util.tensor_extensions.utils import create_ragged_ndarray
@@ -715,7 +720,7 @@ def test_arrow_tensor_array_concat(a1, a2, restore_data_context, tensor_format):
 
     ta1 = ArrowTensorArray.from_numpy(a1)
     ta2 = ArrowTensorArray.from_numpy(a2)
-    ta = ArrowTensorArray._concat_same_type([ta1, ta2])
+    ta = concat_tensor_arrays([ta1, ta2])
     assert len(ta) == a1.shape[0] + a2.shape[0]
     if a1.shape[1:] == a2.shape[1:]:
         if tensor_format == "v1":
@@ -753,8 +758,8 @@ def test_variable_shaped_tensor_array_chunked_concat(
     a2 = np.arange(np.prod(shape2)).reshape(shape2)
     ta1 = ArrowTensorArray.from_numpy(a1)
     ta2 = ArrowTensorArray.from_numpy(a2)
-    chunked_ta = ArrowTensorArray._chunk_tensor_arrays([ta1, ta2])
-    ta = ArrowTensorArray._concat_same_type(chunked_ta.chunks)
+    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)
@@ -828,6 +833,293 @@ def test_tensor_array_string_tensors_simple(restore_data_context, tensor_format)
     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
+
+
+@pytest.mark.skipif(
+    not _extension_array_concat_supported(),
+    reason="ExtensionArrays support concatenation only in Pyarrow >= 12.0",
+)
+def test_arrow_fixed_shape_tensor_type_eq_with_concat(restore_data_context):
+    """Test that ArrowTensorType and ArrowTensorTypeV2 __eq__ methods work correctly
+    when concatenating Arrow arrays with the same tensor type."""
+    from ray.data.context import DataContext
+    from ray.data.extensions.tensor_extension import (
+        ArrowTensorArray,
+        ArrowTensorType,
+        ArrowTensorTypeV2,
+    )
+
+    # Test ArrowTensorType V1
+    tensor_type_v1 = ArrowTensorType((2, 3), pa.int64())
+
+    DataContext.get_current().use_arrow_tensor_v2 = False
+    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_v1 == first.type
+    assert first.type == tensor_type_v1
+
+    # Test concatenation works appropriately
+    concatenated = pa.concat_arrays([first, second])
+    assert len(concatenated) == 5
+    assert concatenated.type == tensor_type_v1
+
+    expected = np.vstack([first.to_numpy(), second.to_numpy()])
+    np.testing.assert_array_equal(concatenated.to_numpy(), expected)
+
+    # Test ArrowTensorTypeV2
+    tensor_type_v2 = ArrowTensorTypeV2((2, 3), pa.int64())
+
+    DataContext.get_current().use_arrow_tensor_v2 = True
+
+    first = ArrowTensorArray.from_numpy(np.ones((2, 2, 3), dtype=np.int64))
+    second = ArrowTensorArray.from_numpy(np.ones((3, 2, 3), dtype=np.int64))
+
+    assert first.type == second.type
+    # Assert commutation
+    assert tensor_type_v2 == first.type
+    assert first.type == tensor_type_v2
+
+    # Test concatenation works appropriately
+    concatenated_v2 = pa.concat_arrays([first, second])
+    assert len(concatenated_v2) == 5
+    assert concatenated_v2.type == tensor_type_v2
+
+    # Assert on the full concatenated array
+    expected = np.vstack([first.to_numpy(), second.to_numpy()])
+    np.testing.assert_array_equal(concatenated_v2.to_numpy(), 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,
+        ArrowVariableShapedTensorType,
+    )
+
+    # Create ArrowVariableShapedTensorType
+    var_tensor_type = ArrowVariableShapedTensorType(pa.int64(), 2)
+
+    # Create arrays with variable-shaped tensors
+    tensors1 = [np.array([[1, 2], [3, 4]]), np.array([[5, 6, 7], [8, 9, 10]])]
+    tensors2 = [np.array([[11, 12, 13, 14]]), np.array([[15], [16], [17]])]
+
+    arr1 = ArrowVariableShapedTensorArray.from_numpy(tensors1)
+    arr2 = ArrowVariableShapedTensorArray.from_numpy(tensors2)
+
+    assert arr1.type == arr2.type
+    # Assert commutation
+    assert var_tensor_type == arr1.type
+    assert arr1.type == var_tensor_type
+
+    # Test concatenation works appropriately
+    concatenated = pa.concat_arrays([arr1, arr2])
+    assert len(concatenated) == 4
+    assert concatenated.type == var_tensor_type
+
+    result = concatenated.to_numpy()
+    expected_shapes = [(2, 2), (2, 3), (1, 4), (3, 1)]
+    for i, expected_shape in enumerate(expected_shapes):
+        assert result[i].shape == expected_shape
+
+
+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)
+
+
 if __name__ == "__main__":
     import sys
Original file line number	Diff line number	Diff line change
`@@ -672,6 +672,8 @@ def filter(self, record):`
`672`	`672`	`"ray._raylet",`
`673`	`673`	`"ray.core.generated",`
`674`	`674`	`"ray.serve.generated",`
	`675`	`+ "ray.air.util.tensor_extensions",`
	`676`	`+ "ray.data._internal.arrow_ops",`
`675`	`677`	`]`
`676`	`678`
`677`	`679`	`for mock_target in autodoc_mock_imports:`