Skip to content

Commit

Permalink
WIP -- First light of round trip of image -> arrow -> image
Browse files Browse the repository at this point in the history
* basic safety included
* only respecting the types we emit
  • Loading branch information
wiredfool committed Jan 21, 2025
1 parent e1ef083 commit 97eb7c0
Show file tree
Hide file tree
Showing 5 changed files with 178 additions and 5 deletions.
7 changes: 6 additions & 1 deletion Tests/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down
20 changes: 20 additions & 0 deletions src/PIL/Image.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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
Expand Down
34 changes: 33 additions & 1 deletion src/_imaging.c
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,6 @@
#endif

#include "libImaging/Imaging.h"
#include "libImaging/Arrow.h"

#define _USE_MATH_DEFINES
#include <math.h>
Expand Down Expand Up @@ -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 */
/* -------------------------------------------------------------------- */
Expand Down Expand Up @@ -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 */
Expand Down
18 changes: 15 additions & 3 deletions src/libImaging/Imaging.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ extern "C" {
#define M_PI 3.1415926535897932384626433832795
#endif

#include "Arrow.h"

/* -------------------------------------------------------------------- */

/*
Expand Down Expand Up @@ -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)])
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand Down
104 changes: 104 additions & 0 deletions src/libImaging/Storage.c
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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 */
Expand Down

0 comments on commit 97eb7c0

Please sign in to comment.