import math import time import numpy as np import pandas as pd import pyarrow as pa import pytest import ray from ray.data._internal.tensor_extensions.arrow import ( MIN_PYARROW_VERSION_FIXED_SHAPE_TENSOR_SCALAR, ArrowTensorArray, FixedShapeTensorFormat, create_arrow_fixed_shape_tensor_type, ) from ray.data._internal.tensor_extensions.utils import _create_possibly_ragged_ndarray from ray.data._internal.utils.arrow_utils import get_pyarrow_version from ray.data.block import BlockAccessor from ray.data.context import DataContext from ray.data.dataset import Schema from ray.data.extensions.tensor_extension import ( ArrowTensorType, ArrowTensorTypeV2, ArrowVariableShapedTensorArray, ArrowVariableShapedTensorType, FixedShapeTensorType, TensorArray, TensorDtype, ) from ray.data.tests.conftest import * # noqa from ray.data.tests.util import extract_values from ray.tests.conftest import * # noqa # https://github.com/ray-project/ray/issues/33695 def test_large_tensor_creation(ray_start_regular_shared, tensor_format_context): """Tests that large tensor read task creation can complete successfully without hanging.""" start_time = time.time() ray.data.range_tensor(1000, override_num_blocks=1000, shape=(80, 80, 100, 100)) end_time = time.time() # Should not take more than 20 seconds. assert end_time - start_time < 20 def test_tensors_basic(ray_start_regular_shared, tensor_format_context): # Determine expected tensor type based on format expected_type = create_arrow_fixed_shape_tensor_type(shape=(3, 5), dtype=pa.int64()) # Create directly. tensor_shape = (3, 5) ds = ray.data.range_tensor(6, shape=tensor_shape, override_num_blocks=6) assert ds.count() == 6 assert ds.schema() == Schema(pa.schema([("data", expected_type)])) # The actual size is slightly larger due to metadata. # We add 6 (one per tensor) offset values of 8 bytes each to account for the # in-memory representation of the PyArrow LargeList type assert math.isclose(ds.size_bytes(), 5 * 3 * 6 * 8 + 6 * 8, rel_tol=0.1) # Test row iterator yields tensors. for tensor in ds.iter_rows(): tensor = tensor["data"] assert isinstance(tensor, np.ndarray) assert tensor.shape == tensor_shape # Test batch iterator yields tensors. for tensor in ds.iter_batches(batch_size=2): tensor = tensor["data"] assert isinstance(tensor, np.ndarray) assert tensor.shape == (2,) + tensor_shape # Native format. def np_mapper(arr): if "data" in arr: arr = arr["data"] else: arr = arr["id"] assert isinstance(arr, np.ndarray) return {"data": arr + 1} res = ray.data.range_tensor(2, shape=(2, 2)).map(np_mapper).take() np.testing.assert_equal( extract_values("data", res), [np.ones((2, 2)), 2 * np.ones((2, 2))] ) # Explicit NumPy format. res = ( ray.data.range_tensor(2, shape=(2, 2)) .map_batches(np_mapper, batch_format="numpy") .take() ) np.testing.assert_equal( extract_values("data", res), [np.ones((2, 2)), 2 * np.ones((2, 2))] ) # Pandas conversion. def pd_mapper(df): assert isinstance(df, pd.DataFrame) return df + 2 res = ray.data.range_tensor(2).map_batches(pd_mapper, batch_format="pandas").take() np.testing.assert_equal(extract_values("data", res), [np.array([2]), np.array([3])]) # Arrow columns in NumPy format. def multi_mapper(col_arrs): assert isinstance(col_arrs, dict) assert list(col_arrs.keys()) == ["a", "b", "c"] assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values()) return {"a": col_arrs["a"] + 1, "b": col_arrs["b"] + 1, "c": col_arrs["c"] + 1} # Multiple columns. t = pa.table( { "a": [1, 2, 3], "b": [4.0, 5.0, 6.0], "c": ArrowTensorArray.from_numpy(np.array([[1, 2], [3, 4], [5, 6]])), } ) res = ( ray.data.from_arrow(t) .map_batches(multi_mapper, batch_size=2, batch_format="numpy") .take() ) np.testing.assert_equal( res, [ {"a": 2, "b": 5.0, "c": np.array([2, 3])}, {"a": 3, "b": 6.0, "c": np.array([4, 5])}, {"a": 4, "b": 7.0, "c": np.array([6, 7])}, ], ) def single_mapper(col_arrs): assert isinstance(col_arrs, dict) assert list(col_arrs.keys()) == ["c"] assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values()) return {"c": col_arrs["c"] + 1} # Single column (should still yield ndarray dict batches). t = t.select(["c"]) res = ( ray.data.from_arrow(t) .map_batches(single_mapper, batch_size=2, batch_format="numpy") .take() ) np.testing.assert_equal( res, [ {"c": np.array([2, 3])}, {"c": np.array([4, 5])}, {"c": np.array([6, 7])}, ], ) # Pandas columns in NumPy format. def multi_mapper(col_arrs): assert isinstance(col_arrs, dict) assert list(col_arrs.keys()) == ["a", "b", "c"] assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values()) return pd.DataFrame( { "a": col_arrs["a"] + 1, "b": col_arrs["b"] + 1, "c": TensorArray(col_arrs["c"] + 1), } ) # Multiple columns. df = pd.DataFrame( { "a": [1, 2, 3], "b": [4.0, 5.0, 6.0], "c": TensorArray(np.array([[1, 2], [3, 4], [5, 6]])), } ) res = ( ray.data.from_pandas(df) .map_batches(multi_mapper, batch_size=2, batch_format="numpy") .take() ) np.testing.assert_equal( res, [ {"a": 2, "b": 5.0, "c": np.array([2, 3])}, {"a": 3, "b": 6.0, "c": np.array([4, 5])}, {"a": 4, "b": 7.0, "c": np.array([6, 7])}, ], ) # Single column (should still yield ndarray dict batches). def single_mapper(col_arrs): assert isinstance(col_arrs, dict) assert list(col_arrs.keys()) == ["c"] assert all(isinstance(col_arr, np.ndarray) for col_arr in col_arrs.values()) return pd.DataFrame({"c": TensorArray(col_arrs["c"] + 1)}) df = df[["c"]] res = ( ray.data.from_pandas(df) .map_batches(single_mapper, batch_size=2, batch_format="numpy") .take() ) np.testing.assert_equal( res, [ {"c": np.array([2, 3])}, {"c": np.array([4, 5])}, {"c": np.array([6, 7])}, ], ) # Simple dataset in NumPy format. def mapper(arr): arr = np_mapper(arr) return arr res = ( ray.data.range(10, override_num_blocks=2) .map_batches(mapper, batch_format="numpy") .take() ) assert extract_values("data", res) == list(range(1, 11)) def test_batch_tensors(ray_start_regular_shared, tensor_format_context): ds = ray.data.from_items( [np.array([0, 0]) for _ in range(40)], override_num_blocks=40 ) batch = next(iter(ds.iter_batches())) assert set(batch) == {"item"} assert batch["item"].shape == (40, 2) def test_tensors_shuffle(ray_start_regular_shared, tensor_format_context): # Test Arrow table representation. tensor_shape = (3, 5) ds = ray.data.range_tensor(6, shape=tensor_shape) shuffled_ds = ds.random_shuffle() shuffled = extract_values("data", shuffled_ds.take()) base = extract_values("data", ds.take()) np.testing.assert_raises( AssertionError, np.testing.assert_equal, shuffled, base, ) np.testing.assert_equal( sorted(shuffled, key=lambda arr: arr.min()), sorted(base, key=lambda arr: arr.min()), ) # Test Pandas table representation. tensor_shape = (3, 5) ds = ray.data.range_tensor(6, shape=tensor_shape) ds = ds.map_batches(lambda df: df, batch_format="pandas") shuffled_ds = ds.random_shuffle() shuffled = extract_values("data", shuffled_ds.take()) base = extract_values("data", ds.take()) np.testing.assert_raises( AssertionError, np.testing.assert_equal, shuffled, base, ) np.testing.assert_equal( sorted(shuffled, key=lambda arr: arr.min()), sorted(base, key=lambda arr: arr.min()), ) def test_tensors_sort(ray_start_regular_shared, tensor_format_context): # Test Arrow table representation. t = pa.table({"a": TensorArray(np.arange(32).reshape((2, 4, 4))), "b": [1, 2]}) ds = ray.data.from_arrow(t) sorted_ds = ds.sort(key="b", descending=True) sorted_arrs = [row["a"] for row in sorted_ds.take()] base = [row["a"] for row in ds.take()] np.testing.assert_raises( AssertionError, np.testing.assert_equal, sorted_arrs, base, ) np.testing.assert_equal( sorted_arrs, sorted(base, key=lambda arr: -arr.min()), ) # Test Pandas table representation. df = pd.DataFrame({"a": TensorArray(np.arange(32).reshape((2, 4, 4))), "b": [1, 2]}) ds = ray.data.from_pandas(df) sorted_ds = ds.sort(key="b", descending=True) sorted_arrs = [np.asarray(row["a"]) for row in sorted_ds.take()] base = [np.asarray(row["a"]) for row in ds.take()] np.testing.assert_raises( AssertionError, np.testing.assert_equal, sorted_arrs, base, ) np.testing.assert_equal( sorted_arrs, sorted(base, key=lambda arr: -arr.min()), ) def test_tensors_inferred_from_map(ray_start_regular_shared, tensor_format_context): tensor_format = tensor_format_context # With tensor_format_context, ARROW_NATIVE only runs when supported, # so to_type() is safe to use expected_type = tensor_format.to_type() # Test map. ds = ray.data.range(10, override_num_blocks=10).map( lambda _: {"data": np.ones((4, 4))} ) ds = ds.materialize() assert ds.count() == 10 schema = ds.schema() assert schema.names == ["data"] dtype = schema.types[0] assert isinstance(dtype, expected_type) assert tuple(dtype.shape) == (4, 4) assert dtype.value_type == pa.float64() # Test map_batches. ds = ray.data.range(16, override_num_blocks=4).map_batches( lambda _: {"data": np.ones((3, 4, 4))}, batch_size=2 ) ds = ds.materialize() assert ds.count() == 24 schema = ds.schema() assert schema.names == ["data"] dtype = schema.types[0] assert isinstance(dtype, expected_type) assert tuple(dtype.shape) == (4, 4) assert dtype.value_type == pa.float64() # Test flat_map. ds = ray.data.range(10, override_num_blocks=10).flat_map( lambda _: [{"data": np.ones((4, 4))}, {"data": np.ones((4, 4))}] ) ds = ds.materialize() assert ds.count() == 20 schema = ds.schema() assert schema.names == ["data"] dtype = schema.types[0] assert isinstance(dtype, expected_type) assert tuple(dtype.shape) == (4, 4) assert dtype.value_type == pa.float64() # Test map_batches ndarray column. ds = ray.data.range(16, override_num_blocks=4).map_batches( lambda _: pd.DataFrame({"a": [np.ones((4, 4))] * 3}), batch_size=2 ) ds = ds.materialize() assert ds.count() == 24 schema = ds.schema() assert schema.names == ["a"] dtype = schema.types[0] assert isinstance(dtype, expected_type) assert tuple(dtype.shape) == (4, 4) assert dtype.value_type == pa.float64() ds = ray.data.range(16, override_num_blocks=4).map_batches( lambda _: pd.DataFrame({"a": [np.ones((2, 2)), np.ones((3, 3))]}), batch_size=2, ) ds = ds.materialize() assert ds.count() == 16 schema = ds.schema() assert schema.names == ["a"] dtype = schema.types[0] assert isinstance(dtype, ArrowVariableShapedTensorType) assert tuple(dtype.shape) == (None, None) assert dtype.value_type == pa.float64() def test_tensor_array_block_slice(tensor_format_context): tensor_format = tensor_format_context # Test that ArrowBlock slicing works with tensor column extension type. def check_for_copy(table1, table2, a, b, is_copy): expected_slice = table1.slice(a, b - a) assert table2.equals(expected_slice) assert table2.schema == table1.schema assert table1.num_columns == table2.num_columns for col1, col2 in zip(table1.columns, table2.columns): assert col1.num_chunks == col2.num_chunks for chunk1, chunk2 in zip(col1.chunks, col2.chunks): bufs1 = chunk1.buffers() bufs2 = chunk2.buffers() expected_offset = 0 if is_copy else a assert chunk2.offset == expected_offset assert len(chunk2) == b - a index = 1 if ( tensor_format == FixedShapeTensorFormat.ARROW_NATIVE and FixedShapeTensorType is not None ): # The buffer for native fixed shaped tensors sits at index 2, not 1 index = 2 if is_copy: assert bufs2[index].address != bufs1[index].address else: assert bufs2[index].address == bufs1[index].address n = 20 one_arr = np.arange(4 * n).reshape(n, 2, 2) df = pd.DataFrame({"one": TensorArray(one_arr), "two": ["a"] * n}) table = pa.Table.from_pandas(df) a, b = 5, 10 block_accessor = BlockAccessor.for_block(table) # Test with copy. table2 = block_accessor.slice(a, b, True) res = table2["one"].chunk(0).to_numpy_ndarray() np.testing.assert_array_equal(res, one_arr[a:b, :, :]) check_for_copy(table, table2, a, b, is_copy=True) # Test without copy. arrow_native requires a copy table2 = block_accessor.slice(a, b, False) res = table2["one"].chunk(0).to_numpy_ndarray() np.testing.assert_array_equal(res, one_arr[a:b, :, :]) check_for_copy(table, table2, a, b, is_copy=False) @pytest.mark.parametrize( "test_data,a,b", [ ([[False, True], [True, False], [True, True], [False, False]], 1, 3), ([[False, True], [True, False], [True, True], [False, False]], 0, 1), ( [ [False, True], [True, False], [True, True], [False, False], [True, False], [False, False], [False, True], [True, True], [False, False], [True, True], [False, True], [True, False], ], 3, 6, ), ( [ [False, True], [True, False], [True, True], [False, False], [True, False], [False, False], [False, True], [True, True], [False, False], [True, True], [False, True], [True, False], ], 7, 11, ), ( [ [False, True], [True, False], [True, True], [False, False], [True, False], [False, False], [False, True], [True, True], [False, False], [True, True], [False, True], [True, False], ], 9, 12, ), # Variable-shaped tensors. ( [[False, True], [True, False, True], [False], [False, False, True, True]], 1, 3, ), ], ) @pytest.mark.parametrize("init_with_pandas", [True, False]) def test_tensor_array_boolean_slice_pandas_roundtrip(init_with_pandas, test_data, a, b): is_variable_shaped = len({len(elem) for elem in test_data}) > 1 n = len(test_data) test_arr = _create_possibly_ragged_ndarray(test_data) df = pd.DataFrame({"one": TensorArray(test_arr), "two": ["a"] * n}) if init_with_pandas: table = pa.Table.from_pandas(df) else: if is_variable_shaped: col = ArrowVariableShapedTensorArray.from_numpy(test_arr) else: col = ArrowTensorArray.from_numpy(test_arr) table = pa.table({"one": col, "two": ["a"] * n}) block_accessor = BlockAccessor.for_block(table) # Test without copy. table2 = block_accessor.slice(a, b, False) out = table2["one"].chunk(0).to_numpy() expected = test_arr[a:b] if is_variable_shaped: for o, e in zip(out, expected): np.testing.assert_array_equal(o, e) else: np.testing.assert_array_equal(out, expected) pd.testing.assert_frame_equal( table2.to_pandas().reset_index(drop=True), df[a:b].reset_index(drop=True) ) # Test with copy. table2 = block_accessor.slice(a, b, True) out = table2["one"].chunk(0).to_numpy() expected = test_arr[a:b] if is_variable_shaped: for o, e in zip(out, expected): np.testing.assert_array_equal(o, e) else: np.testing.assert_array_equal(out, expected) pd.testing.assert_frame_equal( table2.to_pandas().reset_index(drop=True), df[a:b].reset_index(drop=True) ) def test_tensors_in_tables_from_pandas(ray_start_regular_shared, 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": list(arr)}) # Cast column to tensor extension dtype. df["two"] = df["two"].astype(TensorDtype(shape, np.int64)) ds = ray.data.from_pandas([df]) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_from_pandas_variable_shaped( ray_start_regular_shared, tensor_format_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) ] outer_dim = len(arrs) df = pd.DataFrame({"one": list(range(outer_dim)), "two": arrs}) # Cast column to tensor extension dtype. df["two"] = df["two"].astype(TensorDtype(None, np.int64)) ds = ray.data.from_pandas(df) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(range(outer_dim), arrs)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_pandas_roundtrip( ray_start_regular_shared, enable_automatic_tensor_extension_cast, 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)}) ds = ray.data.from_pandas(df) ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas") ds_df = ds.to_pandas() expected_df = df + 1 if enable_automatic_tensor_extension_cast: expected_df["two"] = list(expected_df["two"].to_numpy()) # Roundtrip may use Arrow-backed dtypes (e.g. int64[pyarrow]) for plain columns. expected_df["one"] = expected_df["one"].astype(ds_df["one"].dtype) pd.testing.assert_frame_equal(ds_df, expected_df) def test_tensors_in_tables_pandas_roundtrip_variable_shaped( ray_start_regular_shared, enable_automatic_tensor_extension_cast, tensor_format_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) ] outer_dim = len(arrs) df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arrs)}) ds = ray.data.from_pandas(df) ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas") ds_df = ds.to_pandas() expected_df = df + 1 if enable_automatic_tensor_extension_cast: expected_df["two"] = _create_possibly_ragged_ndarray( expected_df["two"].to_numpy() ) # Roundtrip may use Arrow-backed dtypes (e.g. int64[pyarrow]) for plain columns. expected_df["one"] = expected_df["one"].astype(ds_df["one"].dtype) pd.testing.assert_frame_equal(ds_df, expected_df) def test_tensors_in_tables_parquet_roundtrip( ray_start_regular_shared, tmp_path, 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)}) ds = ray.data.from_pandas(df) ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas") ds.write_parquet(str(tmp_path)) ds = ray.data.read_parquet(str(tmp_path)) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(1, outer_dim + 1)), arr + 1)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_parquet_roundtrip_variable_shaped( ray_start_regular_shared, tmp_path, tensor_format_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) ] outer_dim = len(arrs) df = pd.DataFrame({"one": list(range(outer_dim)), "two": TensorArray(arrs)}) ds = ray.data.from_pandas(df) ds = ds.map_batches(lambda df: df + 1, batch_size=2, batch_format="pandas") ds.write_parquet(str(tmp_path)) ds = ray.data.read_parquet(str(tmp_path)) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(1, outer_dim + 1)), [arr + 1 for arr in arrs])) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_parquet_with_schema( ray_start_regular_shared, tmp_path, 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)}) ds = ray.data.from_pandas([df]) ds.write_parquet(str(tmp_path)) tensor_type = create_arrow_fixed_shape_tensor_type( shape=inner_shape, dtype=pa.from_numpy_dtype(arr.dtype), ) schema = pa.schema([("one", pa.int32()), ("two", tensor_type)]) ds = ray.data.read_parquet(str(tmp_path), schema=schema) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_parquet_pickle_manual_serde( ray_start_regular_shared, tmp_path, tensor_format_context ): import pickle 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": [pickle.dumps(a) for a in arr]} ) ds = ray.data.from_pandas([df]) ds.write_parquet(str(tmp_path)) ds = ray.data.read_parquet(str(tmp_path)) # Manually deserialize the tensor pickle bytes and cast to our tensor # extension type. def deser_mapper(batch: pd.DataFrame): batch["two"] = [pickle.loads(a) for a in batch["two"]] batch["two"] = batch["two"].astype(TensorDtype(shape, np.int64)) return batch casted_ds = ds.map_batches(deser_mapper, batch_format="pandas") values = [[s["one"], s["two"]] for s in casted_ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) # Manually deserialize the pickle tensor bytes and directly cast it to a # TensorArray. def deser_mapper_direct(batch: pd.DataFrame): batch["two"] = TensorArray([pickle.loads(a) for a in batch["two"]]) return batch casted_ds = ds.map_batches(deser_mapper_direct, batch_format="pandas") values = [[s["one"], s["two"]] for s in casted_ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_parquet_bytes_manual_serde( ray_start_regular_shared, tmp_path, 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": [a.tobytes() for a in arr]} ) ds = ray.data.from_pandas([df]) ds.write_parquet(str(tmp_path)) ds = ray.data.read_parquet(str(tmp_path)) tensor_col_name = "two" # Manually deserialize the tensor bytes and cast to a TensorArray. def np_deser_mapper(batch: pa.Table): # NOTE(Clark): We use NumPy to consolidate these potentially # non-contiguous buffers, and to do buffer bookkeeping in general. np_col = np.array( [ np.ndarray(inner_shape, buffer=buf.as_buffer(), dtype=arr.dtype) for buf in batch.column(tensor_col_name) ] ) return batch.set_column( batch._ensure_integer_index(tensor_col_name), tensor_col_name, ArrowTensorArray.from_numpy(np_col), ) ds = ds.map_batches(np_deser_mapper, batch_format="pyarrow") values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_parquet_bytes_manual_serde_udf( ray_start_regular_shared, tmp_path, 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) tensor_col_name = "two" df = pd.DataFrame( {"one": list(range(outer_dim)), tensor_col_name: [a.tobytes() for a in arr]} ) ds = ray.data.from_pandas([df]) ds.write_parquet(str(tmp_path)) ds = ray.data.read_parquet( str(tmp_path), tensor_column_schema={tensor_col_name: (np.dtype(np.int64), inner_shape)}, ) # With tensor_format_context, ARROW_NATIVE only runs when supported, # so to_type() is safe to use without fallback assert isinstance( ds.schema().base_schema.field_by_name(tensor_col_name).type, tensor_format.to_type(), ) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_parquet_bytes_manual_serde_col_schema( ray_start_regular_shared, tmp_path, 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) tensor_col_name = "two" df = pd.DataFrame( {"one": list(range(outer_dim)), tensor_col_name: [a.tobytes() for a in arr]} ) ds = ray.data.from_pandas([df]) ds.write_parquet(str(tmp_path)) def _block_udf(block: pa.Table): df = block.to_pandas() df[tensor_col_name] += 1 return pa.Table.from_pandas(df, schema=block.schema) ds = ray.data.read_parquet( str(tmp_path), tensor_column_schema={tensor_col_name: (arr.dtype, inner_shape)}, _block_udf=_block_udf, ) # With tensor_format_context, ARROW_NATIVE only runs when supported, # so to_type() is safe to use without fallback assert isinstance( ds.schema().base_schema.field_by_name(tensor_col_name).type, tensor_format.to_type(), ) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(outer_dim)), arr + 1)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) @pytest.mark.skip( reason=( "Waiting for Arrow to support registering custom ExtensionType " "casting kernels. See " "https://issues.apache.org/jira/browse/ARROW-5890#" ) ) def test_tensors_in_tables_parquet_bytes_with_schema( ray_start_regular_shared, tmp_path, 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": [a.tobytes() for a in arr]} ) ds = ray.data.from_pandas([df]) ds.write_parquet(str(tmp_path)) tensor_type = create_arrow_fixed_shape_tensor_type( shape=inner_shape, dtype=pa.from_numpy_dtype(arr.dtype), ) schema = pa.schema([("one", pa.int32()), ("two", tensor_type)]) ds = ray.data.read_parquet(str(tmp_path), schema=schema) values = [[s["one"], s["two"]] for s in ds.take()] expected = list(zip(list(range(outer_dim)), arr)) for v, e in zip(sorted(values), expected): np.testing.assert_equal(v, e) def test_tensors_in_tables_iter_batches( ray_start_regular_shared, enable_automatic_tensor_extension_cast, 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) df1 = pd.DataFrame( {"one": TensorArray(arr), "two": TensorArray(arr + 1), "label": [1.0, 2.0, 3.0]} ) arr2 = np.arange(num_items, 2 * num_items).reshape(shape) df2 = pd.DataFrame( { "one": TensorArray(arr2), "two": TensorArray(arr2 + 1), "label": [4.0, 5.0, 6.0], } ) df = pd.concat([df1, df2], ignore_index=True) if enable_automatic_tensor_extension_cast: df["one"] = list(df["one"].to_numpy()) df["two"] = list(df["two"].to_numpy()) ds = ray.data.from_pandas([df1, df2]) batches = list(ds.iter_batches(batch_size=2, batch_format="pandas")) assert len(batches) == 3 expected_batches = [df.iloc[:2], df.iloc[2:4], df.iloc[4:]] for batch, expected_batch in zip(batches, expected_batches): batch = batch.reset_index(drop=True) expected_batch = expected_batch.reset_index(drop=True) pd.testing.assert_frame_equal(batch, expected_batch) def test_ragged_tensors(ray_start_regular_shared, tensor_format_context): """Test Arrow type promotion between ArrowTensorType and ArrowVariableShapedTensorType when a column contains ragged tensors.""" import numpy as np ds = ray.data.from_items( [ {"spam": np.zeros((32, 32, 5))}, {"spam": np.zeros((64, 64, 5))}, ] ) new_type = ds.schema().types[0].value_type assert ds.schema().types == [ ArrowVariableShapedTensorType(dtype=new_type, ndim=3), ] @pytest.mark.parametrize( "write_format", [FixedShapeTensorFormat.V1, FixedShapeTensorFormat.V2], ) @pytest.mark.skipif( get_pyarrow_version() < MIN_PYARROW_VERSION_FIXED_SHAPE_TENSOR_SCALAR, reason="Requires pyarrow>=16 for native FixedShapeTensorType, FixedShapeTensorScalar, FixedShapeTensorArray", ) def test_tensor_format_conversion_v1_v2_to_native( ray_start_regular_shared, tmp_path, restore_data_context, write_format ): """Test that data written in V1/V2 format can be read with native format and written back while preserving types and data. Steps: 1. Write tensor data using V1 or V2 format 2. Set context to use arrow_native format 3. Read the data back (should convert to native format) 4. Write to a different path 5. Read again and verify types/shapes/data are preserved in native format """ ctx = DataContext.get_current() # Step 1: Write data using V1 or V2 format ctx.arrow_fixed_shape_tensor_format = write_format outer_dim = 4 inner_shape = (3, 2) shape = (outer_dim,) + inner_shape num_items = np.prod(np.array(shape)) arr = np.arange(num_items, dtype=np.float64).reshape(shape) # Create dataset with tensor column tensor_col_name = "tensor" df = pd.DataFrame( { "id": list(range(outer_dim)), tensor_col_name: list(arr), } ) ds = ray.data.from_pandas([df]) # Verify initial format matches write_format schema = ds.schema() col_index = schema.names.index(tensor_col_name) initial_type = schema.types[col_index] assert isinstance(initial_type, write_format.to_type()) # Write to first path path1 = tmp_path / "v1_v2_data" ds.write_parquet(str(path1)) # Step 2: Switch context to native format ctx.arrow_fixed_shape_tensor_format = FixedShapeTensorFormat.ARROW_NATIVE # Step 3: Read data back - note that reading preserves original format ds_read = ray.data.read_parquet(str(path1)) # Reading parquet preserves the original extension type (V1 or V2) # The context setting only affects NEW tensor arrays schema_read = ds_read.schema() col_index = schema_read.names.index(tensor_col_name) read_type = schema_read.types[col_index] if write_format == FixedShapeTensorFormat.V1: assert isinstance( read_type, ArrowTensorType ), f"Expected ArrowTensorType when reading V1 data, got {type(read_type).__name__}" else: assert isinstance( read_type, ArrowTensorTypeV2 ), f"Expected ArrowTensorTypeV2 when reading V2 data, got {type(read_type).__name__}" # Step 4: Apply a transformation that recreates tensor arrays in native format # map_batches with identity function will convert to native format def convert_to_native(batch): # This forces recreation of the tensor arrays using the current context format return batch ds_converted = ds_read.map_batches(convert_to_native, batch_format="numpy") # Now verify the format is native schema_converted = ds_converted.schema() col_index = schema_converted.names.index(tensor_col_name) converted_type = schema_converted.types[col_index] assert isinstance(converted_type, FixedShapeTensorType), ( f"Expected FixedShapeTensorType after conversion, " f"got {type(converted_type).__name__}" ) # Verify shape is preserved assert ( tuple(converted_type.shape) == inner_shape ), f"Shape mismatch: expected {inner_shape}, got {converted_type.shape}" # Step 5: Write to different path path2 = tmp_path / "native_data" ds_converted.write_parquet(str(path2)) # Step 6: Read again and verify native format is preserved ds_final = ray.data.read_parquet(str(path2)) schema_final = ds_final.schema() col_index = schema_final.names.index(tensor_col_name) final_type = schema_final.types[col_index] assert isinstance(final_type, FixedShapeTensorType), ( f"Expected FixedShapeTensorType after round-trip, " f"got {type(final_type).__name__}" ) # Verify shape is still correct assert tuple(final_type.shape) == inner_shape, ( f"Shape mismatch after round-trip: expected {inner_shape}, " f"got {final_type.shape}" ) # Verify data is preserved final_values = ds_final.take() for i, row in enumerate(sorted(final_values, key=lambda x: x["id"])): assert row["id"] == i np.testing.assert_array_equal( row[tensor_col_name], arr[i], err_msg=f"Data mismatch at row {i}", ) if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))