Image -> Arrow support (#8330)

Co-authored-by: Andrew Murray <[email protected]>
Co-authored-by: Hugo van Kemenade <[email protected]>

Author: 3 people

.ci/install.sh

@@ -36,6 +36,9 @@ python3 -m pip install -U pytest
 python3 -m pip install -U pytest-cov
 python3 -m pip install -U pytest-timeout
 python3 -m pip install pyroma
+# optional test dependency, only install if there's a binary package.
+# fails on beta 3.14 and PyPy
+python3 -m pip install --only-binary=:all: pyarrow || true
 
 if [[ $(uname) != CYGWIN* ]]; then
     python3 -m pip install numpy

.github/workflows/macos-install.sh

@@ -30,6 +30,9 @@ python3 -m pip install -U pytest-cov
 python3 -m pip install -U pytest-timeout
 python3 -m pip install pyroma
 python3 -m pip install numpy
+# optional test dependency, only install if there's a binary package.
+# fails on beta 3.14 and PyPy
+python3 -m pip install --only-binary=:all: pyarrow || true
 
 # libavif
 pushd depends && ./install_libavif.sh && popd

.github/workflows/test-windows.yml

@@ -88,6 +88,10 @@ jobs:
       run: |
         python3 -m pip install PyQt6
 
+    - name: Install PyArrow dependency
+      run: |
+        python3 -m pip install --only-binary=:all: pyarrow || true
+
     - name: Install dependencies
       id: install
       run: |

Tests/test_arrow.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import pytest
+
+from PIL import Image
+
+from .helper import hopper
+
+
+@pytest.mark.parametrize(
+    "mode, dest_modes",
+    (
+        ("L", ["I", "F", "LA", "RGB", "RGBA", "RGBX", "CMYK", "YCbCr", "HSV"]),
+        ("I", ["L", "F"]),  # Technically I;32 can work for any 4x8bit storage.
+        ("F", ["I", "L", "LA", "RGB", "RGBA", "RGBX", "CMYK", "YCbCr", "HSV"]),
+        ("LA", ["L", "F"]),
+        ("RGB", ["L", "F"]),
+        ("RGBA", ["L", "F"]),
+        ("RGBX", ["L", "F"]),
+        ("CMYK", ["L", "F"]),
+        ("YCbCr", ["L", "F"]),
+        ("HSV", ["L", "F"]),
+    ),
+)
+def test_invalid_array_type(mode: str, dest_modes: list[str]) -> None:
+    img = hopper(mode)
+    for dest_mode in dest_modes:
+        with pytest.raises(ValueError):
+            Image.fromarrow(img, dest_mode, img.size)
+
+
+def test_invalid_array_size() -> None:
+    img = hopper("RGB")
+
+    assert img.size != (10, 10)
+    with pytest.raises(ValueError):
+        Image.fromarrow(img, "RGB", (10, 10))
+
+
+def test_release_schema() -> None:
+    # these should not error out, valgrind should be clean
+    img = hopper("L")
+    schema = img.__arrow_c_schema__()
+    del schema
+
+
+def test_release_array() -> None:
+    # these should not error out, valgrind should be clean
+    img = hopper("L")
+    array, schema = img.__arrow_c_array__()
+    del array
+    del schema
+
+
+def test_readonly() -> None:
+    img = hopper("L")
+    reloaded = Image.fromarrow(img, img.mode, img.size)
+    assert reloaded.readonly == 1
+    reloaded._readonly = 0
+    assert reloaded.readonly == 1
+
+
+def test_multiblock_l_image() -> None:
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in single channel mode
+    size = (4096, 2 * block_size // 4096)
+    img = Image.new("L", size, 128)
+
+    with pytest.raises(ValueError):
+        (schema, arr) = img.__arrow_c_array__()
+
+
+def test_multiblock_rgba_image() -> None:
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in 4 channel mode
+    size = (4096, (block_size // 4096) // 2)
+    img = Image.new("RGBA", size, (128, 127, 126, 125))
+
+    with pytest.raises(ValueError):
+        (schema, arr) = img.__arrow_c_array__()
+
+
+def test_multiblock_l_schema() -> None:
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in single channel mode
+    size = (4096, 2 * block_size // 4096)
+    img = Image.new("L", size, 128)
+
+    with pytest.raises(ValueError):
+        img.__arrow_c_schema__()
+
+
+def test_multiblock_rgba_schema() -> None:
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in 4 channel mode
+    size = (4096, (block_size // 4096) // 2)
+    img = Image.new("RGBA", size, (128, 127, 126, 125))
+
+    with pytest.raises(ValueError):
+        img.__arrow_c_schema__()
+
+
+def test_singleblock_l_image() -> None:
+    Image.core.set_use_block_allocator(1)
+
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in 4 channel mode
+    size = (4096, 2 * (block_size // 4096))
+    img = Image.new("L", size, 128)
+    assert img.im.isblock()
+
+    (schema, arr) = img.__arrow_c_array__()
+    assert schema
+    assert arr
+
+    Image.core.set_use_block_allocator(0)
+
+
+def test_singleblock_rgba_image() -> None:
+    Image.core.set_use_block_allocator(1)
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in 4 channel mode
+    size = (4096, (block_size // 4096) // 2)
+    img = Image.new("RGBA", size, (128, 127, 126, 125))
+    assert img.im.isblock()
+
+    (schema, arr) = img.__arrow_c_array__()
+    assert schema
+    assert arr
+    Image.core.set_use_block_allocator(0)
+
+
+def test_singleblock_l_schema() -> None:
+    Image.core.set_use_block_allocator(1)
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in single channel mode
+    size = (4096, 2 * block_size // 4096)
+    img = Image.new("L", size, 128)
+    assert img.im.isblock()
+
+    schema = img.__arrow_c_schema__()
+    assert schema
+    Image.core.set_use_block_allocator(0)
+
+
+def test_singleblock_rgba_schema() -> None:
+    Image.core.set_use_block_allocator(1)
+    block_size = Image.core.get_block_size()
+
+    # check a 2 block image in 4 channel mode
+    size = (4096, (block_size // 4096) // 2)
+    img = Image.new("RGBA", size, (128, 127, 126, 125))
+    assert img.im.isblock()
+
+    schema = img.__arrow_c_schema__()
+    assert schema
+    Image.core.set_use_block_allocator(0)

Tests/test_pyarrow.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from typing import Any  # undone
+
+import pytest
+
+from PIL import Image
+
+from .helper import (
+    assert_deep_equal,
+    assert_image_equal,
+    hopper,
+)
+
+pyarrow = pytest.importorskip("pyarrow", reason="PyArrow not installed")
+
+TEST_IMAGE_SIZE = (10, 10)
+
+
+def _test_img_equals_pyarray(
+    img: Image.Image, arr: Any, mask: list[int] | None
+) -> None:
+    assert img.height * img.width == len(arr)
+    px = img.load()
+    assert px is not None
+    for x in range(0, img.size[0], int(img.size[0] / 10)):
+        for y in range(0, img.size[1], int(img.size[1] / 10)):
+            if mask:
+                for ix, elt in enumerate(mask):
+                    pixel = px[x, y]
+                    assert isinstance(pixel, tuple)
+                    assert pixel[ix] == arr[y * img.width + x].as_py()[elt]
+            else:
+                assert_deep_equal(px[x, y], arr[y * img.width + x].as_py())
+
+
+# really hard to get a non-nullable list type
+fl_uint8_4_type = pyarrow.field(
+    "_", pyarrow.list_(pyarrow.field("_", pyarrow.uint8()).with_nullable(False), 4)
+).type
+
+
+@pytest.mark.parametrize(
+    "mode, dtype, mask",
+    (
+        ("L", pyarrow.uint8(), None),
+        ("I", pyarrow.int32(), None),
+        ("F", pyarrow.float32(), None),
+        ("LA", fl_uint8_4_type, [0, 3]),
+        ("RGB", fl_uint8_4_type, [0, 1, 2]),
+        ("RGBA", fl_uint8_4_type, None),
+        ("RGBX", fl_uint8_4_type, None),
+        ("CMYK", fl_uint8_4_type, None),
+        ("YCbCr", fl_uint8_4_type, [0, 1, 2]),
+        ("HSV", fl_uint8_4_type, [0, 1, 2]),
+    ),
+)
+def test_to_array(mode: str, dtype: Any, mask: list[int] | None) -> None:
+    img = hopper(mode)
+
+    # Resize to non-square
+    img = img.crop((3, 0, 124, 127))
+    assert img.size == (121, 127)
+
+    arr = pyarrow.array(img)
+    _test_img_equals_pyarray(img, arr, mask)
+    assert arr.type == dtype
+
+    reloaded = Image.fromarrow(arr, mode, img.size)
+
+    assert reloaded
+
+    assert_image_equal(img, reloaded)
+
+
+def test_lifetime() -> None:
+    # valgrind shouldn't error out here.
+    # arrays should be accessible after the image is deleted.
+
+    img = hopper("L")
+
+    arr_1 = pyarrow.array(img)
+    arr_2 = pyarrow.array(img)
+
+    del img
+
+    assert arr_1.sum().as_py() > 0
+    del arr_1
+
+    assert arr_2.sum().as_py() > 0
+    del arr_2
+
+
+def test_lifetime2() -> None:
+    # valgrind shouldn't error out here.
+    # img should remain after the arrays are collected.
+
+    img = hopper("L")
+
+    arr_1 = pyarrow.array(img)
+    arr_2 = pyarrow.array(img)
+
+    assert arr_1.sum().as_py() > 0
+    del arr_1
+
+    assert arr_2.sum().as_py() > 0
+    del arr_2
+
+    img2 = img.copy()
+    px = img2.load()
+    assert px  # make mypy happy
+    assert isinstance(px[0, 0], int)

docs/reference/Image.rst

@@ -79,6 +79,7 @@ Constructing images
 
 .. autofunction:: new
 .. autofunction:: fromarray
+.. autofunction:: fromarrow
 .. autofunction:: frombytes
 .. autofunction:: frombuffer
 
@@ -370,6 +371,8 @@ Protocols
 
 .. autoclass:: SupportsArrayInterface
     :show-inheritance:
+.. autoclass:: SupportsArrowArrayInterface
+    :show-inheritance:
 .. autoclass:: SupportsGetData
     :show-inheritance: