# Hey Cython, this is Python 3!

# cython: language_level=3

###########################################################################################

# 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.

###########################################################################################

from collections import namedtuple

import msgpack

from . cimport ciarray_ext as ciarray

from .ciarray_ext cimport int64_t

import numpy as np

cimport numpy as np

import zarr

import s3fs

import cython

from cpython.pycapsule cimport PyCapsule_New, PyCapsule_GetPointer

from libc.stdlib cimport malloc, free

from libc.string cimport memcpy

import iarray as ia

from iarray import udf

from cpython cimport (

PyObject_GetBuffer,

PyBuffer_Release,

PyBUF_SIMPLE,

)

# dtype conversion tables: udf <-> iarray

udf2ia_dtype = {

"f64": ciarray.IARRAY_DATA_TYPE_DOUBLE,

"f32": ciarray.IARRAY_DATA_TYPE_FLOAT,

"i64": ciarray.IARRAY_DATA_TYPE_INT64,

"i32": ciarray.IARRAY_DATA_TYPE_INT32,

"i16": ciarray.IARRAY_DATA_TYPE_INT16,

"i8": ciarray.IARRAY_DATA_TYPE_INT8,

"u64": ciarray.IARRAY_DATA_TYPE_UINT64,

"u32": ciarray.IARRAY_DATA_TYPE_UINT32,

"u16": ciarray.IARRAY_DATA_TYPE_UINT16,

"u8": ciarray.IARRAY_DATA_TYPE_UINT8,

"bool": ciarray.IARRAY_DATA_TYPE_BOOL,

}

ia2udf_dtype = {v: k for k, v in udf2ia_dtype.items()}

# dtype conversion tables: numpy <-> iarray

np2ia_dtype = {

np.float64: ciarray.IARRAY_DATA_TYPE_DOUBLE,

np.float32: ciarray.IARRAY_DATA_TYPE_FLOAT,

np.int64: ciarray.IARRAY_DATA_TYPE_INT64,

np.int32: ciarray.IARRAY_DATA_TYPE_INT32,

np.int16: ciarray.IARRAY_DATA_TYPE_INT16,

np.int8: ciarray.IARRAY_DATA_TYPE_INT8,

np.uint64: ciarray.IARRAY_DATA_TYPE_UINT64,

np.uint32: ciarray.IARRAY_DATA_TYPE_UINT32,

np.uint16: ciarray.IARRAY_DATA_TYPE_UINT16,

np.uint8: ciarray.IARRAY_DATA_TYPE_UINT8,

np.bool_: ciarray.IARRAY_DATA_TYPE_BOOL,

}

ia2np_dtype = {v: k for k, v in np2ia_dtype.items()}

# datetime64 and timedelta64

dtype_str2dtype = {bo + l + '[' + size + ']': np.int64 for bo in ['<', '>'] for l in ['M8', 'm8']

for size in ['Y', 'M', 'D', 'h', 's', 'ms', 'us', 'μs', 'ns', 'ps', 'fs', 'as']}

# NaT (not a time)

dtype_str2dtype.update({bo + l: np.int64 for bo in ['<', '>'] for l in ['M8', 'm8']})

# integers

dtype_str2dtype.update({bo + 'i' + size: dtype for bo in ['<', '>']

for (size, dtype) in [('1', np.int8), ('2', np.int16), ('4', np.int32), ('8', np.int64)]})

# unsigned integers

dtype_str2dtype.update({bo + 'u' + size: dtype for bo in ['<', '>']

for (size, dtype) in [('1', np.uint8), ('2', np.uint16), ('4', np.uint32), ('8', np.uint64)]})

# floats

dtype_str2dtype.update({bo + 'f' + size: dtype for bo in ['<', '>']

for (size, dtype) in [('2', np.int16), ('4', np.float32), ('8', np.float64)]})

# booleans

dtype_str2dtype.update({'|b1': np.bool_, '|u1': np.bool_})

def f_np2ia_dtype(dtype):

if type(dtype) != type:

return np2ia_dtype[dtype.type]

return np2ia_dtype[dtype]

def compress_squeeze(data, selectors):

return tuple(d for d, s in zip(data, selectors) if not s)

class IArrayError(Exception):

pass

def iarray_check(error):

if error != 0:

raise IArrayError(str(ciarray.iarray_err_strerror(error)))

IARRAY_ERR_EVAL_ENGINE_FAILED = ciarray.IARRAY_ERR_EVAL_ENGINE_FAILED

IARRAY_ERR_EVAL_ENGINE_NOT_COMPILED = ciarray.IARRAY_ERR_EVAL_ENGINE_NOT_COMPILED

IARRAY_ERR_EVAL_ENGINE_OUT_OF_RANGE = ciarray.IARRAY_ERR_EVAL_ENGINE_OUT_OF_RANGE

cdef set_storage(cfg, ciarray.iarray_storage_t *cstore):

cstore.contiguous = cfg.contiguous

for i in range(len(cfg.chunks)):

cstore.chunkshape[i] = cfg.chunks[i]

cstore.blockshape[i] = cfg.blocks[i]

if cfg.urlpath is not None:

urlpath = cfg.urlpath.encode("utf-8") if isinstance(cfg.urlpath, str) else cfg.urlpath

cstore.urlpath = cfg.urlpath

else:

cstore.urlpath = NULL

cdef class ReadBlockIter:

cdef ciarray.iarray_iter_read_block_t *ia_read_iter

cdef ciarray.iarray_iter_read_block_value_t ia_block_val

cdef Container container

cdef int dtype

cdef int flag

cdef object Info

def __cinit__(self, container, block):

self.container = container

cdef ciarray.int64_t block_[ciarray.IARRAY_DIMENSION_MAX]

if block is None:

block = container.chunks

for i in range(len(block)):

block_[i] = block[i]

iarray_check(ciarray.iarray_iter_read_block_new(self.container.context.ia_ctx, &self.ia_read_iter,

self.container.ia_container, block_, &self.ia_block_val, False))

self.dtype = f_np2ia_dtype(self.container.dtype)

self.Info = namedtuple('Info', 'index elemindex nblock shape size slice')

def __dealloc__(self):

ciarray.iarray_iter_read_block_free(&self.ia_read_iter)

def __iter__(self):

return self

def __next__(self):

if ciarray.iarray_iter_read_block_has_next(self.ia_read_iter) != 0:

raise StopIteration

iarray_check(ciarray.iarray_iter_read_block_next(self.ia_read_iter, NULL, 0))

shape = tuple(self.ia_block_val.block_shape[i] for i in range(self.container.ndim))

size = np.prod(shape)

if self.dtype == ciarray.IARRAY_DATA_TYPE_DOUBLE:

view = <np.float64_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_FLOAT:

view = <np.float32_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT64:

view = <np.int64_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT32:

view = <np.int32_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT16:

view = <np.int16_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT8:

view = <np.int8_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT64:

view = <np.uint64_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT32:

view = <np.uint32_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT16:

view = <np.uint16_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT8:

view = <np.uint8_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_BOOL:

view = <ciarray.bool[:size]> self.ia_block_val.block_pointer

a = np.asarray(view)

if self.container.np_dtype is not None:

a = a.astype(dtype=self.container.np_dtype)

elem_index = tuple(self.ia_block_val.elem_index[i] for i in range(self.container.ndim))

index = tuple(self.ia_block_val.block_index[i] for i in range(self.container.ndim))

nblock = self.ia_block_val.nblock

slice_ = tuple([slice(i, i + s) for i, s in zip(elem_index, shape)])

info = self.Info(index=index, elemindex=elem_index, nblock=nblock, shape=shape,

size=size, slice=slice_)

return info, a.reshape(shape)

cdef class WriteBlockIter:

cdef ciarray.iarray_iter_write_block_t *ia_write_iter

cdef ciarray.iarray_iter_write_block_value_t ia_block_val

cdef Container container

cdef int dtype

cdef int flag

cdef object Info

def __cinit__(self, c, block=None):

self.container = c

cdef ciarray.int64_t block_[ciarray.IARRAY_DIMENSION_MAX]

if block is None:

# The block for iteration has always be provided

block = c.chunks

for i in range(len(block)):

block_[i] = block[i]

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(self.container.urlpath, self.container.mode, True)

iarray_check(ciarray.iarray_iter_write_block_new(self.container.context.ia_ctx,

&self.ia_write_iter,

self.container.ia_container,

block_,

&self.ia_block_val,

False))

self.dtype = f_np2ia_dtype(self.container.dtype)

self.Info = namedtuple('Info', 'index elemindex nblock shape size')

def __dealloc__(self):

ciarray.iarray_iter_write_block_free(&self.ia_write_iter)

def __iter__(self):

return self

def __next__(self):

if ciarray.iarray_iter_write_block_has_next(self.ia_write_iter) != 0:

raise StopIteration

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(self.container.urlpath, self.container.mode, True)

iarray_check(ciarray.iarray_iter_write_block_next(self.ia_write_iter, NULL, 0))

shape = tuple(self.ia_block_val.block_shape[i] for i in range(self.container.ndim))

size = np.prod(shape)

if self.dtype == ciarray.IARRAY_DATA_TYPE_DOUBLE:

view = <np.float64_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_FLOAT:

view = <np.float32_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT64:

view = <np.int64_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT32:

view = <np.int32_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT16:

view = <np.int16_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_INT8:

view = <np.int8_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT64:

view = <np.uint64_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT32:

view = <np.uint32_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT16:

view = <np.uint16_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_UINT8:

view = <np.uint8_t[:size]> self.ia_block_val.block_pointer

elif self.dtype == ciarray.IARRAY_DATA_TYPE_BOOL:

view = <ciarray.bool[:size]> self.ia_block_val.block_pointer

a = np.asarray(view)

elem_index = tuple(self.ia_block_val.elem_index[i] for i in range(self.container.ndim))

index = tuple(self.ia_block_val.block_index[i] for i in range(self.container.ndim))

nblock = self.ia_block_val.nblock

info = self.Info(index=index, elemindex=elem_index, nblock=nblock, shape=shape, size=size)

return info, a.reshape(shape)

cdef class IArrayInit:

def __cinit__(self):

iarray_check(ciarray.iarray_init())

def __dealloc__(self):

ciarray.iarray_destroy()

cdef class Config:

cdef ciarray.iarray_config_t config

def __init__(self, compression_codec, compression_meta, compression_level, compression_favor,

use_dict, filters, max_num_threads, fp_mantissa_bits, eval_method, btune, split_mode):

self.config.compression_codec = compression_codec.value

# Avoid error in case compression_meta < 0

cdef ciarray.uint8_t compression_meta_

if compression_meta is not None:

compression_meta_ = <ciarray.int8_t> compression_meta

self.config.compression_meta = compression_meta_

self.config.compression_level = compression_level

self.config.compression_favor = compression_favor.value

self.config.use_dict = 1 if use_dict else 0

cdef int filter_flags = 0

# TODO: filters are really a pipeline, and here we are just ORing them, which is tricky.

# This should be fixed (probably at C iArray level and then propagating the change here).

# At any rate, `filters` should be a list for displaying purposes in high level Config().

for f in filters:

filter_flags |= f.value

self.config.filter_flags = filter_flags

if eval_method == ia.Eval.AUTO:

method = ciarray.IARRAY_EVAL_METHOD_AUTO

elif eval_method == ia.Eval.ITERBLOSC:

method = ciarray.IARRAY_EVAL_METHOD_ITERBLOSC

elif eval_method == ia.Eval.ITERCHUNK:

method = ciarray.IARRAY_EVAL_METHOD_ITERCHUNK

else:

raise ValueError("eval_method method not recognized:", eval_method)

self.config.eval_method = method

self.config.max_num_threads = max_num_threads

self.config.fp_mantissa_bits = fp_mantissa_bits

self.config.btune = btune

self.config.splitmode = split_mode.value

def _to_dict(self):

return <object> self.config

cdef class Context:

cdef ciarray.iarray_context_t *ia_ctx

cdef public object cfg

def __init__(self, cfg):

cdef ciarray.iarray_config_t cfg_ = cfg._to_dict()

# Set default contiguous correctly

if cfg.contiguous is None:

cfg.contiguous = False

iarray_check(ciarray.iarray_context_new(&cfg_, &self.ia_ctx))

self.cfg = cfg

def __dealloc__(self):

ciarray.iarray_context_free(&self.ia_ctx)

def to_capsule(self):

return PyCapsule_New(self.ia_ctx, <char*>"iarray_context_t*", NULL)

cdef class IaDTShape:

cdef ciarray.iarray_dtshape_t ia_dtshape

def __cinit__(self, dtshape):

self.ia_dtshape.ndim = len(dtshape.shape)

self.ia_dtshape.dtype = f_np2ia_dtype(dtshape.dtype)

self.ia_dtshape.dtype_size = np.dtype(dtshape.dtype).itemsize

for i in range(len(dtshape.shape)):

self.ia_dtshape.shape[i] = dtshape.shape[i]

cdef to_dict(self):

return <object> self.ia_dtshape

@property

def ndim(self):

return self.ia_dtshape.ndim

@property

def dtype(self):

return ia2np_dtype[self.ia_dtshape.dtype]

@property

def dtype_size(self):

return self.ia_dtshape.dtype_size

@property

def shape(self):

shape = []

for i in range(self.ndim):

shape.append(self.ia_dtshape.shape[i])

return tuple(shape)

def __str__(self):

return self.ia_dtshape

cdef class RandomContext:

cdef ciarray.iarray_random_ctx_t *random_ctx

cdef Context context

def __init__(self, ctx, seed, rng):

self.context = ctx

cdef ciarray.iarray_random_ctx_t* r_ctx

if rng == ia.RandomGen.MRG32K3A:

iarray_check(ciarray.iarray_random_ctx_new(self.context.ia_ctx, seed, ciarray.IARRAY_RANDOM_RNG_MRG32K3A, &r_ctx))

else:

raise ValueError("Random generator unknown")

self.random_ctx = r_ctx

def __dealloc__(self):

if self.context is not None and self.context.ia_ctx != NULL:

ciarray.iarray_random_ctx_free(self.context.ia_ctx, &self.random_ctx)

self.context = None

def to_capsule(self):

return PyCapsule_New(self.random_ctx, <char*>"iarray_random_ctx_t*", NULL)

def _is_s3_store(urlpath):

if urlpath[:5] == "s3://":

return True

return False

cpdef _zarray_from_proxy(urlpath):

if _is_s3_store(urlpath):

s3 = s3fs.S3FileSystem(anon=True)

store = s3fs.S3Map(root=urlpath, s3=s3)

return zarr.open(store)

else:

return zarr.open(urlpath)

cdef class Container:

cdef ciarray.iarray_container_t *ia_container

cdef Context context

cdef Py_ssize_t bp_shape[ciarray.IARRAY_DIMENSION_MAX]

cdef Py_ssize_t bp_strides[ciarray.IARRAY_DIMENSION_MAX]

cdef int view_count

def __init__(self, ctx, c):

if ctx is None:

raise ValueError("You must pass a context to the Container constructor")

if c is None:

raise ValueError("You must pass a Capsule to the C container struct of the Container constructor")

self.context = ctx

self.ia_container = <ciarray.iarray_container_t*> PyCapsule_GetPointer(c, <char*>"iarray_container_t*")

self.buffer = None

self.view_count = 0

def __dealloc__(self):

if self.context is not None and self.context.ia_ctx != NULL:

# if self.view_count > 0:

# # TODO: set Blosc flag `cframe_avoid_free = True`

# pass

ciarray.iarray_container_free(self.context.ia_ctx, &self.ia_container)

self.context = None

# THERE ARE COLLISIONS WITH LAZY EXPRESSIONS

#

# def __getbuffer__(self, Py_buffer *buffer, int flags):

# dtype = np.dtype(self.dtype)

#

# ctx = Context(ia.Config())

# cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t *> PyCapsule_GetPointer(

# ctx.to_capsule(),

# <char *> "iarray_context_t*")

#

# cdef ciarray.uint8_t *cframe

# cdef ciarray.int64_t cframe_len

# cdef ciarray.bool needs_free

# iarray_check(ciarray.iarray_to_cframe(ctx_, self.ia_container, &cframe, &cframe_len, &needs_free))

#

# self.bp_shape[0] = cframe_len

# self.bp_strides[0] = 1

#

# buffer.buf = <char *> cframe

# buffer.format = 'B' # unsigned bytes (compressed array)

# buffer.internal = NULL # see References

# buffer.readonly = 1

# buffer.obj = self

# buffer.itemsize = 1

# buffer.len = cframe_len

# buffer.ndim = 1

# buffer.shape = self.bp_shape

# buffer.strides = self.bp_strides

# buffer.suboffsets = NULL

# if not needs_free:

# self.view_count += 1

# def __releasebuffer__(self, Py_buffer *buffer):

# self.view_count -= 1

def to_capsule(self):

return PyCapsule_New(self.ia_container, <char*>"iarray_container_t*", NULL)

def to_cframe(self):

""" Create a cframe.

Returns

-------

A bytes object containing the cframe

"""

return get_cframe(self)

@property

def ndim(self):

"""Number of array dimensions."""

cdef ciarray.iarray_dtshape_t dtshape

iarray_check(ciarray.iarray_get_dtshape(self.context.ia_ctx, self.ia_container, &dtshape))

return dtshape.ndim

@property

def shape(self):

"""Tuple of array dimensions."""

cdef ciarray.iarray_dtshape_t dtshape

iarray_check(ciarray.iarray_get_dtshape(self.context.ia_ctx, self.ia_container, &dtshape))

shape = [dtshape.shape[i] for i in range(self.ndim)]

return tuple(shape)

@property

def chunks(self):

"""Tuple of chunk dimensions."""

cdef ciarray.iarray_storage_t storage

iarray_check(ciarray.iarray_get_storage(self.context.ia_ctx, self.ia_container, &storage))

if self.is_view:

return None

chunks = [storage.chunkshape[i] for i in range(self.ndim)]

return tuple(chunks)

@property

def blocks(self):

"""Tuple of block dimensions."""

cdef ciarray.iarray_storage_t storage

iarray_check(ciarray.iarray_get_storage(self.context.ia_ctx, self.ia_container, &storage))

if self.is_view:

return None

blocks = [storage.blockshape[i] for i in range(self.ndim)]

return tuple(blocks)

@property

def dtype(self):

"""Data-type of the array’s elements."""

cdef ciarray.iarray_dtshape_t dtshape

iarray_check(ciarray.iarray_get_dtshape(self.context.ia_ctx, self.ia_container, &dtshape))

return ia2np_dtype[dtshape.dtype]

@property

def np_dtype(self):

"""The array-protocol typestring of the np.dtype object to use."""

return self.attrs["np_dtype"] if "np_dtype" in self.attrs.keys() else None

@np_dtype.setter

def np_dtype(self, value):

if value is not None:

self.attrs["np_dtype"] = np.dtype(value).str

@property

def dtshape(self):

"""The :py:obj:`DTShape` of the array."""

return ia.DTShape(self.shape, self.dtype)

@property

def cratio(self):

"""Array compression ratio."""

# Return zarr array values if it is a zproxy

if "zproxy_urlpath" in self.attrs:

urlpath = self.attrs["zproxy_urlpath"]

z = _zarray_from_proxy(urlpath)

return z.nbytes / z.nbytes_stored

# It is a normal iarray

cdef ciarray.int64_t nbytes, cbytes

iarray_check(ciarray.iarray_container_info(self.ia_container, &nbytes, &cbytes))

return <double>nbytes / <double>cbytes

@property

def cfg(self):

return self.context.cfg

@property

def urlpath(self):

return self.context.cfg.urlpath

@property

def mode(self):

return self.context.cfg.mode

def __getitem__(self, key):

# key has been massaged already

start, stop, squeeze_mask = key

with ia.config(cfg=self.cfg) as cfg:

return get_slice(cfg, self, start, stop, squeeze_mask, True, None)

@property

def is_view(self):

"""Whether the :ref:`IArray` is a view or not.

"""

cdef ciarray.bool view

iarray_check(ciarray.iarray_is_view(self.context.ia_ctx, self.ia_container, &view))

return view

cdef class Expression:

cdef object expression

cdef ciarray.iarray_expression_t *ia_expr

cdef Context context

cdef ciarray.bool zproxy_op

def __init__(self, cfg):

self.cfg = cfg

self.context = Context(cfg)

cdef ciarray.iarray_expression_t* e

iarray_check(

ciarray.iarray_expr_new(self.context.ia_ctx, f_np2ia_dtype(cfg.dtype), &e)

)

self.ia_expr = e

self.expression = None

self.dtshape = None

self.zproxy_op = False

def __dealloc__(self):

if self.context is not None and self.context.ia_ctx != NULL:

ciarray.iarray_expr_free(self.context.ia_ctx, &self.ia_expr)

self.context = None

def bind(self, var, c):

var2 = var.encode("utf-8") if isinstance(var, str) else var

if "zproxy_urlpath" in c.attrs:

# To not release the GIL when evaluating

self.zproxy_op = True

cdef ciarray.iarray_container_t *c_ = <ciarray.iarray_container_t*> PyCapsule_GetPointer(

c.to_capsule(), <char*>"iarray_container_t*")

iarray_check(ciarray.iarray_expr_bind(self.ia_expr, var2, c_))

def bind_param(self, value, type_):

cdef ciarray.iarray_user_param_t user_param;

if type_ is udf.float64:

user_param.f64 = value

elif type_ is udf.float32:

user_param.f32 = value

elif type_ is udf.int64:

user_param.i64 = value

elif type_ is udf.int32:

user_param.i32 = value

elif type_ is udf.bool:

user_param.b = value

iarray_check(

ciarray.iarray_expr_bind_param(self.ia_expr, user_param)

)

def bind_out_properties(self, dtshape):

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

cdef ciarray.iarray_storage_t store_

set_storage(self.cfg, &store_)

iarray_check(ciarray.iarray_expr_bind_out_properties(self.ia_expr, &dtshape_, &store_))

self.dtshape = dtshape

def compile(self, expr):

expr = expr.encode("utf-8") if isinstance(expr, str) else expr

iarray_check(ciarray.iarray_expr_compile(self.ia_expr, expr))

self.expression = expr

def compile_bc(self, bc, name):

name = name.encode()

cdef int bc_len = len(bc)

iarray_check(ciarray.iarray_expr_compile_udf(self.ia_expr, bc_len, bc, name))

self.expression = "user_defined_function"

def compile_udf(self, func):

self.compile_bc(func.bc, func.name)

def eval(self):

cdef ciarray.iarray_container_t *c;

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(self.cfg.urlpath, self.cfg.mode)

# Update the chunks and blocks with the correct values

self.update_chunks_blocks()

if self.zproxy_op:

error = ciarray.iarray_eval(self.ia_expr, &c)

else:

with nogil:

error = ciarray.iarray_eval(self.ia_expr, &c)

iarray_check(error)

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

return ia.IArray(self.context, c_c)

def update_chunks_blocks(self):

cdef int nvars = self.ia_expr.nvars;

cdef ciarray.int8_t ndim = <ciarray.int8_t> self.dtshape["ndim"]

cdef ciarray.iarray_storage_t storage_0

cdef ciarray.iarray_storage_t storage_i

if self.cfg.chunks is None:

# Set blocks and chunks to the ones from the operands in case all of them are equal

if nvars > 0:

equal = True

ciarray.iarray_get_storage(self.ia_expr.ctx, self.ia_expr.vars[0].c, &storage_0)

chunks_0 = list(storage_0.chunkshape)[:ndim]

blocks_0 = list(storage_0.blockshape)[:ndim]

for i in range(1, nvars):

ciarray.iarray_get_storage(self.ia_expr.ctx, self.ia_expr.vars[i].c, &storage_i)

chunks_i = list(storage_i.chunkshape)[:ndim]

blocks_i = list(storage_i.blockshape)[:ndim]

if chunks_i != chunks_0 or blocks_i != blocks_0:

equal = False

break

if equal:

self.ia_expr.out_store_properties.chunkshape = storage_0.chunkshape

self.ia_expr.out_store_properties.blockshape = storage_0.blockshape

chunks = list(self.ia_expr.out_store_properties.chunkshape)[:ndim]

blocks = list(self.ia_expr.out_store_properties.blockshape)[:ndim]

self.cfg.chunks = chunks

self.cfg.blocks = blocks

#

# Iarray container constructors

#

def copy(cfg, src):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(

ctx.to_capsule(), <char*>"iarray_context_t*")

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_container_t *c

cdef ciarray.iarray_container_t *src_ = <ciarray.iarray_container_t *> PyCapsule_GetPointer(

src.to_capsule(), <char*>"iarray_container_t*")

if "zproxy_urlpath" in src.attrs:

error = ciarray.iarray_copy(ctx_, src_, False, &store_, &c)

else:

with nogil:

error = ciarray.iarray_copy(ctx_, src_, False, &store_, &c)

iarray_check(error)

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

def uninit(cfg, dtshape):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_uninit(ctx_, &dtshape_, &store_, &c))

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

def arange(cfg, slice_, dtshape):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

start, stop, step = slice_.start, slice_.stop, slice_.step

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_arange(ctx_, &dtshape_, start, step, &store_, &c))

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

def linspace(cfg, start, stop, dtshape):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_linspace(ctx_, &dtshape_, start, stop, &store_, &c))

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

def zeros(cfg, dtshape):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_zeros(ctx_, &dtshape_, &store_, &c))

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

def ones(cfg, dtshape):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_ones(ctx_, &dtshape_, &store_, &c))

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

def full(cfg, fill_value, dtshape):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

dtshape = IaDTShape(dtshape).to_dict()

cdef ciarray.iarray_dtshape_t dtshape_ = <ciarray.iarray_dtshape_t> dtshape

cdef ciarray.iarray_storage_t store_

set_storage(cfg, &store_)

# Check that we are not inadvertently overwriting anything

ia._check_access_mode(cfg.urlpath, cfg.mode)

cdef ciarray.iarray_container_t *c

dtype = ia2np_dtype[dtshape_.dtype]

# The ciarray.iarray_fill function requires a void pointer

nparr = np.array([fill_value], dtype=dtype)

cdef Py_buffer *val = <Py_buffer *> malloc(sizeof(Py_buffer))

PyObject_GetBuffer(nparr, val, PyBUF_SIMPLE)

iarray_check(ciarray.iarray_fill(ctx_, &dtshape_, val.buf, &store_, &c))

PyBuffer_Release(val)

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

a = ia.IArray(ctx, c_c)

a.np_dtype = cfg.np_dtype

return a

cdef get_cfg_from_container(cfg, ciarray.iarray_context_t *ctx, ciarray.iarray_container_t *c, urlpath):

cdef ciarray.iarray_config_t cfg_

ciarray.iarray_get_cfg(ctx, c, &cfg_)

clevel = cfg_.compression_level

codec = ia.Codec(cfg_.compression_codec)

zfp_meta = cfg_.compression_meta

mantissa_bits = cfg_.fp_mantissa_bits

filters = []

if cfg_.filter_flags & ciarray.IARRAY_COMP_TRUNC_PREC:

filters.append(ia.Filter.TRUNC_PREC)

if cfg_.filter_flags & ciarray.IARRAY_COMP_DELTA:

filters.append(ia.Filter.DELTA)

if cfg_.filter_flags & ciarray.IARRAY_COMP_BITSHUFFLE:

filters.append(ia.Filter.BITSHUFFLE)

if cfg_.filter_flags & ciarray.IARRAY_COMP_SHUFFLE:

filters.append(ia.Filter.SHUFFLE)

cdef ciarray.iarray_dtshape_t dtshape;

ciarray.iarray_get_dtshape(ctx, c, &dtshape)

dtype = ia2np_dtype[dtshape.dtype]

cdef const char *name = "np_dtype"

cdef ciarray.bool exists

iarray_check(ciarray.iarray_vlmeta_exists(ctx, c, name, &exists))

cdef ciarray.iarray_metalayer_t meta

if exists:

iarray_check(ciarray.iarray_vlmeta_get(ctx, c, name, &meta))

np_dtype = meta.sdata[:meta.size]

np_dtype = msgpack.unpackb(np_dtype)

else:

np_dtype = None

cdef ciarray.iarray_storage_t storage;

ciarray.iarray_get_storage(ctx, c, &storage)

chunks = tuple(storage.chunkshape[i] for i in range(dtshape.ndim))

blocks = tuple(storage.blockshape[i] for i in range(dtshape.ndim))

contiguous = storage.contiguous

# The config params should already have been checked

ia._defaults.check_compat = False

c_cfg = ia.Config(

codec=codec,

zfp_meta=zfp_meta,

clevel=clevel,

filters=filters,

fp_mantissa_bits = mantissa_bits,

use_dict=False,

favor=cfg.favor,

nthreads=cfg.nthreads,

eval_method=cfg.eval_method,

seed=cfg.seed,

random_gen=cfg.random_gen,

btune=False, # we have not used btune to load/open

dtype=dtype,

np_dtype=np_dtype,

chunks=chunks,

blocks=blocks,

urlpath=urlpath,

contiguous=contiguous,

mode=cfg.mode,

)

return c_cfg

def load(cfg, urlpath):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

urlpath = urlpath.encode("utf-8") if isinstance(urlpath, str) else urlpath

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_container_load(ctx_, urlpath, &c))

# Fetch config from the new container

c_cfg = get_cfg_from_container(cfg, ctx_, c, None)

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

return ia.IArray(Context(c_cfg), c_c)

def open(cfg, urlpath):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

urlpath = urlpath.encode("utf-8") if isinstance(urlpath, str) else urlpath

cdef ciarray.iarray_container_t *c

iarray_check(ciarray.iarray_container_open(ctx_, urlpath, &c))

# Fetch config from the recently open container

c_cfg = get_cfg_from_container(cfg, ctx_, c, urlpath)

c_c = PyCapsule_New(c, <char*>"iarray_container_t*", NULL)

iarr = ia.IArray(Context(c_cfg), c_c)

if "zproxy_urlpath" in iarr.attrs:

set_zproxy_postfilter(iarr)

return iarr

def set_orthogonal_selection(cfg, dst, selection, ndarray):

ctx = Context(cfg)

cdef ciarray.iarray_context_t *ctx_ = <ciarray.iarray_context_t*> PyCapsule_GetPointer(ctx.to_capsule(),

<char*>"iarray_context_t*")

cdef ciarray.iarray_container_t *data_ = <ciarray.iarray_container_t*> PyCapsule_GetPointer(dst.to_capsule(),

<char*>"iarray_container_t*")

ndim = dst.ndim

interface = ndarray.__array_interface__