Function test_tensors_basic

python/ray/data/tests/test_tensor.py:47–224

(ray_start_regular_shared, tensor_format_context)

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