From 97eb7c09bad7ce19a959156c6d47a9c5ea650085 Mon Sep 17 00:00:00 2001 From: wiredfool Date: Tue, 21 Jan 2025 21:27:41 +0000 Subject: [PATCH] WIP -- First light of round trip of image -> arrow -> image * basic safety included * only respecting the types we emit --- Tests/test_arrow.py | 7 ++- src/PIL/Image.py | 20 ++++++++ src/_imaging.c | 34 ++++++++++++- src/libImaging/Imaging.h | 18 +++++-- src/libImaging/Storage.c | 104 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 178 insertions(+), 5 deletions(-) diff --git a/Tests/test_arrow.py b/Tests/test_arrow.py index 57b22443b12..e8eff880db4 100644 --- a/Tests/test_arrow.py +++ b/Tests/test_arrow.py @@ -6,7 +6,7 @@ from PIL import Image -from .helper import assert_deep_equal, assert_image, hopper, skip_unless_feature +from .helper import assert_deep_equal, assert_image, hopper, skip_unless_feature, assert_image_equal from typing import Any # undone @@ -64,6 +64,11 @@ def test_to_array(mode: str, dtype: Any, mask: Any ) -> None: _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(): # valgrind shouldn't error out here. diff --git a/src/PIL/Image.py b/src/PIL/Image.py index 49bd19ae449..b1dcf0e0fa4 100644 --- a/src/PIL/Image.py +++ b/src/PIL/Image.py @@ -3258,6 +3258,14 @@ class SupportsArrayInterface(Protocol): def __array_interface__(self) -> dict[str, Any]: raise NotImplementedError() +class SupportsArrowArrayInterface(Protocol): + """ + An object that has an ``__arrow_c_array__`` method corresponding to the arrow c data interface. + """ + + def __arrow_c_array__(self, requested_schema:"PyCapsule"=None) -> tuple["PyCapsule", "PyCapsule"]: + raise NotImplementedError() + def fromarray(obj: SupportsArrayInterface, mode: str | None = None) -> Image: """ @@ -3347,6 +3355,18 @@ def fromarray(obj: SupportsArrayInterface, mode: str | None = None) -> Image: return frombuffer(mode, size, obj, "raw", rawmode, 0, 1) +def fromarrow(obj: SupportsArrowArrayIngerface, mode, size) -> ImageFile.ImageFile: + if not hasattr(obj, '__arrow_c_array__'): + raise ValueError("arrow_c_array interface not found") + + (schema_capsule, array_capsule) = obj.__arrow_c_array__() + _im = core.new_arrow(mode, size, schema_capsule, array_capsule) + if (_im): + return Image()._new(_im) + + return None + + def fromqimage(im: ImageQt.QImage) -> ImageFile.ImageFile: """Creates an image instance from a QImage image""" from . import ImageQt diff --git a/src/_imaging.c b/src/_imaging.c index eea719a5f91..c8bc21ac3f3 100644 --- a/src/_imaging.c +++ b/src/_imaging.c @@ -89,7 +89,6 @@ #endif #include "libImaging/Imaging.h" -#include "libImaging/Arrow.h" #define _USE_MATH_DEFINES #include @@ -261,6 +260,38 @@ PyObject* ExportArrowArrayPyCapsule(ImagingObject *self) { } +static PyObject * +_new_arrow(PyObject *self, PyObject *args) { + char *mode; + int xsize, ysize; + PyObject *schema_capsule, *array_capsule; + PyObject *ret; + + if (!PyArg_ParseTuple(args, "s(ii)OO", &mode, &xsize, &ysize, + &schema_capsule, &array_capsule)) { + return NULL; + } + + struct ArrowSchema* schema = + (struct ArrowSchema*)PyCapsule_GetPointer(schema_capsule, "arrow_schema"); + + struct ArrowArray* array = + (struct ArrowArray*)PyCapsule_GetPointer(array_capsule, "arrow_array"); + + ret = PyImagingNew(ImagingNewArrow(mode, xsize, ysize, schema, array)); + if (schema->release){ + schema->release(schema); + schema->release = NULL; + } + if (!ret && array->release) { + array->release(array); + array->release = NULL; + } + return ret; +} + + + /* -------------------------------------------------------------------- */ /* EXCEPTION REROUTING */ /* -------------------------------------------------------------------- */ @@ -4258,6 +4289,7 @@ static PyMethodDef functions[] = { {"fill", (PyCFunction)_fill, METH_VARARGS}, {"new", (PyCFunction)_new, METH_VARARGS}, {"new_block", (PyCFunction)_new_block, METH_VARARGS}, + {"new_arrow", (PyCFunction)_new_arrow, METH_VARARGS}, {"merge", (PyCFunction)_merge, METH_VARARGS}, /* Functions */ diff --git a/src/libImaging/Imaging.h b/src/libImaging/Imaging.h index 7f074292737..cf43873f390 100644 --- a/src/libImaging/Imaging.h +++ b/src/libImaging/Imaging.h @@ -20,6 +20,8 @@ extern "C" { #define M_PI 3.1415926535897932384626433832795 #endif +#include "Arrow.h" + /* -------------------------------------------------------------------- */ /* @@ -107,11 +109,17 @@ struct ImagingMemoryInstance { /* arrow */ int arrow_borrow; /* Number of arrow arrays that have been allocated */ + char band_names[4][3]; /* names of bands, max 2 char + null terminator */ + char arrow_band_format[2]; /* single character + null terminator */ + + int read_only; /* flag for read-only. set for arrow borrowed arrays */ + struct ArrowArray *arrow_array_capsule; /* upstream arrow array source */ + int blocks_count; /* Number of blocks that have been allocated */ int lines_per_block; /* Number of lines in a block have been allocated */ - char band_names[4][3]; /* names of bands, max 2 char + null terminator */ - char arrow_band_format[2]; /* single character + null terminator */ + + }; #define IMAGING_PIXEL_1(im, x, y) ((im)->image8[(y)][(x)]) @@ -195,6 +203,11 @@ ImagingDelete(Imaging im); extern Imaging ImagingNewBlock(const char *mode, int xsize, int ysize); +extern Imaging +ImagingNewArrow(const char *mode, int xsize, int ysize, + struct ArrowSchema *schema, + struct ArrowArray *external_array); + extern Imaging ImagingNewPrologue(const char *mode, int xsize, int ysize); extern Imaging @@ -712,7 +725,6 @@ _imaging_tell_pyFd(PyObject *fd); /* Arrow */ -#include "Arrow.h" extern int export_imaging_array(Imaging im, struct ArrowArray* array); extern int export_imaging_schema(Imaging im, struct ArrowSchema* schema); extern void export_uint32_type(struct ArrowSchema* schema); diff --git a/src/libImaging/Storage.c b/src/libImaging/Storage.c index 5505d44d5e7..f2e03424753 100644 --- a/src/libImaging/Storage.c +++ b/src/libImaging/Storage.c @@ -278,6 +278,7 @@ ImagingNewPrologueSubtype(const char *mode, int xsize, int ysize, int size) { break; } + // UNDONE -- not accurate for arrow MUTEX_LOCK(&ImagingDefaultArena.mutex); ImagingDefaultArena.stats_new_count += 1; MUTEX_UNLOCK(&ImagingDefaultArena.mutex); @@ -556,6 +557,62 @@ ImagingAllocateBlock(Imaging im) { return im; } +/* Borrowed Arrow Storage Type */ +/* --------------------------- */ +/* Don't allocate the image. */ + + +static void +ImagingDestroyArrow(Imaging im) { + if (im->arrow_array_capsule && im->arrow_array_capsule->release) { + im->arrow_array_capsule->release(im->arrow_array_capsule); + im->arrow_array_capsule->release = NULL; + im->arrow_array_capsule = NULL; + } +} + +Imaging +ImagingAllocateArrow(Imaging im, struct ArrowArray *external_array) { + /* don't really allocate, but naming patterns */ + Py_ssize_t y, i; + + char* borrowed_buffer = NULL; + struct ArrowArray* arr = external_array; + + /* overflow check for malloc */ + if (im->linesize && im->ysize > INT_MAX / im->linesize) { + return (Imaging)ImagingError_MemoryError(); + } + + if (arr->n_children == 1) { + arr = arr->children[0]; + } + if (arr->n_buffers == 2) { + // undone offset. offset here is # of elements, + // so depends on the size of the element + // undone image is char*, arrow is void * + // buffer 0 is the null list + // buffer 1 is the data + borrowed_buffer = (char *)arr->buffers[1]; + } + + if (! borrowed_buffer) { + // UNDONE better error here. + // currently, only wanting one where + return (Imaging)ImagingError_MemoryError(); + } + + for (y = i = 0; y < im->ysize; y++) { + im->image[y] = borrowed_buffer + i; + i += im->linesize; + } + im->read_only = 1; + im->arrow_array_capsule = external_array; + im->destroy = ImagingDestroyArrow; + + return im; +} + /* -------------------------------------------------------------------- * Create a new, internally allocated, image. */ @@ -627,6 +684,53 @@ ImagingNewBlock(const char *mode, int xsize, int ysize) { return NULL; } +Imaging +ImagingNewArrow(const char *mode, int xsize, int ysize, + struct ArrowSchema *schema, + struct ArrowArray *external_array) { + /* A borrowed arrow array */ + Imaging im; + + if (xsize < 0 || ysize < 0) { + return (Imaging)ImagingError_ValueError("bad image size"); + } + + im = ImagingNewPrologue(mode, xsize, ysize); + if (!im) { + return NULL; + } + + int64_t pixels = (int64_t)xsize * (int64_t)ysize; + + if (((strcmp(schema->format, "i") == 0 && + im->pixelsize == 4) || + (strcmp(schema->format, im->arrow_band_format) == 0 && + im->bands == 1)) + && pixels == external_array->length) { + // one arrow element per, and it matches a pixelsize*char + if (ImagingAllocateArrow(im, external_array)) { + return im; + } + } + if (strcmp(schema->format, "+w:4") == 0 + && im->pixelsize == 4 + && schema->n_children > 0 + && schema->children + && strcmp(schema->children[0]->format, "C") == 0 + && pixels == external_array->length + && external_array->n_children == 1 + && external_array->children + && 4 * pixels == external_array->children[0]->length) { + // 4 up element of char into pixelsize == 4 + if (ImagingAllocateArrow(im, external_array)) { + return im; + } + } + + ImagingDelete(im); + return NULL; +} + Imaging ImagingNew2Dirty(const char *mode, Imaging imOut, Imaging imIn) { /* allocate or validate output image */