/*

* Copyright ironArray SL 2021.

*

* All rights reserved.

*

* This software is the confidential and proprietary information of ironArray SL

* ("Confidential Information"). You shall not disclose such Confidential

* Information and shall use it only in accordance with the terms of the license agreement.

*

*/

#include "iarray_private.h"

#include <libiarray/iarray.h>

#include <minjugg.h>

#include <caterva.h>

#if defined(_OPENMP)

#include <omp.h>

#endif

typedef enum iarray_expr_input_class_e {

IARRAY_EXPR_EQ = 0u, // Same chunkshape/blockshape

IARRAY_EXPR_EQ_NCOMP = 1u, // Same chunkshape/blockshape and no-compressed data

IARRAY_EXPR_NEQ = 2u // Different chunkshape/blockshape

} iarray_expr_input_class_t;

// Struct to be used as info container for dealing with the expression

typedef struct iarray_expr_pparams_s {

bool compressed_inputs;

int ninputs; // number of data inputs

iarray_expr_input_class_t input_class[IARRAY_EXPR_OPERANDS_MAX]; // whether the inputs are compressed or not

uint8_t* inputs[IARRAY_EXPR_OPERANDS_MAX]; // the data inputs

int32_t input_csizes[IARRAY_EXPR_OPERANDS_MAX]; // the compressed size of data inputs

int32_t input_typesizes[IARRAY_EXPR_OPERANDS_MAX]; // the typesizes for data inputs

iarray_expression_t *e;

iarray_iter_write_block_value_t out_value;

} iarray_expr_pparams_t;

// Struct to be used as argument to the evaluation function

typedef struct iarray_eval_pparams_s {

int ninputs; // number of data inputs

uint8_t* inputs[IARRAY_EXPR_OPERANDS_MAX]; // the data inputs

int32_t input_typesizes[IARRAY_EXPR_OPERANDS_MAX]; // the typesizes for data inputs

void *user_data; // a pointer to an iarray_expr_pparams_t struct

uint8_t *out; // the output buffer

int32_t out_size; // the size of output buffer (in bytes)

int32_t out_typesize; // the typesize of output

int8_t ndim; // the number of dimensions for inputs / output arrays

int32_t *window_shape; // the shape of the window for the input arrays (NULL if not available)

int64_t *window_start; // the start coordinates for the window shape (NULL if not available)

int32_t *window_strides; // the strides for the window shape (NULL if not available)

iarray_user_param_t user_params[IARRAY_EXPR_USER_PARAMS_MAX]; // the input user parameters

} iarray_eval_pparams_t;

typedef int (*iarray_eval_fn)(iarray_eval_pparams_t *params);

INA_API(ina_rc_t) iarray_expr_new(iarray_context_t *ctx, iarray_data_type_t data_type, iarray_expression_t **e) {

INA_VERIFY_NOT_NULL(ctx);

INA_VERIFY_NOT_NULL(e);

*e = ina_mem_alloc(sizeof(iarray_expression_t));

INA_RETURN_IF_NULL(e);

(*e)->ctx = ctx;

(*e)->ctx->expr_vars = NULL;

(*e)->expr = NULL;

(*e)->nvars = 0;

(*e)->max_out_len = 0; // helper for leftovers

(*e)->nuser_params = 0;

ina_mem_set(&(*e)->vars, 0, sizeof(_iarray_jug_var_t) * IARRAY_EXPR_OPERANDS_MAX);

// map dtype to JUG type

jug_expression_dtype_t dtype;

switch (data_type) {

case IARRAY_DATA_TYPE_BOOL:

// how to support?

dtype = JUG_EXPRESSION_DTYPE_SINT8;// is that accurate?

break;

case IARRAY_DATA_TYPE_DOUBLE:

dtype = JUG_EXPRESSION_DTYPE_DOUBLE;

break;

case IARRAY_DATA_TYPE_FLOAT:

dtype = JUG_EXPRESSION_DTYPE_FLOAT;

break;

// case IARRAY_DATA_TYPE_FLOAT16:

// // not supported yet

// return INA_ERROR(INA_ERR_INVALID_ARGUMENT);

// case IARRAY_DATA_TYPE_FLOAT8:

// // not supported yet

// return INA_ERROR(INA_ERR_INVALID_ARGUMENT);

case IARRAY_DATA_TYPE_INT8:

dtype = JUG_EXPRESSION_DTYPE_SINT8;

break;

case IARRAY_DATA_TYPE_INT16:

dtype = JUG_EXPRESSION_DTYPE_SINT16;

break;

case IARRAY_DATA_TYPE_INT32:

dtype = JUG_EXPRESSION_DTYPE_SINT32;

break;

case IARRAY_DATA_TYPE_INT64:

dtype = JUG_EXPRESSION_DTYPE_SINT64;

break;

default:

return INA_ERROR(INA_ERR_INVALID_ARGUMENT);

}

jug_expression_new(&(*e)->jug_expr, dtype);

return INA_SUCCESS;

}

INA_API(void) iarray_expr_free(iarray_context_t *ctx, iarray_expression_t **e)

{

INA_VERIFY_FREE(e);

if ((*e)->jug_expr != NULL) {

jug_expression_free(&(*e)->jug_expr);

}

for (int nvar=0; nvar < (*e)->nvars; nvar++) {

free((void*)((*e)->vars[nvar].var));

}

INA_MEM_FREE(ctx->expr_vars);

ina_str_free((*e)->expr);

INA_MEM_FREE_SAFE(*e);

}

INA_API(ina_rc_t) iarray_expr_bind(iarray_expression_t *e, const char *var, iarray_container_t *val)

{

INA_VERIFY_NOT_NULL(e);

INA_VERIFY_NOT_NULL(var);

INA_VERIFY_NOT_NULL(val);

e->vars[e->nvars].var = strdup(var); // yes, we want a copy here!

e->vars[e->nvars].c = val;

e->nvars++;

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_expr_bind_param(iarray_expression_t *e, iarray_user_param_t val)

{

INA_VERIFY_NOT_NULL(e);

if (e->nuser_params >= IARRAY_EXPR_USER_PARAMS_MAX) {

return INA_ERROR(INA_ERR_FULL);

}

e->user_params[e->nuser_params] = val;

e->nuser_params++;

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_expr_bind_out_properties(iarray_expression_t *e, iarray_dtshape_t *dtshape, iarray_storage_t *store)

{

e->out_dtshape = ina_mem_alloc(sizeof(iarray_dtshape_t));

ina_mem_cpy(e->out_dtshape, dtshape, sizeof(iarray_dtshape_t));

e->out_store_properties = ina_mem_alloc(sizeof(iarray_storage_t));

ina_mem_cpy(e->out_store_properties, store, sizeof(iarray_storage_t));

if (store->urlpath != NULL) {

e->out_store_properties->urlpath = strdup(store->urlpath);

} else {

e->out_store_properties->urlpath = NULL;

}

return INA_SUCCESS;

}

int caterva_blosc_array_repart_chunk(int8_t *rchunk, int64_t rchunksize, void *chunk,

int64_t chunksize, caterva_array_t *array) {

if (rchunksize != array->extchunknitems * array->itemsize) {

CATERVA_ERROR(CATERVA_ERR_INVALID_ARGUMENT);

}

if (chunksize != array->chunknitems * array->itemsize) {

CATERVA_ERROR(CATERVA_ERR_INVALID_ARGUMENT);

}

const int8_t *src_b = (int8_t *) chunk;

memset(rchunk, 0, (size_t) rchunksize);

int32_t d_pshape[CATERVA_MAX_DIM];

int64_t d_epshape[CATERVA_MAX_DIM];

int32_t d_spshape[CATERVA_MAX_DIM];

uint8_t d_ndim = array->ndim;

for (int i = 0; i < CATERVA_MAX_DIM; ++i) {

d_pshape[(CATERVA_MAX_DIM - d_ndim + i) % CATERVA_MAX_DIM] = array->chunkshape[i];

d_epshape[(CATERVA_MAX_DIM - d_ndim + i) % CATERVA_MAX_DIM] = array->extchunkshape[i];

d_spshape[(CATERVA_MAX_DIM - d_ndim + i) % CATERVA_MAX_DIM] = array->blockshape[i];

}

int64_t aux[CATERVA_MAX_DIM];

aux[7] = d_epshape[7] / d_spshape[7];

for (int i = CATERVA_MAX_DIM - 2; i >= 0; i--) {

aux[i] = d_epshape[i] / d_spshape[i] * aux[i + 1];

}

/* Fill each block buffer */

int32_t orig[CATERVA_MAX_DIM];

int64_t actual_spsize[CATERVA_MAX_DIM];

for (int32_t sci = 0; sci < array->extchunknitems / array->blocknitems; sci++) {

/*Calculate the coord. of the block first element */

orig[7] = sci % ((int32_t)d_epshape[7] / d_spshape[7]) * d_spshape[7];

for (int i = CATERVA_MAX_DIM - 2; i >= 0; i--) {

orig[i] = (int32_t)(sci % (aux[i]) / (aux[i + 1]) * d_spshape[i]);

}

/* Calculate if padding with 0s is needed for this block */

for (int i = CATERVA_MAX_DIM - 1; i >= 0; i--) {

if (orig[i] + d_spshape[i] > d_pshape[i]) {

actual_spsize[i] = (d_pshape[i] - orig[i]);

} else {

actual_spsize[i] = d_spshape[i];

}

}

int32_t seq_copylen = (int32_t) (actual_spsize[7] * array->itemsize);

/* Reorder each line of data from src_b to chunk */

int64_t ii[CATERVA_MAX_DIM];

int64_t ncopies = 1;

for (int i = 0; i < CATERVA_MAX_DIM - 1; ++i) {

ncopies *= actual_spsize[i];

}

for (int ncopy = 0; ncopy < ncopies; ++ncopy) {

iarray_index_unidim_to_multidim_shape(CATERVA_MAX_DIM - 1, actual_spsize, ncopy, ii);

int64_t d_a = d_spshape[7];

int64_t d_coord_f = sci * array->blocknitems;

for (int i = CATERVA_MAX_DIM - 2; i >= 0; i--) {

d_coord_f += ii[i] * d_a;

d_a *= d_spshape[i];

}

int64_t s_coord_f = orig[7];

int64_t s_a = d_pshape[7];

for (int i = CATERVA_MAX_DIM - 2; i >= 0; i--) {

s_coord_f += (orig[i] + ii[i]) * s_a;

s_a *= d_pshape[i];

}

memcpy(rchunk + d_coord_f * array->itemsize, src_b + s_coord_f * array->itemsize,

seq_copylen);

}

}

return CATERVA_SUCCEED;

}

static ina_rc_t _iarray_expr_prepare(iarray_expression_t *e)

{

uint32_t eval_method = e->ctx->cfg->eval_method & 0x3u;

if (eval_method == IARRAY_EVAL_METHOD_AUTO) {

eval_method = IARRAY_EVAL_METHOD_ITERBLOSC;

}

e->ctx->cfg->eval_method = eval_method;

switch (e->out_dtshape->dtype) {

case IARRAY_DATA_TYPE_DOUBLE:

e->typesize = sizeof(double);

break;

case IARRAY_DATA_TYPE_FLOAT:

e->typesize = sizeof(float);

break;

case IARRAY_DATA_TYPE_BOOL:

case IARRAY_DATA_TYPE_INT8:

case IARRAY_DATA_TYPE_UINT8:

e->typesize = sizeof(int8_t);

break;

case IARRAY_DATA_TYPE_INT16:

case IARRAY_DATA_TYPE_UINT16:

e->typesize = sizeof(int16_t);

break;

case IARRAY_DATA_TYPE_INT32:

case IARRAY_DATA_TYPE_UINT32:

e->typesize = sizeof(int32_t);

break;

case IARRAY_DATA_TYPE_INT64:

case IARRAY_DATA_TYPE_UINT64:

e->typesize = sizeof(int64_t);

break;

default:

return INA_ERROR(IARRAY_ERR_INVALID_DTYPE);

}

size_t size;

IARRAY_RETURN_IF_FAILED(iarray_shape_size(e->out_dtshape, &size));

e->nbytes = (int64_t) size * e->typesize;

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_expr_compile_udf(

iarray_expression_t *e,

int llvm_bc_len,

const char *llvm_bc,

const char* name)

{

INA_VERIFY_NOT_NULL(e);

INA_VERIFY_NOT_NULL(llvm_bc);

IARRAY_RETURN_IF_FAILED(_iarray_expr_prepare(e));

IARRAY_RETURN_IF_FAILED(

jug_udf_compile(e->jug_expr, llvm_bc_len, llvm_bc, name, &e->jug_expr_func)

);

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_expr_compile(iarray_expression_t *e, const char *expr)

{

INA_VERIFY_NOT_NULL(e);

INA_VERIFY_NOT_NULL(expr);

e->expr = ina_str_new_fromcstr(expr);

IARRAY_RETURN_IF_FAILED(_iarray_expr_prepare(e));

jug_te_variable *jug_vars = ina_mem_alloc(e->nvars * sizeof(jug_te_variable));

memset(jug_vars, 0, e->nvars * sizeof(jug_te_variable));

for (int nvar = 0; nvar < e->nvars; nvar++) {

jug_vars[nvar].name = e->vars[nvar].var;

}

IARRAY_RETURN_IF_FAILED(jug_expression_compile(e->jug_expr, ina_str_cstr(e->expr), e->nvars,

jug_vars, &e->jug_expr_func));

return INA_SUCCESS;

}

int prefilter_func(blosc2_prefilter_params *pparams)

{

iarray_expr_pparams_t *expr_pparams = (iarray_expr_pparams_t*)pparams->user_data;

struct iarray_expression_s *e = expr_pparams->e;

int ninputs = expr_pparams->ninputs;

// Populate the eval_pparams

iarray_eval_pparams_t eval_pparams = {0};

eval_pparams.ninputs = ninputs;

memcpy(eval_pparams.input_typesizes, expr_pparams->input_typesizes, ninputs * sizeof(int32_t));

eval_pparams.user_data = expr_pparams;

eval_pparams.out = pparams->out;

eval_pparams.out_size = pparams->out_size;

eval_pparams.out_typesize = pparams->out_typesize;

eval_pparams.ndim = expr_pparams->e->out_dtshape->ndim;

int32_t blocksize = pparams->out_size;

int32_t typesize = pparams->out_typesize;

int8_t ndim = e->out->dtshape->ndim;

// Element strides (in elements)

int32_t strides[IARRAY_DIMENSION_MAX];

strides[ndim - 1] = 1;

for (int i = ndim - 2; i >= 0 ; --i) {

strides[i] = strides[i+1] * e->out->catarr->blockshape[i+1];

}

// Block strides (in blocks)

int32_t strides_block[IARRAY_DIMENSION_MAX];

strides_block[ndim - 1] = 1;

for (int i = ndim - 2; i >= 0 ; --i) {

strides_block[i] = strides_block[i+1] * (int32_t) (e->out->catarr->extchunkshape[i+1] / e->out->catarr->blockshape[i+1]);

}

// Flattened block number

int32_t nblock = pparams->out_offset / pparams->out_size;

// Multidimensional block number

int32_t nblock_ndim[IARRAY_DIMENSION_MAX];

for (int i = ndim - 1; i >= 0; --i) {

if (i != 0) {

nblock_ndim[i] = (nblock % strides_block[i-1]) / strides_block[i];

} else {

nblock_ndim[i] = (nblock % ((int32_t)e->out->catarr->extchunknitems / e->out->catarr->blocknitems)) / strides_block[i];

}

}

// Position of the first element of the block (inside current chunk)

int64_t start_in_chunk[IARRAY_DIMENSION_MAX];

for (int i = 0; i < ndim; ++i) {

start_in_chunk[i] = nblock_ndim[i] * e->out->catarr->blockshape[i];

}

// Position of the first element of the block (inside container)

int64_t start_in_container[IARRAY_DIMENSION_MAX];

for (int i = 0; i < ndim; ++i) {

start_in_container[i] = start_in_chunk[i] + expr_pparams->out_value.block_index[i] * e->out->catarr->chunkshape[i];

}

// Check if the block is out of bounds

bool out_of_bounds = false;

for (int i = 0; i < ndim; ++i) {

if (start_in_container[i] > e->out->catarr->shape[i]) {

out_of_bounds = true;

break;

}

}

// Shape of the current block

int32_t shape[IARRAY_DIMENSION_MAX];

for (int i = 0; i < ndim; ++i) {

if (out_of_bounds) {

shape[i] = 0;

} else if (start_in_container[i] + e->out->catarr->blockshape[i] > e->out->catarr->shape[i]) {

shape[i] = (int32_t) (e->out->catarr->shape[i] - start_in_container[i]);

} else if (start_in_chunk[i] + e->out->catarr->blockshape[i] > e->out->catarr->chunkshape[i]) {

shape[i] = (int32_t) (e->out->catarr->chunkshape[i] - start_in_chunk[i]);

} else {

shape[i] = e->out->catarr->blockshape[i];

}

}

// We can only set the visible shape of the output for the ITERBLOSC eval method.

eval_pparams.window_shape = shape;

eval_pparams.window_start = start_in_container;

eval_pparams.window_strides = strides;

// The code below only works for the case where inputs and output have the same typesize.

// More love is needed in the future, where we would want to allow mixed types in expressions.

bool inputs_malloced[IARRAY_DIMENSION_MAX];

for (int i = 0; i < ninputs; i++) {

inputs_malloced[i] = false;

switch (expr_pparams->input_class[i]) {

case IARRAY_EXPR_EQ_NCOMP:

eval_pparams.inputs[i] = expr_pparams->inputs[i] + BLOSC_EXTENDED_HEADER_LENGTH + pparams->out_offset;

break;

case IARRAY_EXPR_EQ:

eval_pparams.inputs[i] = ina_mem_alloc_aligned(64, blocksize);

inputs_malloced[i] = true;

int64_t nitems = blocksize / typesize;

int64_t offset_index = pparams->out_offset / typesize;

blosc2_dparams dparams = {.nthreads = 1,

.schunk = e->vars[i].c->catarr->sc,

.postfilter = e->vars[i].c->catarr->sc->dctx->postfilter,

.postparams = e->vars[i].c->catarr->sc->dctx->postparams,

};

blosc2_context *dctx = blosc2_create_dctx(dparams);

int64_t rbytes = blosc2_getitem_ctx(dctx, expr_pparams->inputs[i], expr_pparams->input_csizes[i],

(int) offset_index, (int) nitems,

eval_pparams.inputs[i], blocksize);

blosc2_free_ctx(dctx);

if (rbytes != blocksize) {

fprintf(stderr, "Read from inputs failed inside pipeline\n");

return -1;

}

break;

case IARRAY_EXPR_NEQ:

eval_pparams.inputs[i] = expr_pparams->inputs[i] + pparams->out_offset;

break;

default:

return -1;

}

}

for (unsigned int i = 0; i < e->nuser_params; i++) {

eval_pparams.user_params[i] = e->user_params[i];

}

// Eval the expression for this chunk

int ret;

ret = ((iarray_eval_fn) e->jug_expr_func)(&eval_pparams);

switch (ret) {

case 0:

// 0 means success

break;

case 1:

IARRAY_TRACE1(iarray.error, "Out of bounds in LLVM eval engine");

return -2;

default:

IARRAY_TRACE1(iarray.error, "Error in executing LLVM eval engine");

return -3;

}

for (int i = 0; i < ninputs; i++) {

if (inputs_malloced[i]) {

INA_MEM_FREE_SAFE(eval_pparams.inputs[i]);

}

}

return 0;

}

ina_rc_t iarray_eval_cleanup(iarray_expression_t *e, int64_t nitems_written)

{

int64_t nitems_in_schunk = e->nbytes / e->typesize;

if (nitems_written != nitems_in_schunk) {

IARRAY_TRACE1(iarray.error, "The number of items written is different from items in final container");

return INA_ERROR(INA_ERR_NOT_COMPLETE);

}

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_eval_iterchunk(iarray_expression_t *e, iarray_container_t *ret, int64_t *out_chunkshape)

{

int nvars = e->nvars;

int64_t nitems_written = 0;

// Create and initialize an iterator per variable

iarray_config_t cfg = IARRAY_CONFIG_DEFAULTS;

iarray_context_t *ctx = NULL;

iarray_context_new(&cfg, &ctx);

iarray_iter_read_block_t **iter_var = ina_mem_alloc(nvars * sizeof(iarray_iter_read_block_t));

iarray_iter_read_block_value_t *iter_value = ina_mem_alloc(nvars * sizeof(iarray_iter_read_block_value_t));

for (int nvar = 0; nvar < nvars; nvar++) {

iarray_container_t *var = e->vars[nvar].c;

IARRAY_RETURN_IF_FAILED(iarray_iter_read_block_new(ctx, &iter_var[nvar], var, out_chunkshape, &iter_value[nvar],

false));

}

// Write iterator for output

iarray_iter_write_block_t *iter_out;

iarray_iter_write_block_value_t out_value;

IARRAY_RETURN_IF_FAILED(iarray_iter_write_block_new(ctx, &iter_out, ret, out_chunkshape, &out_value, false));

// Create expr pparams

iarray_expr_pparams_t expr_pparams = {0};

expr_pparams.e = e;

expr_pparams.ninputs = nvars;

// Create eval pparams

iarray_eval_pparams_t eval_pparams = {0};

eval_pparams.ninputs = nvars;

eval_pparams.out_typesize = (int32_t) e->out->catarr->itemsize;

eval_pparams.ndim = e->out->dtshape->ndim;

eval_pparams.user_data = &expr_pparams;

for (int i = 0; i < nvars; ++i) {

eval_pparams.input_typesizes[i] = (int32_t) e->vars[i].c->catarr->itemsize;

expr_pparams.input_typesizes[i] = (int32_t) e->vars[i].c->catarr->itemsize;

expr_pparams.input_class[i] = IARRAY_EXPR_NEQ;

}

// Evaluate the expression for all the chunks in variables

while (INA_SUCCEED(iarray_iter_write_block_has_next(iter_out))) {

IARRAY_RETURN_IF_FAILED(iarray_iter_write_block_next(iter_out, NULL, 0));

int32_t out_items = (int32_t)(iter_out->cur_block_size);

// Decompress chunks in variables into temporaries

for (int nvar = 0; nvar < nvars; nvar++) {

IARRAY_RETURN_IF_FAILED(iarray_iter_read_block_next(iter_var[nvar], NULL, 0));

eval_pparams.inputs[nvar] = iter_value[nvar].block_pointer;

expr_pparams.inputs[nvar] = iter_value[nvar].block_pointer;

}

// Eval the expression for this chunk

e->max_out_len = out_items; // so as to prevent operating beyond the limits

eval_pparams.out = out_value.block_pointer;

eval_pparams.out_size = (int32_t)out_value.block_size * e->typesize;

expr_pparams.out_value = out_value;

int32_t shape[IARRAY_DIMENSION_MAX];

int64_t start[IARRAY_DIMENSION_MAX];

int32_t strides[IARRAY_DIMENSION_MAX];

strides[ret->dtshape->ndim - 1] = 1;

for (int i = ret->dtshape->ndim - 1; i >= 0; --i) {

shape[i] = (int32_t) out_value.block_shape[i];

start[i] = out_value.elem_index[i];

if (i != ret->dtshape->ndim - 1)

strides[i] = strides[i+1] * shape[i+1];

}

eval_pparams.window_shape = shape;

eval_pparams.window_start = start;

eval_pparams.window_strides = strides;

int err = ((iarray_eval_fn) e->jug_expr_func)(&eval_pparams);

if (err != 0) {

return INA_ERROR(IARRAY_ERR_EVAL_ENGINE_FAILED);

}

nitems_written += out_items;

}

IARRAY_ITER_FINISH();

for (int nvar = 0; nvar < nvars; nvar++) {

iarray_iter_read_block_free(&(iter_var[nvar]));

}

iarray_iter_write_block_free(&iter_out);

INA_MEM_FREE_SAFE(iter_var);

INA_MEM_FREE_SAFE(iter_value);

iarray_context_free(&ctx);

IARRAY_RETURN_IF_FAILED(iarray_eval_cleanup(e, nitems_written));

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_eval_iterblosc(iarray_expression_t *e, iarray_container_t *ret, int64_t *out_chunkshape)

{

int nvars = e->nvars;

// Setup a new cparams with a prefilter

blosc2_prefilter_params pparams = {0};

iarray_expr_pparams_t expr_pparams = {0};

expr_pparams.e = e;

expr_pparams.ninputs = nvars;

pparams.user_data = (void *) &expr_pparams;

// Initialize the typesize for each variable

for (int nvar = 0; nvar < nvars; nvar++) {

iarray_container_t *var = e->vars[nvar].c;

expr_pparams.input_typesizes[nvar] = (int32_t) var->catarr->itemsize;

}

// Determine the class of each container

iarray_container_t *out = e->out;

char *name = "zproxy_urlpath";

bool is_zproxy;

bool iterblosc_allowed;

for (int nvar = 0; nvar < nvars; ++nvar) {

iterblosc_allowed = true;

iarray_container_t *var = e->vars[nvar].c;

IARRAY_RETURN_IF_FAILED(iarray_vlmeta_exists(e->ctx, e->vars[nvar].c, name, &is_zproxy));

if (is_zproxy || var->transposed) {

// If it is a zproxy or transposed, iterblosc cannot be used

iterblosc_allowed = false;

}

else {

if (var->container_viewed != NULL) {

// If shape is not the same we cannot use iterblosc

// See https://github.com/inaos/iron-array/issues/581

for (int i = 0; i < var->container_viewed->dtshape->ndim; ++i) {

if (var->dtshape->shape[i] != var->container_viewed->dtshape->shape[i]) {

iterblosc_allowed = false;

break;

}

}

}

}

if (iterblosc_allowed) {

for (int i = 0; i < var->dtshape->ndim; ++i) {

if (out->storage->chunkshape[i] != var->storage->chunkshape[i]) {

iterblosc_allowed = false;

break;

}

if (out->storage->blockshape[i] != var->storage->blockshape[i]) {

iterblosc_allowed = false;

break;

}

}

}

if (iterblosc_allowed == false) {

expr_pparams.input_class[nvar] = IARRAY_EXPR_NEQ;

}

else {

expr_pparams.input_class[nvar] = IARRAY_EXPR_EQ;

}

}

iarray_context_t *ctx = e->ctx;

// Need for not compatible containers

iarray_iter_read_block_t **iter_var = ina_mem_alloc(nvars * sizeof(iarray_iter_read_block_t));

iarray_iter_read_block_value_t *iter_value = ina_mem_alloc(nvars * sizeof(iarray_iter_read_block_value_t));

uint8_t **external_buffers = ina_mem_alloc(nvars * sizeof(void *));

for (int nvar = 0; nvar < nvars; nvar++) {

if (expr_pparams.input_class[nvar] == IARRAY_EXPR_NEQ) {

iarray_container_t *var = e->vars[nvar].c;

IARRAY_RETURN_IF_FAILED(iarray_iter_read_block_new(ctx, &iter_var[nvar], var, out_chunkshape, &iter_value[nvar],

false));

external_buffers[nvar] = ina_mem_alloc(ret->catarr->extchunknitems * ret->catarr->itemsize);

iter_var[nvar]->padding = true;

}

}

// Need for compatible containers

uint8_t **var_chunks = ina_mem_alloc(nvars * sizeof(void*));

bool *var_needs_free = ina_mem_alloc(nvars * sizeof(bool));

// Write iterator for output

ctx->prefilter_fn = (blosc2_prefilter_fn)prefilter_func;

ctx->prefilter_params = &pparams;

iarray_iter_write_block_t *iter_out;

iarray_iter_write_block_value_t out_value;

int32_t external_buffer_size = (int32_t) (ret->catarr->extchunknitems * ret->catarr->sc->typesize + BLOSC2_MAX_OVERHEAD);

void *external_buffer = NULL; // to inform the iterator that we are passing an external buffer

IARRAY_RETURN_IF_FAILED(iarray_iter_write_block_new(ctx, &iter_out, ret, out_chunkshape, &out_value, true));

// Evaluate the expression for all the chunks in variables

int64_t nchunk = 0;

while (INA_SUCCEED(iarray_iter_write_block_has_next(iter_out))) {

// The external buffer is needed *inside* the write iterator because

// this will end as a (realloc'ed) compressed chunk of a final container

// (we do so in order to avoid copies as much as possible)

// calloc to keep unwritten values as zeros

external_buffer = calloc(1, external_buffer_size);

IARRAY_RETURN_IF_FAILED(iarray_iter_write_block_next(iter_out, external_buffer, external_buffer_size));

// int32_t out_items = (int32_t)(iter_out->cur_block_size); // TODO: add a protection against cur_block_size > 2**31

// Get the chunk for each variable

for (int nvar = 0; nvar < nvars; nvar++) {

if (expr_pparams.input_class[nvar] != IARRAY_EXPR_NEQ) {

blosc2_schunk *schunk = e->vars[nvar].c->catarr->sc;

int csize = blosc2_schunk_get_lazychunk(schunk, nchunk, &var_chunks[nvar], &var_needs_free[nvar]);

if (csize < 0) {

IARRAY_TRACE1(iarray.error, "Error in retrieving chunk from schunk");

return INA_ERROR(INA_ERR_NOT_SUPPORTED);

}

bool memcpyed = *(var_chunks[nvar] + 2) & (uint8_t)BLOSC_MEMCPYED;

if (memcpyed && schunk->storage->urlpath == NULL) {

expr_pparams.input_class[nvar] = IARRAY_EXPR_EQ_NCOMP;

} else {

expr_pparams.input_class[nvar] = IARRAY_EXPR_EQ;

}

expr_pparams.inputs[nvar] = var_chunks[nvar];

expr_pparams.input_csizes[nvar] = csize;

} else {

IARRAY_RETURN_IF_FAILED(iarray_iter_read_block_next(iter_var[nvar], NULL, 0));

IARRAY_ERR_CATERVA(caterva_blosc_array_repart_chunk((int8_t *) external_buffers[nvar],

ret->catarr->extchunknitems * ret->catarr->itemsize,

iter_value[nvar].block_pointer,

ret->catarr->chunknitems * ret->catarr->itemsize,

ret->catarr));

expr_pparams.inputs[nvar] = external_buffers[nvar];

}

}

// Eval the expression for this chunk

expr_pparams.out_value = out_value; // useful for the prefilter function

// Assign the prefilter to the super-chunk context

blosc2_context *cctx = ret->catarr->sc->cctx;

blosc2_prefilter_fn old_prefilter = cctx->prefilter;

blosc2_prefilter_params *old_pparams = cctx->preparams;

cctx->prefilter = ctx->prefilter_fn;

cctx->preparams = ctx->prefilter_params;

// Do the compression with prefilters

int csize = blosc2_compress_ctx(cctx, NULL, (int32_t)ret->catarr->extchunknitems * e->typesize,

out_value.block_pointer,

(int32_t)ret->catarr->extchunknitems * e->typesize + BLOSC2_MAX_OVERHEAD);

// Reset prefilters to a possible previous value

cctx->prefilter = old_prefilter;

cctx->preparams = old_pparams;

if (csize <= 0) {

IARRAY_TRACE1(iarray.error, "Error compressing a blosc chunk");

return INA_ERROR(IARRAY_ERR_BLOSC_FAILED);

}

// Free temporary chunks

for (int nvar = 0; nvar < e->nvars; nvar++) {

if (var_needs_free[nvar] && expr_pparams.input_class[nvar] != IARRAY_EXPR_NEQ) {

free(var_chunks[nvar]);

}

}

iter_out->compressed_chunk_buffer = true;

nchunk += 1;

}

IARRAY_ITER_FINISH();

iarray_iter_write_block_free(&iter_out);

// Free initialized iterators

for (int nvar = 0; nvar < nvars; ++nvar) {

if (expr_pparams.input_class[nvar] == IARRAY_EXPR_NEQ) {

iarray_iter_read_block_free(&iter_var[nvar]);

ina_mem_free(external_buffers[nvar]);

}

}

INA_MEM_FREE_SAFE(external_buffers);

INA_MEM_FREE_SAFE(var_chunks);

INA_MEM_FREE_SAFE(var_needs_free);

INA_MEM_FREE_SAFE(iter_var);

INA_MEM_FREE_SAFE(iter_value);

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_eval(iarray_expression_t *e, iarray_container_t **container)

{

INA_VERIFY_NOT_NULL(e);

INA_VERIFY_NOT_NULL(container);

IARRAY_RETURN_IF_FAILED(iarray_empty(e->ctx, e->out_dtshape, e->out_store_properties,

container));

e->out = *container;

iarray_container_t *ret = *container;

int64_t out_chunkshape[IARRAY_DIMENSION_MAX];

for (int i = 0; i < ret->dtshape->ndim; ++i) {

out_chunkshape[i] = ret->storage->chunkshape[i];

}

uint32_t eval_method = e->ctx->cfg->eval_method & 0x3u;

switch (eval_method) {

case IARRAY_EVAL_METHOD_ITERCHUNK:

IARRAY_RETURN_IF_FAILED( iarray_eval_iterchunk(e, ret, out_chunkshape));

break;

case IARRAY_EVAL_METHOD_ITERBLOSC:

IARRAY_RETURN_IF_FAILED(iarray_eval_iterblosc(e, ret, out_chunkshape));

break;

default:

IARRAY_TRACE1(iarray.error, "Invalid eval method");

return INA_ERROR(IARRAY_ERR_INVALID_EVAL_METHOD);

}

return INA_SUCCESS;

}

ina_rc_t iarray_shape_size(iarray_dtshape_t *dtshape, size_t *size)

{

INA_VERIFY_NOT_NULL(dtshape);

INA_VERIFY_NOT_NULL(size);

*size = 1;

for (int i = 0; i < dtshape->ndim; ++i) {

*size *= dtshape->shape[i];

}

return INA_SUCCESS;

}

INA_API(ina_rc_t) iarray_expr_get_nthreads(iarray_expression_t *e, int *nthreads)

{

INA_VERIFY_NOT_NULL(e);

INA_VERIFY_NOT_NULL(nthreads);

*nthreads = e->ctx->cfg->max_num_threads;

return INA_SUCCESS;

}