From 3d77723a0c9d245f61e124c58a11e3a1779b3c0d Mon Sep 17 00:00:00 2001
From: wiredfool <eric-github@soroos.net>
Date: Thu, 17 Apr 2025 21:42:42 +0100
Subject: [PATCH] Added arrow support for a flat array of 4*uint8 for image32
 modes

---
 Tests/test_pyarrow.py    | 66 +++++++++++++++++++++++++++++++++++++---
 src/libImaging/Storage.c | 14 +++++++++
 2 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/Tests/test_pyarrow.py b/Tests/test_pyarrow.py
index ece9f8f26..e7f2bc5f9 100644
--- a/Tests/test_pyarrow.py
+++ b/Tests/test_pyarrow.py
@@ -18,18 +18,25 @@ TEST_IMAGE_SIZE = (10, 10)
 
 
 def _test_img_equals_pyarray(
-    img: Image.Image, arr: Any, mask: list[int] | None
+    img: Image.Image, arr: Any, mask: list[int] | None, elts_per_pixel: int = 1
 ) -> None:
-    assert img.height * img.width == len(arr)
+    assert img.height * img.width * elts_per_pixel == len(arr)
     px = img.load()
     assert px is not None
+    if elts_per_pixel > 1 and mask is None:
+        # have to do element wise comparison when we're comparing
+        # flattened r,g,b,a to a pixel.
+        mask = list(range(elts_per_pixel))
     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:
+                pixel = px[x, y]
+                assert isinstance(pixel, tuple)
                 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]
+                    if elts_per_pixel == 1:
+                        assert pixel[ix] == arr[y * img.width + x].as_py()[elt]
+                    else:
+                        assert pixel[ix] == arr[(y * img.width + x) * elts_per_pixel + elt].as_py()
             else:
                 assert_deep_equal(px[x, y], arr[y * img.width + x].as_py())
 
@@ -110,3 +117,52 @@ def test_lifetime2() -> None:
     px = img2.load()
     assert px  # make mypy happy
     assert isinstance(px[0, 0], int)
+
+
+UINT_ARR = (
+    fl_uint8_4_type,
+    [1,2,3,4],
+    1
+)
+UINT = (
+    pyarrow.uint8(),
+    3,
+    4
+)
+
+
+
+@pytest.mark.parametrize(
+    "mode, data_tp, mask",
+    (
+        ("L", (pyarrow.uint8(), 3, 1), None),
+        ("I", (pyarrow.int32(), 1<<24, 1), None),
+        ("F", (pyarrow.float32(), 3.14159, 1), None),
+        ("LA", UINT_ARR, [0, 3]),
+        ("LA", UINT, [0, 3]),
+        ("RGB", UINT_ARR, [0, 1, 2]),
+        ("RGBA", UINT_ARR, None),
+        ("RGBA", UINT_ARR, None),
+        ("CMYK", UINT_ARR, None),
+        ("YCbCr", UINT_ARR,  [0, 1, 2]),
+        ("HSV", UINT_ARR,  [0, 1, 2]),
+        ("RGB", UINT, [0, 1, 2]),
+        ("RGBA", UINT, None),
+        ("RGBA", UINT, None),
+        ("CMYK", UINT, None),
+        ("YCbCr", UINT,  [0, 1, 2]),
+        ("HSV", UINT,  [0, 1, 2]),
+    ),
+)
+def test_fromarray(mode: str,
+                   data_tp: tuple,
+                   mask:list[int] | None) -> None:
+    (dtype,
+     elt,
+     elts_per_pixel) = data_tp
+
+    ct_pixels = TEST_IMAGE_SIZE[0] * TEST_IMAGE_SIZE[1]
+    arr = pyarrow.array([elt]*(ct_pixels*elts_per_pixel), type=dtype)
+    img = Image.fromarrow(arr, mode, TEST_IMAGE_SIZE)
+
+    _test_img_equals_pyarray(img, arr, mask, elts_per_pixel)
diff --git a/src/libImaging/Storage.c b/src/libImaging/Storage.c
index 4fa4ecd1c..7f8d9c4a0 100644
--- a/src/libImaging/Storage.c
+++ b/src/libImaging/Storage.c
@@ -723,6 +723,8 @@ ImagingNewArrow(
     int64_t pixels = (int64_t)xsize * (int64_t)ysize;
 
     // fmt:off   // don't reformat this
+    // stored as a single array, one element per pixel, either single band
+    // or multiband, where each pixel is an I32.
     if (((strcmp(schema->format, "I") == 0  // int32
           && im->pixelsize == 4             // 4xchar* storage
           && im->bands >= 2)                // INT32 into any INT32 Storage mode
@@ -735,6 +737,7 @@ ImagingNewArrow(
             return im;
         }
     }
+    // Stored as [[r,g,b,a],....]
     if (strcmp(schema->format, "+w:4") == 0  // 4 up array
         && im->pixelsize == 4                // storage as 32 bpc
         && schema->n_children > 0            // make sure schema is well formed.
@@ -750,6 +753,17 @@ ImagingNewArrow(
             return im;
         }
     }
+    // Stored as [r,g,b,a,r,g,b,a....]
+    if (strcmp(schema->format, "C") == 0  // uint8
+        && im->pixelsize == 4                // storage as 32 bpc
+        && schema->n_children == 0            // make sure schema is well formed.
+        && strcmp(im->arrow_band_format, "C") == 0        // Expected Format
+        && 4* pixels == external_array->length) {         // expected length
+        // single flat array, interleaved storage.
+        if (ImagingBorrowArrow(im, external_array, 1, array_capsule)) {
+            return im;
+        }
+    }
     // fmt: on
     ImagingDelete(im);
     return NULL;