diff --git a/Tests/test_arrow.py b/Tests/test_arrow.py index 69193394d..51b680aa1 100644 --- a/Tests/test_arrow.py +++ b/Tests/test_arrow.py @@ -158,3 +158,44 @@ def test_readonly(): assert reloaded.readonly == 1 reloaded._readonly = 0 assert reloaded.readonly == 1 + +def test_multiblock_l_image(): + 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(): + 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(): + 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 = img.__arrow_c_schema__() + +def test_multiblock__rgba_schema(): + 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= img.__arrow_c_schema__() diff --git a/src/_imaging.c b/src/_imaging.c index c7ea0b932..3d2b034b7 100644 --- a/src/_imaging.c +++ b/src/_imaging.c @@ -227,6 +227,20 @@ PyImaging_GetBuffer(PyObject *buffer, Py_buffer *view) { /* Arrow HANDLING */ /* -------------------------------------------------------------------- */ +PyObject * +ArrowError(int err) { + if (err == IMAGING_CODEC_MEMORY) { + return ImagingError_MemoryError(); + } + if (err == IMAGING_ARROW_INCOMPATIBLE_MODE) { + return ImagingError_ValueError("Incompatible Pillow mode for Arrow Array"); + } + if (err == IMAGING_ARROW_MEMORY_LAYOUT) { + return ImagingError_ValueError("Image is in multiple array blocks, use imaging_new_block for zero copy"); + } + return ImagingError_ValueError("Unknown Error"); +} + void ReleaseArrowSchemaPyCapsule(PyObject *capsule) { struct ArrowSchema *schema = @@ -241,8 +255,12 @@ PyObject * ExportArrowSchemaPyCapsule(ImagingObject *self) { struct ArrowSchema *schema = (struct ArrowSchema *)calloc(1, sizeof(struct ArrowSchema)); - export_imaging_schema(self->image, schema); - return PyCapsule_New(schema, "arrow_schema", ReleaseArrowSchemaPyCapsule); + int err = export_imaging_schema(self->image, schema); + if (err == 0) { + return PyCapsule_New(schema, "arrow_schema", ReleaseArrowSchemaPyCapsule); + } + free(schema); + return ArrowError(err); } void @@ -259,10 +277,16 @@ PyObject * ExportArrowArrayPyCapsule(ImagingObject *self) { struct ArrowArray *array = (struct ArrowArray *)calloc(1, sizeof(struct ArrowArray)); - export_imaging_array(self->image, array); - return PyCapsule_New(array, "arrow_array", ReleaseArrowArrayPyCapsule); + int err = export_imaging_array(self->image, array); + if (err == 0) { + return PyCapsule_New(array, "arrow_array", ReleaseArrowArrayPyCapsule); + } + free(array); + // raise error here + return ArrowError(err); } + static PyObject * _new_arrow(PyObject *self, PyObject *args) { char *mode; diff --git a/src/libImaging/Arrow.c b/src/libImaging/Arrow.c index 7dcd04d5c..dd6c8fc38 100644 --- a/src/libImaging/Arrow.c +++ b/src/libImaging/Arrow.c @@ -24,21 +24,26 @@ ReleaseExportedSchema(struct ArrowSchema *array) { free((void *)array->name); array->name = NULL; } + if (array->metadata) { + free((void *)array->metadata); + array->metadata = NULL; + } // Release children for (int64_t i = 0; i < array->n_children; ++i) { struct ArrowSchema *child = array->children[i]; if (child->release != NULL) { child->release(child); - // assert(child->release == NULL); + child->release = NULL; } + // UNDONE -- should I be releasing the children? } // Release dictionary struct ArrowSchema *dict = array->dictionary; if (dict != NULL && dict->release != NULL) { dict->release(dict); - // assert(dict->release == NULL); + dict->release = NULL; } // TODO here: release and/or deallocate all data directly owned by @@ -58,28 +63,28 @@ export_named_type(struct ArrowSchema *schema, char *format, char *name) { formatp = calloc(format_len, 1); if (!formatp) { - return 1; + return IMAGING_CODEC_MEMORY; } namep = calloc(name_len, 1); if (!namep) { free(formatp); - return 1; + return IMAGING_CODEC_MEMORY; } strncpy(formatp, format, format_len); strncpy(namep, name, name_len); *schema = (struct ArrowSchema){// Type description - .format = formatp, - .name = namep, - .metadata = NULL, - .flags = 0, - .n_children = 0, - .children = NULL, - .dictionary = NULL, - // Bookkeeping - .release = &ReleaseExportedSchema + .format = formatp, + .name = namep, + .metadata = NULL, + .flags = 0, + .n_children = 0, + .children = NULL, + .dictionary = NULL, + // Bookkeeping + .release = &ReleaseExportedSchema }; return 0; } @@ -89,7 +94,12 @@ export_imaging_schema(Imaging im, struct ArrowSchema *schema) { int retval = 0; if (strcmp(im->arrow_band_format, "") == 0) { - return 1; + return IMAGING_ARROW_INCOMPATIBLE_MODE; + } + + /* for now, single block images */ + if (!(im->blocks_count == 0 || im->blocks_count == 1)){ + return IMAGING_ARROW_MEMORY_LAYOUT; } if (im->bands == 1) { @@ -97,17 +107,18 @@ export_imaging_schema(Imaging im, struct ArrowSchema *schema) { } retval = export_named_type(schema, "+w:4", ""); - if (retval) { + if (retval != 0) { return retval; } // if it's not 1 band, it's an int32 at the moment. 4 unint8 bands. schema->n_children = 1; schema->children = calloc(1, sizeof(struct ArrowSchema *)); schema->children[0] = (struct ArrowSchema *)calloc(1, sizeof(struct ArrowSchema)); - if (export_named_type(schema->children[0], im->arrow_band_format, "pixel")) { + retval = export_named_type(schema->children[0], im->arrow_band_format, "pixel"); + if (retval != 0) { free(schema->children[0]); schema->release(schema); - return 2; + return retval; } return 0; } @@ -118,11 +129,27 @@ static void release_const_array(struct ArrowArray *array) { Imaging im = (Imaging)array->private_data; - ImagingDelete(im); + if (array->n_children == 0) { + ImagingDelete(im); + } - // assert(array->n_buffers == 2); // Free the buffers and the buffers array - free(array->buffers); + if (array->buffers) { + free(array->buffers); + array->buffers = NULL; + } + if (array->children) { + // undone -- does arrow release all the children recursively. + for (int i=0; i < array->n_children; i++) { + if (array->children[i]->release){ + array->children[i]->release(array->children[i]); + array->children[i]->release = NULL; + free(array->children[i]); + } + } + free(array->children); + array->children = NULL; + } // Mark released array->release = NULL; } @@ -131,8 +158,10 @@ int export_single_channel_array(Imaging im, struct ArrowArray *array) { int length = im->xsize * im->ysize; - /* undone -- for now, single block images */ - // assert (im->block_count = 0 || im->block_count = 1); + /* for now, single block images */ + if (!(im->blocks_count == 0 || im->blocks_count == 1)){ + return IMAGING_ARROW_MEMORY_LAYOUT; + } if (im->lines_per_block && im->lines_per_block < im->ysize) { length = im->xsize * im->lines_per_block; @@ -172,8 +201,10 @@ int export_fixed_pixel_array(Imaging im, struct ArrowArray *array) { int length = im->xsize * im->ysize; - /* undone -- for now, single block images */ - // assert (im->block_count = 0 || im->block_count = 1); + /* for now, single block images */ + if (!(im->blocks_count == 0 || im->blocks_count == 1)) { + return IMAGING_ARROW_MEMORY_LAYOUT; + } if (im->lines_per_block && im->lines_per_block < im->ysize) { length = im->xsize * im->lines_per_block; @@ -200,13 +231,22 @@ export_fixed_pixel_array(Imaging im, struct ArrowArray *array) { // Allocate list of buffers array->buffers = (const void **)calloc(1, sizeof(void *) * array->n_buffers); + if (!array->buffers) { + goto err; + } // assert(array->buffers != NULL); array->buffers[0] = NULL; // no nulls, null bitmap can be omitted // if it's not 1 band, it's an int32 at the moment. 4 unint8 bands. array->n_children = 1; array->children = calloc(1, sizeof(struct ArrowArray *)); + if (!array->children) { + goto err; + } array->children[0] = (struct ArrowArray *)calloc(1, sizeof(struct ArrowArray)); + if (!array->children[0]) { + goto err; + } MUTEX_LOCK(im->mutex); im->refcount++; @@ -233,12 +273,24 @@ export_fixed_pixel_array(Imaging im, struct ArrowArray *array) { array->children[0]->buffers[1] = im->blocks[0].ptr; } return 0; + + err: + if (array->children[0]) { + free(array->children[0]); + } + if (array->children) { + free(array->children); + } + if (array->buffers) { + free(array->buffers); + } + return IMAGING_CODEC_MEMORY; } int export_imaging_array(Imaging im, struct ArrowArray *array) { if (strcmp(im->arrow_band_format, "") == 0) { - return 1; + return IMAGING_ARROW_INCOMPATIBLE_MODE; } if (im->bands == 1) { diff --git a/src/libImaging/Imaging.h b/src/libImaging/Imaging.h index 750e8a71e..7136e7ff9 100644 --- a/src/libImaging/Imaging.h +++ b/src/libImaging/Imaging.h @@ -744,6 +744,9 @@ export_imaging_schema(Imaging im, struct ArrowSchema *schema); #define IMAGING_CODEC_UNKNOWN -3 #define IMAGING_CODEC_CONFIG -8 #define IMAGING_CODEC_MEMORY -9 +#define IMAGING_ARROW_INCOMPATIBLE_MODE -10 +#define IMAGING_ARROW_MEMORY_LAYOUT -11 + #include "ImagingUtils.h" extern UINT8 *clip8_lookups; diff --git a/src/libImaging/Storage.c b/src/libImaging/Storage.c index 9e6439b88..fd75a5584 100644 --- a/src/libImaging/Storage.c +++ b/src/libImaging/Storage.c @@ -703,7 +703,7 @@ ImagingNewArrow( if (((strcmp(schema->format, "I") == 0 // int32 && im->pixelsize == 4 // 4xchar* storage && im->bands >= 2) // INT32 into any INT32 Storage mode - || + || // (()||()) && (strcmp(schema->format, im->arrow_band_format) == 0 // same mode && im->bands == 1)) // Single band match && pixels == external_array->length) {