Skip to content

Commit

Permalink
Error handling:
Browse files Browse the repository at this point in the history
* Return memory error for allocation errors
* Return value error for invalid layout (block) or bad mode)
* Free children on releasing the array
* Only decrement refcount on the leaf array release
  • Loading branch information
wiredfool committed Feb 3, 2025
1 parent 4fc8328 commit e7bd152
Show file tree
Hide file tree
Showing 5 changed files with 150 additions and 30 deletions.
41 changes: 41 additions & 0 deletions Tests/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
32 changes: 28 additions & 4 deletions src/_imaging.c
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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
Expand All @@ -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;
Expand Down
102 changes: 77 additions & 25 deletions src/libImaging/Arrow.c
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
}
Expand All @@ -89,25 +94,31 @@ 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) {
return export_named_type(schema, im->arrow_band_format, im->band_names[0]);
}

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;
}
Expand All @@ -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;
}
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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++;
Expand All @@ -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) {
Expand Down
3 changes: 3 additions & 0 deletions src/libImaging/Imaging.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion src/libImaging/Storage.c
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down

0 comments on commit e7bd152

Please sign in to comment.