arraykit 0.6.1__tar.gz → 0.6.3__tar.gz

{arraykit-0.6.1/arraykit.egg-info → arraykit-0.6.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arraykit
-Version: 0.6.1
+Version: 0.6.3
 Summary: Array utilities for StaticFrame
 Home-page: https://github.com/static-frame/arraykit
 Author: Christopher Ariza, Brandt Bucher, Charles Burkland

{arraykit-0.6.1 → arraykit-0.6.3}/README.rst

@@ -37,6 +37,20 @@ ArrayKit requires the following:
 What is New in ArrayKit
 -------------------------
 
+0.6.3
+............
+
+Optimized memory allocation strategy for ``nonzero_1d()``.
+
+
+0.6.2
+............
+
+Extended ``nonzero_1d()`` to support non-contiguous arrays.
+
+Optimizations to ``TriMap`` when mapping to object and flexible dtypes.
+
+
 0.6.1
 ............
 

{arraykit-0.6.1 → arraykit-0.6.3/arraykit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arraykit
-Version: 0.6.1
+Version: 0.6.3
 Summary: Array utilities for StaticFrame
 Home-page: https://github.com/static-frame/arraykit
 Author: Christopher Ariza, Brandt Bucher, Charles Burkland

{arraykit-0.6.1 → arraykit-0.6.3}/setup.py

@@ -5,7 +5,7 @@ from setuptools import Extension  # type: ignore
 from setuptools import setup
 from pathlib import Path
 
-AK_VERSION = '0.6.1'
+AK_VERSION = '0.6.3'
 
 def get_long_description() -> str:
     return '''The ArrayKit library provides utilities for creating and transforming NumPy arrays, implementing performance-critical StaticFrame operations as Python C extensions.

{arraykit-0.6.1 → arraykit-0.6.3}/src/_arraykit.c

@@ -836,9 +836,9 @@ AK_TP_ResolveLineResetField(AK_TypeParser* tp,
 static char* TRUE_LOWER = "true";
 static char* TRUE_UPPER = "TRUE";
 
-#define ERROR_NO_DIGITS 1
-#define ERROR_OVERFLOW 2
-#define ERROR_INVALID_CHARS 3
+#define AK_ERROR_NO_DIGITS 1
+#define AK_ERROR_OVERFLOW 2
+#define AK_ERROR_INVALID_CHARS 3
 
 // Convert a Py_UCS4 array to a signed integer. Extended from pandas/_libs/src/parser/tokenizer.c. Sets `error` to values greater than 0 on error; never sets error on success.
 static inline npy_int64
@@ -866,7 +866,7 @@ AK_UCS4_to_int64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
 
     // Check that there is a first digit.
     if (!AK_is_digit(*p)) {
-        *error = ERROR_NO_DIGITS;
+        *error = AK_ERROR_NO_DIGITS;
         return 0;
     }
     if (isneg) {
@@ -891,7 +891,7 @@ AK_UCS4_to_int64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
                     if (p >= end) return number;
                     d = *p;
                 } else {
-                    *error = ERROR_OVERFLOW;
+                    *error = AK_ERROR_OVERFLOW;
                     return 0;
                 }
             }
@@ -904,7 +904,7 @@ AK_UCS4_to_int64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
                     if (p >= end) return number;
                     d = *p;
                 } else {
-                    *error = ERROR_OVERFLOW;
+                    *error = AK_ERROR_OVERFLOW;
                     return 0;
                 }
             }
@@ -931,7 +931,7 @@ AK_UCS4_to_int64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
                     if (p >= end) return number;
                     d = *p;
                 } else {
-                    *error = ERROR_OVERFLOW;
+                    *error = AK_ERROR_OVERFLOW;
                     return 0;
                 }
             }
@@ -944,7 +944,7 @@ AK_UCS4_to_int64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
                     if (p >= end) return number;
                     d = *p;
                 } else {
-                    *error = ERROR_OVERFLOW;
+                    *error = AK_ERROR_OVERFLOW;
                     return 0;
                 }
             }
@@ -952,7 +952,7 @@ AK_UCS4_to_int64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
     }
     while (p < end) {
         if (!AK_is_space(*p)) {
-            *error = ERROR_INVALID_CHARS;
+            *error = AK_ERROR_INVALID_CHARS;
             return 0;
         }
         p++;
@@ -975,7 +975,7 @@ AK_UCS4_to_uint64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
         if (p >= end) return number;
     }
     if (*p == '-') {
-        *error = ERROR_INVALID_CHARS;
+        *error = AK_ERROR_INVALID_CHARS;
         return 0;
     } else if (*p == '+') {
         p++;
@@ -984,7 +984,7 @@ AK_UCS4_to_uint64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
 
     // Check that there is a first digit.
     if (!AK_is_digit(*p)) {
-        *error = ERROR_NO_DIGITS;
+        *error = AK_ERROR_NO_DIGITS;
         return 0;
     }
     // If number is less than pre_max, at least one more digit can be processed without overflowing.
@@ -1006,7 +1006,7 @@ AK_UCS4_to_uint64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
                 if (p >= end) return number;
                 d = *p;
             } else {
-                *error = ERROR_OVERFLOW;
+                *error = AK_ERROR_OVERFLOW;
                 return 0;
             }
         }
@@ -1019,14 +1019,14 @@ AK_UCS4_to_uint64(Py_UCS4 *p_item, Py_UCS4 *end, int *error, char tsep)
                 if (p >= end) return number;
                 d = *p;
             } else {
-                *error = ERROR_OVERFLOW;
+                *error = AK_ERROR_OVERFLOW;
                 return 0;
             }
         }
     }
     while (p < end) {
         if (!AK_is_space(*p)) {
-            *error = ERROR_INVALID_CHARS;
+            *error = AK_ERROR_INVALID_CHARS;
             return 0;
         }
         p++;
@@ -3535,21 +3535,9 @@ resolve_dtype_iter(PyObject *Py_UNUSED(m), PyObject *arg) {
 //------------------------------------------------------------------------------
 // general utility
 
-#define NONZERO_APPEND_INDEX {                                               \
-    if (AK_UNLIKELY(count == capacity)) {                                    \
-        capacity <<= 1;                                                      \
-        indices = (npy_int64*)realloc(indices, sizeof(npy_int64) * capacity);\
-        if (indices == NULL) {                                               \
-            return NULL;                                                     \
-        }                                                                    \
-    }                                                                        \
-    indices[count++] = p - p_start;                                          \
-}                                                                            \
-
-// Given a Boolean, contiguous 1D array, return the index positions in an int64 array.
+// Given a Boolean, contiguous 1D array, return the index positions in an int64 array. Through experimentation it has been verified that doing full-size allocation of memory provides the best performance at all scales. Using NpyIter, or using, bit masks does not improve performance over pointer arithmetic. Prescanning for all empty is very effective. Note that NumPy benefits from first counting the nonzeros, then allocating only enough data for the expexted number of indices.
 static inline PyObject*
 AK_nonzero_1d(PyArrayObject* array) {
-    // the maxiumum number of indices we could return is the size of the array; if this is under a certain number, probably better to just allocate that rather than reallocate
     PyObject* final;
     npy_intp count_max = PyArray_SIZE(array);
 
@@ -3562,116 +3550,75 @@ AK_nonzero_1d(PyArrayObject* array) {
     lldiv_t size_div = lldiv((long long)count_max, 8); // quot, rem
 
     Py_ssize_t count = 0;
-    // the maximum number of collected integers is equal to or less than count_max; for small count_max, we can just set that value; for large size, we set it to half the size
-    Py_ssize_t capacity = count_max < 1024 ? count_max : count_max / 8;
+    // the maximum number of collected integers is equal to or less than count_max
+    Py_ssize_t capacity = count_max;
     npy_int64* indices = (npy_int64*)malloc(sizeof(npy_int64) * capacity);
 
-    // array is contiguous, 1d, boolean
-    npy_bool* p_start = (npy_bool*)PyArray_DATA(array);
-    npy_bool* p = p_start;
-    npy_bool* p_end = p + count_max;
-    npy_bool* p_end_roll = p_end - size_div.rem;
-
     NPY_BEGIN_THREADS_DEF;
     NPY_BEGIN_THREADS;
-    // Through experimentation it has been verified that doing full-size allocation of memory does not permit outperforming NumPy at 10_000_000 scale; but doing less optimizations does help.
-    // Using bit masks does not improve perforamnce over pointer arithmetic.
-    // Prescanning for all empty is very effective.
 
-    while (p < p_end_roll) {
-        if (*(npy_uint64*)p == 0) {
-            p += 8; // no true within this 8 byte roll region
-            continue;
+    if (PyArray_IS_C_CONTIGUOUS(array)) {
+        npy_bool* p_start = (npy_bool*)PyArray_DATA(array);
+        npy_bool* p = p_start;
+        npy_bool* p_end = p + count_max;
+        npy_bool* p_end_roll = p_end - size_div.rem;
+
+        while (p < p_end_roll) {
+            if (*(npy_uint64*)p == 0) {
+                p += 8; // no true within this 8 byte roll region
+                continue;
+            }
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
+        }
+        while (p < p_end) {
+            if (*p) {indices[count++] = p - p_start;}
+            p++;
         }
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
     }
-    while (p < p_end) {
-        if (*p) {NONZERO_APPEND_INDEX;}
-        p++;
+    else {
+        npy_intp i = 0; // position within Boolean array
+        npy_intp i_end = count_max;
+        npy_intp i_end_roll = count_max - size_div.rem;
+        while (i < i_end_roll) {
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+        }
+        while (i < i_end) {
+            if (*(npy_bool*)PyArray_GETPTR1(array, i)) {indices[count++] = i;}
+            i++;
+        }
     }
     NPY_END_THREADS;
 
-    // npy_uint64 roll;
-    // while (p < p_end_roll) {
-    //     roll = *(npy_uint64*)p;
-    //     if (roll == 0) {
-    //         p += 8; // no true within this 8 byte roll region
-    //         continue;
-    //     }
-    //     // this order depends on byte order
-    //     if (roll & 0xFF) {NONZERO_APPEND_OFFSET(0);}
-    //     if (roll & 0xFF00) {NONZERO_APPEND_OFFSET(1);}
-    //     if (roll & 0xFF0000) {NONZERO_APPEND_OFFSET(2);}
-    //     if (roll & 0xFF000000) {NONZERO_APPEND_OFFSET(3);}
-    //     if (roll & 0xFF00000000) {NONZERO_APPEND_OFFSET(4);}
-    //     if (roll & 0xFF0000000000) {NONZERO_APPEND_OFFSET(5);}
-    //     if (roll & 0xFF000000000000) {NONZERO_APPEND_OFFSET(6);}
-    //     if (roll & 0xFF00000000000000) {NONZERO_APPEND_OFFSET(7);}
-    //     p += 8;
-    // }
-    // while (p < p_end) {
-    //     if (*p) {NONZERO_APPEND_OFFSET(0);}
-    //     p++;
-    // }
-
-
-    // while (p < p_end_roll) {
-    //     if (*(npy_uint64*)p == 0) {
-    //         p += 8; // no true within this roll region
-    //         continue;
-    //     }
-    //     if (AK_UNLIKELY(count + 8 >= capacity)) {
-    //         capacity <<= 1;
-    //         indices = (npy_int64*)realloc(indices, sizeof(npy_int64) * capacity);
-    //         if (indices == NULL) {
-    //             return NULL;
-    //         }
-    //     }
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    // }
-    // // at most three more indices remain
-    // if (AK_UNLIKELY(count + 7 >= capacity)) {
-    //     capacity <<= 1;
-    //     indices = (npy_int64*)realloc(indices, sizeof(npy_int64) * capacity);
-    //     if (indices == NULL) {
-    //         return NULL;
-    //     }
-    // }
-    // while (p < p_end) {
-    //     if (*p) {indices[count++] = p - p_start;}
-    //     p++;
-    // }
-
     npy_intp dims = {count};
     final = PyArray_SimpleNewFromData(1, &dims, NPY_INT64, (void*)indices);
     if (!final) {
@@ -3683,7 +3630,6 @@ AK_nonzero_1d(PyArrayObject* array) {
     PyArray_CLEARFLAGS((PyArrayObject*)final, NPY_ARRAY_WRITEABLE);
     return final;
 }
-#undef NONZERO_APPEND_INDEX
 
 static PyObject*
 nonzero_1d(PyObject *Py_UNUSED(m), PyObject *a) {
@@ -3697,10 +3643,6 @@ nonzero_1d(PyObject *Py_UNUSED(m), PyObject *a) {
         PyErr_SetString(PyExc_ValueError, "Array must be of type bool");
         return NULL;
     }
-    if (!PyArray_IS_C_CONTIGUOUS(array)) {
-        PyErr_SetString(PyExc_ValueError, "Array must be contiguous");
-        return NULL;
-    }
     return AK_nonzero_1d(array);
 }
 
@@ -5869,8 +5811,6 @@ typedef struct TriMapObject {
     Py_ssize_t src_len;
     Py_ssize_t dst_len;
     Py_ssize_t len;
-    // Py_ssize_t src_connected;
-    // Py_ssize_t dst_connected;
     bool is_many;
     bool finalized;
 
@@ -5932,8 +5872,6 @@ TriMap_init(PyObject *self, PyObject *args, PyObject *kwargs) {
     tm->is_many = false;
     tm->finalized = false;
     tm->len = 0;
-    // tm->src_connected = 0;
-    // tm->dst_connected = 0;
 
     // we create arrays, and also pre-extract pointers to array data for fast insertion; we keep the array for optimal summing routines
     npy_intp dims_src_len[] = {src_len};
@@ -6039,7 +5977,7 @@ TriMap_repr(TriMapObject *self) {
             is_finalized);
 }
 
-// Provide the integer positions connecting the `src` to the `dst`. If there is no match to `src` or `dst`, the unmatched position can be provided with -1. From each side, a connection is documented to the current `len`. Each time this is called `len` is incremented, indicating the inrease in position of th `final`. Return NULL on error
+// Provide the integer positions connecting the `src` to the `dst`. If there is no match to `src` or `dst`, the unmatched position can be provided with -1. From each side, a connection is documented to the current `len`. Each time this is called `len` is incremented, indicating the inrease in position in the `final`. Return NULL on error.
 
 // Inner function for calling from C; returns 0 on success, -1 on error. Exceptions will be set on error.
 static inline int
@@ -6062,7 +6000,6 @@ AK_TM_register_one(TriMapObject* tm, Py_ssize_t src_from, Py_ssize_t dst_from) {
         }
         tm->src_one[tm->src_one_count] = (TriMapOne){src_from, tm->len};
         tm->src_one_count += 1;
-        // tm->src_connected += 1;
     }
     if (dst_matched) {
         if (AK_UNLIKELY(tm->dst_one_count == tm->dst_one_capacity)) {
@@ -6076,7 +6013,6 @@ AK_TM_register_one(TriMapObject* tm, Py_ssize_t src_from, Py_ssize_t dst_from) {
         }
         tm->dst_one[tm->dst_one_count] = (TriMapOne){dst_from, tm->len};
         tm->dst_one_count += 1;
-        // tm->dst_connected += 1;
     }
     if (src_matched && dst_matched) {
         if (!tm->is_many) {
@@ -6141,14 +6077,12 @@ TriMap_register_unmatched_dst(TriMapObject *self) {
             return NULL;
         }
         // derive indices for unmatched locations, call each with register_one
-        // PyObject* nonzero = PyArray_Nonzero(dst_unmatched);
         PyArrayObject* indices = (PyArrayObject*)AK_nonzero_1d(dst_unmatched);
         if (indices == NULL) {
             Py_DECREF((PyObject*)dst_unmatched);
             return NULL;
         }
         // borrow ref to array in 1-element tuple
-        // PyArrayObject *indices = (PyArrayObject*)PyTuple_GET_ITEM(nonzero, 0);
         npy_int64 *index_data = (npy_int64 *)PyArray_DATA(indices);
         npy_intp index_len = PyArray_SIZE(indices);
 
@@ -6215,8 +6149,6 @@ TriMap_register_many(TriMapObject *self, PyObject *args) {
         npy_int64 pos = *(npy_int64*)PyArray_GETPTR1(dst_from, i); // always int64
         self->dst_match_data[pos] = NPY_TRUE;
     }
-    // self->src_connected += increment;
-    // self->dst_connected += increment;
     self->len += increment;
     self->is_many = true;
     Py_RETURN_NONE;
@@ -6359,7 +6291,7 @@ TriMap_dst_no_fill(TriMapObject *self, PyObject *Py_UNUSED(unused)) {
     Py_RETURN_FALSE;
 }
 
-# define TRANSFER_SCALARS(npy_type_to, npy_type_from) {                \
+# define AK_TM_TRANSFER_SCALAR(npy_type_to, npy_type_from) do {        \
     npy_type_to* array_to_data = (npy_type_to*)PyArray_DATA(array_to); \
     TriMapOne* o = one_pairs;                                          \
     TriMapOne* o_end = o + one_count;                                  \
@@ -6396,57 +6328,11 @@ TriMap_dst_no_fill(TriMapObject *self, PyObject *Py_UNUSED(unused)) {
             }                                                          \
         }                                                              \
     }                                                                  \
-}                                                                      \
-
-// The elsize of the `to` array will be equal to or greater than the from array
-#define TRANSFER_FLEXIBLE(c_type) {                                        \
-    npy_intp t_element_size = PyArray_DESCR(array_to)->elsize;             \
-    npy_intp t_element_cp = t_element_size / sizeof(c_type);               \
-    npy_intp f_element_size = PyArray_DESCR(array_from)->elsize;           \
-    npy_intp f_element_cp = f_element_size / sizeof(c_type);               \
-    npy_intp gap = t_element_size - f_element_size;                        \
-    c_type* array_to_data = (c_type*)PyArray_DATA(array_to);               \
-    c_type* f;                                                             \
-    c_type* t;                                                             \
-    c_type* t_end;                                                         \
-    npy_intp dst_pos;                                                      \
-    npy_int64 f_pos;                                                       \
-    PyArrayObject* dst;                                                    \
-    TriMapOne* o = one_pairs;                                              \
-    TriMapOne* o_end = o + one_count;                                      \
-    for (; o < o_end; o++) {                                               \
-        f = (c_type*)PyArray_GETPTR1(array_from, o->from);                 \
-        t = array_to_data + t_element_cp * o->to;                          \
-        memcpy(t, f, f_element_size);                                      \
-        memset(t + f_element_cp, '\0', gap);                               \
-    }                                                                      \
-    for (Py_ssize_t i = 0; i < tm->many_count; i++) {                      \
-        t = array_to_data + t_element_cp * tm->many_to[i].start;           \
-        t_end = array_to_data + t_element_cp * tm->many_to[i].stop;        \
-        if (from_src) {                                                    \
-            f = (c_type*)PyArray_GETPTR1(array_from, tm->many_from[i].src);\
-            for (; t < t_end; t += t_element_cp) {                         \
-                memcpy(t, f, f_element_size);                              \
-                memset(t + f_element_cp, '\0', gap);                       \
-            }                                                              \
-        }                                                                  \
-        else {                                                             \
-            dst_pos = 0;                                                   \
-            dst = tm->many_from[i].dst;                                    \
-            for (; t < t_end; t += t_element_cp) {                         \
-                f_pos = *(npy_int64*)PyArray_GETPTR1(dst, dst_pos);        \
-                f = (c_type*)PyArray_GETPTR1(array_from, f_pos);           \
-                memcpy(t, f, f_element_size);                              \
-                memset(t + f_element_cp, '\0', gap);                       \
-                dst_pos++;                                                 \
-            }                                                              \
-        }                                                                  \
-    }                                                                      \
-}                                                                          \
+} while (0)                                                            \
 
 // Based on `tm` state, transfer from src or from dst (depending on `from_src`) to a `array_to`, a newly created contiguous array that is compatible with the values in `array_from`. Returns -1 on error. This only needs to match to / from type combinations that are possible from `resolve_dtype`, i.e., bool never goes to integer.
 static inline int
-AK_TM_transfer(TriMapObject* tm,
+AK_TM_transfer_scalar(TriMapObject* tm,
         bool from_src,
         PyArrayObject* array_from,
         PyArrayObject* array_to) {
@@ -6455,206 +6341,194 @@ AK_TM_transfer(TriMapObject* tm,
 
     switch(PyArray_TYPE(array_to)){
         case NPY_BOOL:
-            TRANSFER_SCALARS(npy_bool, npy_bool);
+            AK_TM_TRANSFER_SCALAR(npy_bool, npy_bool);
             return 0;
         case NPY_INT64:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_INT64:
-                    TRANSFER_SCALARS(npy_int64, npy_int64);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_int64);
                     return 0;
                 case NPY_INT32:
-                    TRANSFER_SCALARS(npy_int64, npy_int32);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_int32);
                     return 0;
                 case NPY_INT16:
-                    TRANSFER_SCALARS(npy_int64, npy_int16);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_int16);
                     return 0;
                 case NPY_INT8:
-                    TRANSFER_SCALARS(npy_int64, npy_int8);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_int8);
                     return 0;
                 case NPY_UINT32:
-                    TRANSFER_SCALARS(npy_int64, npy_uint32);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_uint32);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_int64, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_int64, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_int64, npy_uint8);
                     return 0;
             }
             break;
         case NPY_INT32:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_INT32:
-                    TRANSFER_SCALARS(npy_int32, npy_int32);
+                    AK_TM_TRANSFER_SCALAR(npy_int32, npy_int32);
                     return 0;
                 case NPY_INT16:
-                    TRANSFER_SCALARS(npy_int32, npy_int16);
+                    AK_TM_TRANSFER_SCALAR(npy_int32, npy_int16);
                     return 0;
                 case NPY_INT8:
-                    TRANSFER_SCALARS(npy_int32, npy_int8);
+                    AK_TM_TRANSFER_SCALAR(npy_int32, npy_int8);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_int32, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_int32, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_int32, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_int32, npy_uint8);
                     return 0;
             }
             break;
         case NPY_INT16:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_INT16:
-                    TRANSFER_SCALARS(npy_int16, npy_int16);
+                    AK_TM_TRANSFER_SCALAR(npy_int16, npy_int16);
                     return 0;
                 case NPY_INT8:
-                    TRANSFER_SCALARS(npy_int16, npy_int8);
+                    AK_TM_TRANSFER_SCALAR(npy_int16, npy_int8);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_int16, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_int16, npy_uint8);
                     return 0;
             }
             break;
         case NPY_INT8:
-            TRANSFER_SCALARS(npy_int8, npy_int8);
+            AK_TM_TRANSFER_SCALAR(npy_int8, npy_int8);
             return 0;
         case NPY_UINT64:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_UINT64:
-                    TRANSFER_SCALARS(npy_uint64, npy_uint64);
+                    AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint64);
                     return 0;
                 case NPY_UINT32:
-                    TRANSFER_SCALARS(npy_uint64, npy_uint32);
+                    AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint32);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_uint64, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_uint64, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint8);
                     return 0;
             }
             break;
         case NPY_UINT32:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_UINT32:
-                    TRANSFER_SCALARS(npy_uint32, npy_uint32);
+                    AK_TM_TRANSFER_SCALAR(npy_uint32, npy_uint32);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_uint32, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_uint32, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_uint32, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_uint32, npy_uint8);
                     return 0;
             }
             break;
         case NPY_UINT16:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_uint16, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_uint16, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_uint16, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_uint16, npy_uint8);
                     return 0;
             }
             break;
         case NPY_UINT8:
-            TRANSFER_SCALARS(npy_uint8, npy_uint8);
+            AK_TM_TRANSFER_SCALAR(npy_uint8, npy_uint8);
             return 0;
         case NPY_FLOAT64:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_FLOAT64:
-                    TRANSFER_SCALARS(npy_float64, npy_float64);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_float64);
                     return 0;
                 case NPY_FLOAT32:
-                    TRANSFER_SCALARS(npy_float64, npy_float32);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_float32);
                     return 0;
                 case NPY_FLOAT16:
-                    TRANSFER_SCALARS(npy_float64, npy_float16);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_float16);
                     return 0;
                 case NPY_INT64:
-                    TRANSFER_SCALARS(npy_float64, npy_int64);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_int64);
                     return 0;
                 case NPY_INT32:
-                    TRANSFER_SCALARS(npy_float64, npy_int32);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_int32);
                     return 0;
                 case NPY_INT16:
-                    TRANSFER_SCALARS(npy_float64, npy_int16);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_int16);
                     return 0;
                 case NPY_INT8:
-                    TRANSFER_SCALARS(npy_float64, npy_int8);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_int8);
                     return 0;
                 case NPY_UINT64:
-                    TRANSFER_SCALARS(npy_float64, npy_uint64);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint64);
                     return 0;
                 case NPY_UINT32:
-                    TRANSFER_SCALARS(npy_float64, npy_uint32);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint32);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_float64, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_float64, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint8);
                     return 0;
             }
             break;
         case NPY_FLOAT32:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_FLOAT32:
-                    TRANSFER_SCALARS(npy_float32, npy_float32);
+                    AK_TM_TRANSFER_SCALAR(npy_float32, npy_float32);
                     return 0;
                 case NPY_FLOAT16:
-                    TRANSFER_SCALARS(npy_float32, npy_float16);
+                    AK_TM_TRANSFER_SCALAR(npy_float32, npy_float16);
                     return 0;
                 case NPY_INT16:
-                    TRANSFER_SCALARS(npy_float32, npy_int16);
+                    AK_TM_TRANSFER_SCALAR(npy_float32, npy_int16);
                     return 0;
                 case NPY_INT8:
-                    TRANSFER_SCALARS(npy_float32, npy_int8);
+                    AK_TM_TRANSFER_SCALAR(npy_float32, npy_int8);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_float32, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_float32, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_float32, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_float32, npy_uint8);
                     return 0;
             }
             break;
         case NPY_FLOAT16:
             switch (PyArray_TYPE(array_from)) {
                 case NPY_FLOAT16:
-                    TRANSFER_SCALARS(npy_float16, npy_float16);
+                    AK_TM_TRANSFER_SCALAR(npy_float16, npy_float16);
                     return 0;
                 case NPY_INT8:
-                    TRANSFER_SCALARS(npy_float16, npy_int8);
+                    AK_TM_TRANSFER_SCALAR(npy_float16, npy_int8);
                     return 0;
                 case NPY_UINT16:
-                    TRANSFER_SCALARS(npy_float16, npy_uint16);
+                    AK_TM_TRANSFER_SCALAR(npy_float16, npy_uint16);
                     return 0;
                 case NPY_UINT8:
-                    TRANSFER_SCALARS(npy_float16, npy_uint8);
+                    AK_TM_TRANSFER_SCALAR(npy_float16, npy_uint8);
                     return 0;
             }
             break;
-        case NPY_UNICODE: {
-            TRANSFER_FLEXIBLE(Py_UCS4);
-            return 0;
-        }
-        case NPY_STRING: {
-            TRANSFER_FLEXIBLE(char);
-            return 0;
-        }
         case NPY_DATETIME: {
-            TRANSFER_SCALARS(npy_int64, npy_int64);
+            AK_TM_TRANSFER_SCALAR(npy_int64, npy_int64);
             return 0;
         }
     }
-    // AK_DEBUG_MSG_OBJ("array_to", (PyObject*)array_to);
-    // AK_DEBUG_MSG_OBJ("array_from", (PyObject*)array_from);
     PyErr_SetString(PyExc_TypeError, "No handling for types");
     return -1;
 }
-
-#undef TRANSFER_SCALARS
-#undef TRANSFER_FLEXIBLE
+#undef AK_TM_TRANSFER_SCALAR
 
 // Returns -1 on error. Specialized transfer from any type of an array to an object array.
 static inline int
@@ -6740,12 +6614,10 @@ AK_TM_fill_object(TriMapObject* tm,
 
     PyArrayObject* final_fill = (PyArrayObject*)(from_src
             ? tm->final_src_fill : tm->final_dst_fill);
-
     PyObject** array_to_data = (PyObject**)PyArray_DATA(array_to);
     npy_int64* p = (npy_int64*)PyArray_DATA(final_fill);
     npy_int64* p_end = p + PyArray_SIZE(final_fill);
     PyObject** target;
-    // npy_int64 pos;
     while (p < p_end) {
         target = array_to_data + *p++;
         Py_INCREF(fill_value);
@@ -6754,29 +6626,117 @@ AK_TM_fill_object(TriMapObject* tm,
     return 0;
 }
 
-// TODO: AK_TM_fill_flexible
-// this manually inserts string
-// if (t_is_flexible) {
-//     // insert fill values
-//     Py_UCS4* t = (Py_UCS4*)PyArray_DATA(array_to);
-//     npy_intp t_cp = PyArray_DESCR(array_to)->elsize / UCS4_SIZE;
-//     Py_ssize_t len = PyUnicode_GET_LENGTH(fill_value) * UCS4_SIZE; // code points
-//     Py_ssize_t count = from_src ? tm->src_len : tm->dst_len;
-//     // NOTE: matches do not tell where a fill is needed
-//     npy_bool* d = from_src ? tm->src_match_data : tm->dst_match_data;
-//     npy_bool* d_end = d + count;
-//     while (d < d_end) {
-//         if (*d == NPY_FALSE) {
-//             if (PyUnicode_AsUCS4(fill_value, t, len, 0) == NULL) {
-//                 Py_DECREF((PyObject*)array_to);
-//                 return NULL;
-//             }
-//         }
-//         t += t_cp;
-//         d++;
-//     }
-// }
+#define AK_TM_TRANSFER_FLEXIBLE(c_type) do {                               \
+    Py_ssize_t one_count = from_src ? tm->src_one_count : tm->dst_one_count;\
+    TriMapOne* one_pairs = from_src ? tm->src_one : tm->dst_one;           \
+    npy_intp t_element_size = PyArray_DESCR(array_to)->elsize;             \
+    npy_intp t_element_cp = t_element_size / sizeof(c_type);               \
+    npy_intp f_element_size = PyArray_DESCR(array_from)->elsize;           \
+    c_type* array_to_data = (c_type*)PyArray_DATA(array_to);               \
+    c_type* f;                                                             \
+    c_type* t;                                                             \
+    c_type* t_end;                                                         \
+    npy_intp dst_pos;                                                      \
+    npy_int64 f_pos;                                                       \
+    PyArrayObject* dst;                                                    \
+    TriMapOne* o = one_pairs;                                              \
+    TriMapOne* o_end = o + one_count;                                      \
+    for (; o < o_end; o++) {                                               \
+        f = (c_type*)PyArray_GETPTR1(array_from, o->from);                 \
+        t = array_to_data + t_element_cp * o->to;                          \
+        memcpy(t, f, f_element_size);                                      \
+    }                                                                      \
+    for (Py_ssize_t i = 0; i < tm->many_count; i++) {                      \
+        t = array_to_data + t_element_cp * tm->many_to[i].start;           \
+        t_end = array_to_data + t_element_cp * tm->many_to[i].stop;        \
+        if (from_src) {                                                    \
+            f = (c_type*)PyArray_GETPTR1(array_from, tm->many_from[i].src);\
+            for (; t < t_end; t += t_element_cp) {                         \
+                memcpy(t, f, f_element_size);                              \
+            }                                                              \
+        }                                                                  \
+        else {                                                             \
+            dst_pos = 0;                                                   \
+            dst = tm->many_from[i].dst;                                    \
+            for (; t < t_end; t += t_element_cp) {                         \
+                f_pos = *(npy_int64*)PyArray_GETPTR1(dst, dst_pos);        \
+                f = (c_type*)PyArray_GETPTR1(array_from, f_pos);           \
+                memcpy(t, f, f_element_size);                              \
+                dst_pos++;                                                 \
+            }                                                              \
+        }                                                                  \
+    }                                                                      \
+} while (0)                                                                \
+
+// Returns -1 on error.
+static inline int
+AK_TM_fill_unicode(TriMapObject* tm,
+        bool from_src,
+        PyArrayObject* array_to,
+        PyObject* fill_value) {
+    PyArrayObject* final_fill = (PyArrayObject*)(from_src
+            ? tm->final_src_fill : tm->final_dst_fill);
+
+    Py_UCS4* array_to_data = (Py_UCS4*)PyArray_DATA(array_to);
+    // code points per element
+    npy_intp cp = PyArray_DESCR(array_to)->elsize / UCS4_SIZE;
+
+    bool decref_fill_value = false;
+    if (PyBytes_Check(fill_value)) {
+        fill_value = PyUnicode_FromEncodedObject(fill_value, "utf-8", NULL);
+        if (fill_value == NULL) {
+            return -1;
+        }
+        decref_fill_value = true;
+    }
+    else if (!PyUnicode_Check(fill_value)) {
+        return -1;
+    }
+    Py_ssize_t fill_cp = PyUnicode_GET_LENGTH(fill_value) * UCS4_SIZE; // code points
+    // p is the index position to fill
+    npy_int64* p = (npy_int64*)PyArray_DATA(final_fill);
+    npy_int64* p_end = p + PyArray_SIZE(final_fill);
+    Py_UCS4* target;
+    while (p < p_end) {
+        target = array_to_data + (*p * cp);
+        // disabling copying a null
+        if (PyUnicode_AsUCS4(fill_value, target, fill_cp, 0) == NULL) {
+            return -1;
+        }
+        p++;
+    }
+    if (decref_fill_value) {
+        Py_DECREF(fill_value);
+    }
+    return 0;
+}
+
+// Returns -1 on error.
+static inline int
+AK_TM_fill_string(TriMapObject* tm,
+        bool from_src,
+        PyArrayObject* array_to,
+        PyObject* fill_value) {
+    PyArrayObject* final_fill = (PyArrayObject*)(from_src
+            ? tm->final_src_fill : tm->final_dst_fill);
 
+    char* array_to_data = (char*)PyArray_DATA(array_to);
+    npy_intp cp = PyArray_DESCR(array_to)->elsize;
+    if (!PyBytes_Check(fill_value)) {
+        return -1;
+    }
+    Py_ssize_t fill_cp = PyBytes_GET_SIZE(fill_value);
+    const char* fill_data = PyBytes_AS_STRING(fill_value);
+    // p is the index position to fill
+    npy_int64* p = (npy_int64*)PyArray_DATA(final_fill);
+    npy_int64* p_end = p + PyArray_SIZE(final_fill);
+    char* target;
+    while (p < p_end) {
+        target = array_to_data + (*p++ * cp);
+        memcpy(target, fill_data, fill_cp);
+    }
+    return 0;
+}
 
 // Returns NULL on error.
 static inline PyObject*
@@ -6790,6 +6750,9 @@ AK_TM_map_no_fill(TriMapObject* tm,
     npy_intp dims[] = {tm->len};
     PyArrayObject* array_to;
     bool dtype_is_obj = PyArray_TYPE(array_from) == NPY_OBJECT;
+    bool dtype_is_unicode = PyArray_TYPE(array_from) == NPY_UNICODE;
+    bool dtype_is_string = PyArray_TYPE(array_from) == NPY_STRING;
+
     // create to array
     if (dtype_is_obj) { // initializes values to NULL
         array_to = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_OBJECT);
@@ -6810,8 +6773,14 @@ AK_TM_map_no_fill(TriMapObject* tm,
             return NULL;
         }
     }
+    else if (dtype_is_unicode) {
+        AK_TM_TRANSFER_FLEXIBLE(Py_UCS4);
+    }
+    else if (dtype_is_string) {
+        AK_TM_TRANSFER_FLEXIBLE(char);
+    }
     else {
-        if (AK_TM_transfer(tm, from_src, array_from, array_to)) {
+        if (AK_TM_transfer_scalar(tm, from_src, array_from, array_to)) {
             Py_DECREF((PyObject*)array_to);
             return NULL;
         }
@@ -6864,15 +6833,22 @@ AK_TM_map_fill(TriMapObject* tm,
     // passing a borrowed ref; returns a new ref
     PyArray_Descr* dtype = AK_resolve_dtype(PyArray_DESCR(array_from), fill_value_dtype);
     bool dtype_is_obj = dtype->type_num == NPY_OBJECT;
+    bool dtype_is_unicode = dtype->type_num == NPY_UNICODE;
+    bool dtype_is_string = dtype->type_num == NPY_STRING;
 
     npy_intp dims[] = {tm->len};
     PyArrayObject* array_to;
+
     if (dtype_is_obj) {
         Py_DECREF(dtype); // not needed
         // will initialize to NULL, not None
         array_to = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_OBJECT);
         Py_INCREF(array_from); // normalize refs when casting
     }
+    else if (dtype_is_unicode || dtype_is_string) {
+        array_to = (PyArrayObject*)PyArray_Zeros(1, dims, dtype, 0); // steals dtype ref
+        Py_INCREF(array_from); // normalize refs when casting
+    }
     else {
         array_to = (PyArrayObject*)PyArray_Empty(1, dims, dtype, 0); // steals dtype ref
         if (PyArray_TYPE(array_from) == NPY_DATETIME &&
@@ -6893,44 +6869,51 @@ AK_TM_map_fill(TriMapObject* tm,
         Py_DECREF((PyObject*)array_from);
         return NULL;
     }
+    // array_from, array_to inc refed and dec refed on error
     if (dtype_is_obj) {
         if (AK_TM_transfer_object(tm, from_src, array_from, array_to)) {
-            Py_DECREF((PyObject*)array_to);
-            Py_DECREF((PyObject*)array_from);
-            return NULL;
+            goto error;
         }
         if (AK_TM_fill_object(tm, from_src, array_to, fill_value)) {
-            Py_DECREF((PyObject*)array_to);
-            Py_DECREF((PyObject*)array_from);
-            return NULL;
+            goto error;
+        }
+    }
+    else if (dtype_is_unicode) {
+        AK_TM_TRANSFER_FLEXIBLE(Py_UCS4);
+        if (AK_TM_fill_unicode(tm, from_src, array_to, fill_value)) {
+            goto error;
+        }
+    }
+    else if (dtype_is_string) {
+        AK_TM_TRANSFER_FLEXIBLE(char);
+        if (AK_TM_fill_string(tm, from_src, array_to, fill_value)) {
+            goto error;
         }
     }
-    // TODO: add special hanldig for unicode/bytes
     else {
         // Most simple is to fill with scalar, then overwrite values as needed; for object and flexible dtypes this is not efficient; for object dtypes, this obbligates us to decref the filled value when assigning
         if (PyArray_FillWithScalar(array_to, fill_value)) { // -1 on error
-            Py_DECREF((PyObject*)array_to);
-            Py_DECREF((PyObject*)array_from);
-            return NULL;
+            goto error;
         }
-        if (AK_TM_transfer(tm, from_src, array_from, array_to)) {
-            Py_DECREF((PyObject*)array_to);
-            Py_DECREF((PyObject*)array_from);
-            return NULL;
+        if (AK_TM_transfer_scalar(tm, from_src, array_from, array_to)) {
+            goto error;
         }
     }
-
     Py_DECREF((PyObject*)array_from); // ref inc for this function
     PyArray_CLEARFLAGS(array_to, NPY_ARRAY_WRITEABLE);
     return (PyObject*)array_to;
+error:
+    Py_DECREF((PyObject*)array_to);
+    Py_DECREF((PyObject*)array_from);
+    return NULL;
 }
+#undef AK_TM_TRANSFER_FLEXIBLE
 
 static PyObject*
 TriMap_map_src_fill(TriMapObject *self, PyObject *args) {
     PyArrayObject* array_from;
     PyObject* fill_value;
     PyArray_Descr* fill_value_dtype;
-
     if (!PyArg_ParseTuple(args,
             "O!OO!:map_src_fill",
             &PyArray_Type, &array_from,
@@ -6952,7 +6935,6 @@ TriMap_map_dst_fill(TriMapObject *self, PyObject *args) {
     PyArrayObject* array_from;
     PyObject* fill_value;
     PyArray_Descr* fill_value_dtype;
-
     if (!PyArg_ParseTuple(args,
             "O!OO!:map_dst_fill",
             &PyArray_Type, &array_from,
@@ -6969,7 +6951,6 @@ TriMap_map_dst_fill(TriMapObject *self, PyObject *args) {
     return AK_TM_map_fill(self, from_src, array_from, fill_value, fill_value_dtype);
 }
 
-
 static PyMethodDef TriMap_methods[] = {
     {"register_one", (PyCFunction)TriMap_register_one, METH_VARARGS, NULL},
     {"register_unmatched_dst", (PyCFunction)TriMap_register_unmatched_dst, METH_NOARGS, NULL},

{arraykit-0.6.1 → arraykit-0.6.3}/test/test_nonzero_1d.py

@@ -83,4 +83,19 @@ class TestUnit(unittest.TestCase):
         a1[999] = True
         self.assertEqual(nonzero_1d(a1).tolist(), [999, 9_999_999])
         a1[0] = True
-        self.assertEqual(nonzero_1d(a1).tolist(), [0, 999, 9_999_999])
+        self.assertEqual(nonzero_1d(a1).tolist(), [0, 999, 9_999_999])
+
+    def test_nonzero_1d_f(self) -> None:
+        # non-contiguous
+        a1 = np.arange(40).reshape(10, 4) % 3 == 0
+        a2 = a1[:, 3]
+        self.assertEqual(nonzero_1d(a2).tolist(), [0, 3, 6, 9])
+
+        a3 = a1[:, 1]
+        self.assertEqual(nonzero_1d(a3).tolist(), [2, 5, 8])
+
+    def test_nonzero_1d_g(self) -> None:
+        a1 = np.arange(20).reshape(4, 5) % 3 == 0
+        a2 = a1[:, 4]
+        # array([False,  True, False, False])
+        self.assertEqual(nonzero_1d(a2).tolist(), [1])

{arraykit-0.6.1 → arraykit-0.6.3}/test/test_tri_map.py

@@ -975,6 +975,21 @@ class TestUnit(unittest.TestCase):
         post_dst = tm.map_dst_no_fill(dst)
         self.assertEqual(post_dst.tolist(), [b'a', b'bbb', b'cc', b'cc', b'dddd', b'a'])
 
+    def test_tri_map_map_bytes_a(self) -> None:
+        src = np.array([b'a', b'bbb', b'cc'], dtype=np.bytes_)
+        dst = np.array([b'cc', b'dddd', b'eee'], dtype=np.bytes_)
+
+        tm = TriMap(len(src), len(dst))
+        tm.register_one(0, -1)
+        tm.register_one(1, -1)
+        tm.register_one(2, 0)
+        tm.register_unmatched_dst()
+        tm.finalize()
+
+        post_src = tm.map_src_fill(src, b'--', np.dtype(np.bytes_))
+        post_dst = tm.map_dst_fill(dst, b'--', np.dtype(np.bytes_))
+        self.assertEqual(post_src.tolist(), [b'a', b'bbb', b'cc', b'--', b'--'])
+        self.assertEqual(post_dst.tolist(), [b'--', b'--', b'cc', b'dddd', b'eee'])
     #---------------------------------------------------------------------------
 
     def test_tri_map_map_unicode_a(self) -> None:
@@ -994,7 +1009,6 @@ class TestUnit(unittest.TestCase):
         post_dst = tm.map_dst_fill(dst, '====', np.array('====').dtype)
         self.assertEqual(post_dst.tolist(), ['a', 'a', 'a', '====', 'cc', 'cc', '===='])
 
-
     def test_tri_map_map_unicode_b(self) -> None:
         src = np.array(['a', 'bbb', 'cc', 'dddd'])
         dst = np.array(['cc', 'a', 'a', 'a', 'cc'])