first rendering

Author: flexatone

README.rst

@@ -35,6 +35,12 @@ ArrayKit requires the following:
 What is New in ArrayKit
 -------------------------
 
+1.12.0
+............
+
+Added ``map_object()``.
+
+
 1.11.0
 ............
 

src/__init__.py

@@ -28,6 +28,7 @@
 from ._arraykit import factorize as factorize
 from ._arraykit import group_ordering as group_ordering
 from ._arraykit import group_reduce as group_reduce
+from ._arraykit import map_object as map_object
 from ._arraykit import fill_directional as fill_directional
 from ._arraykit import count_iteration as count_iteration
 from ._arraykit import first_true_1d as first_true_1d

src/__init__.pyi

@@ -236,6 +236,7 @@ def group_ordering(
 def group_reduce(
     codes: np.ndarray, size: int, values: np.ndarray, op: str
 ) -> np.ndarray: ...
+def map_object(array: np.ndarray, func: tp.Callable[[tp.Any], tp.Any]) -> np.ndarray: ...
 def fill_directional(
     array: np.ndarray,
     target: np.ndarray,

src/_arraykit.c

@@ -82,6 +82,10 @@ static PyMethodDef arraykit_methods[] =  {
             (PyCFunction)group_reduce,
             METH_VARARGS | METH_KEYWORDS,
             NULL},
+    {"map_object",
+            (PyCFunction)map_object,
+            METH_VARARGS | METH_KEYWORDS,
+            NULL},
     {"fill_directional",
             (PyCFunction)fill_directional,
             METH_VARARGS | METH_KEYWORDS,

src/methods.c

@@ -1403,6 +1403,149 @@ group_reduce(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
     return out_arr;
 }
 
+// int magnitude beyond which a Python int is no longer losslessly coercible to float;
+// mirrors static_frame.core.util.INT_MAX_COERCIBLE_TO_FLOAT
+#define AK_INT_MAX_COERCIBLE_TO_FLOAT 1000000000000000LL
+
+static char *map_object_kwarg_names[] = {
+    "array",
+    "func",
+    NULL
+};
+
+// Apply a Python callable to each element of a 1D array (elements boxed as numpy scalars,
+// matching NumPy/Series iteration) and return a new 1D array, inferring the result dtype
+// with the same rules as static_frame.core.util.prepare_iter_for_array: the result is an
+// object array when the applied values mix strings and non-strings, include a sized object
+// (tuple/list/array), an Enum, or mix a large Python int with a Python float/complex;
+// otherwise NumPy auto-detects the dtype (e.g. str -> '<U', float -> float64). This fuses
+// the per-element apply, the type inspection, and the array build into one C pass.
+PyObject *
+map_object(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
+{
+    PyArrayObject *array = NULL;
+    PyObject *func = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+            "O!O:map_object",
+            map_object_kwarg_names,
+            &PyArray_Type, &array,
+            &func
+            )) {
+        return NULL;
+    }
+    if (PyArray_NDIM(array) != 1) {
+        PyErr_SetString(PyExc_ValueError, "array must be 1-dimensional");
+        return NULL;
+    }
+    if (!PyCallable_Check(func)) {
+        PyErr_SetString(PyExc_TypeError, "func must be callable");
+        return NULL;
+    }
+    npy_intp n = PyArray_SIZE(array);
+    int is_object = PyArray_TYPE(array) == NPY_OBJECT;
+
+    PyObject *values = PyList_New(n); // collected results; owns references
+    if (values == NULL) {
+        return NULL;
+    }
+    // enum.Enum for the rare Enum-result case; on failure proceed without the check
+    PyObject *enum_type = NULL;
+    PyObject *enum_mod = PyImport_ImportModule("enum");
+    if (enum_mod != NULL) {
+        enum_type = PyObject_GetAttrString(enum_mod, "Enum");
+        Py_DECREF(enum_mod);
+    }
+    if (enum_type == NULL) {
+        PyErr_Clear();
+    }
+
+    // prepare_iter_for_array inference state
+    int has_str = 0;
+    int has_non_str = 0;
+    int has_inexact = 0;
+    int has_big_int = 0;
+    int needs_object = 0;
+
+    for (npy_intp i = 0; i < n; i++) {
+        PyObject *elem;
+        if (is_object) {
+            elem = *(PyObject**)PyArray_GETPTR1(array, i);
+            Py_INCREF(elem);
+        }
+        else {
+            elem = PyArray_ToScalar(PyArray_GETPTR1(array, i), array);
+            if (elem == NULL) {
+                goto fail;
+            }
+        }
+        PyObject *r = PyObject_CallOneArg(func, elem);
+        Py_DECREF(elem);
+        if (r == NULL) {
+            goto fail;
+        }
+        PyList_SET_ITEM(values, i, r); // steals reference to r
+
+        if (needs_object) {
+            continue; // dtype already resolved; keep collecting only
+        }
+        PyTypeObject *rt = Py_TYPE(r);
+        // exact str/bytes (Python or numpy scalar) -> string; subclasses fall through
+        if (PyUnicode_CheckExact(r) || PyBytes_CheckExact(r)
+                || PyArray_IsScalar(r, Unicode) || PyArray_IsScalar(r, String)) {
+            has_str = 1;
+        }
+        // a sized object (tuple, list, array, SF container, str subclass) -> object
+        else if ((rt->tp_as_sequence && rt->tp_as_sequence->sq_length)
+                || (rt->tp_as_mapping && rt->tp_as_mapping->mp_length)) {
+            needs_object = 1;
+        }
+        else {
+            has_non_str = 1;
+            if (rt == &PyFloat_Type || rt == &PyComplex_Type) {
+                has_inexact = 1;
+            }
+            else if (rt == &PyLong_Type) {
+                int overflow = 0;
+                long long lv = PyLong_AsLongLongAndOverflow(r, &overflow);
+                if (overflow || llabs(lv) > AK_INT_MAX_COERCIBLE_TO_FLOAT) {
+                    has_big_int = 1;
+                }
+            }
+            else if (PyArray_IsScalar(r, Generic)) {
+                ; // any other numpy scalar: non-str, no inexact/big-int, not an Enum
+            }
+            else if (enum_type != NULL && PyObject_IsInstance(r, enum_type) == 1) {
+                needs_object = 1;
+            }
+        }
+        if ((has_str && has_non_str) || (has_big_int && has_inexact)) {
+            needs_object = 1;
+        }
+    }
+    Py_XDECREF(enum_type);
+
+    PyObject *result;
+    if (needs_object) {
+        // build an object array of the collected values
+        result = PyArray_FROM_OTF(values, NPY_OBJECT, NPY_ARRAY_C_CONTIGUOUS);
+    }
+    else {
+        // let NumPy auto-detect the dtype from the values (str -> '<U', float -> f8, ...)
+        result = PyArray_FromAny(values, NULL, 1, 1, NPY_ARRAY_C_CONTIGUOUS, NULL);
+    }
+    Py_DECREF(values);
+    if (result == NULL) {
+        return NULL;
+    }
+    PyArray_CLEARFLAGS((PyArrayObject*)result, NPY_ARRAY_WRITEABLE);
+    return result;
+
+fail:
+    Py_XDECREF(enum_type);
+    Py_DECREF(values);
+    return NULL;
+}
+
 // Fill one strided lane in place: walk positions in the fill direction, carrying
 // the most recent non-target value into each target position (subject to `limit`
 // consecutive fills per run). `elem_base`/`elem_stride` address elements in bytes;

src/methods.h

@@ -75,6 +75,9 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
 PyObject *
 group_reduce(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
 
+PyObject *
+map_object(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
+
 PyObject *
 fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
 

test/test_map_object.py

@@ -0,0 +1,139 @@
+import unittest
+from enum import Enum
+
+import numpy as np
+from arraykit import map_object
+
+
+class Color(Enum):
+    R = 1
+    G = 2
+
+
+# reference: prepare_iter_for_array's inference + build, as static_frame applies it
+_INEXACT = (float, complex, np.inexact)
+_BIG = 1_000_000_000_000_000
+
+
+def _reference(arr, func):
+    vals = [func(v) for v in arr]
+    resolved = None
+    has_str = has_non = has_inx = has_big = False
+    for v in vals:
+        vt = v.__class__
+        if vt is str or vt is np.str_ or vt is bytes or vt is np.bytes_:
+            has_str = True
+        elif hasattr(v, '__len__') or isinstance(v, Enum):
+            resolved = object
+            break
+        else:
+            has_non = True
+            if vt in _INEXACT:
+                has_inx = True
+            elif vt is int and abs(v) > _BIG:
+                has_big = True
+        if (has_str and has_non) or (has_big and has_inx):
+            resolved = object
+            break
+    return np.array(vals) if resolved is None else np.array(vals, dtype=object)
+
+
+class TestUnit(unittest.TestCase):
+    def _check(self, arr, func):
+        post = map_object(arr, func)
+        exp = _reference(arr, func)
+        self.assertEqual(post.dtype, exp.dtype, (arr.dtype, exp.dtype))
+        self.assertTrue(np.array_equal(post, exp))
+        self.assertFalse(post.flags.writeable)
+        return post
+
+    def test_map_object_str_from_float(self) -> None:
+        post = self._check(np.array([1.5, 2.25, 3.0]), lambda x: str(x))
+        self.assertEqual(post.dtype.kind, 'U')
+
+    def test_map_object_str_from_bool(self) -> None:
+        post = self._check(np.array([True, False, True]), lambda x: str(x))
+        self.assertEqual(post.tolist(), ['True', 'False', 'True'])
+        self.assertEqual(post.dtype.kind, 'U')
+
+    def test_map_object_str_from_int(self) -> None:
+        self._check(np.array([1, 2, 3], dtype=np.int64), lambda x: str(x))
+
+    def test_map_object_native_float(self) -> None:
+        post = self._check(np.array([1.5, 2.5]), lambda x: float(x) * 2)
+        self.assertEqual(post.dtype, np.dtype(np.float64))
+
+    def test_map_object_native_int(self) -> None:
+        post = self._check(np.array([1, 2, 3]), lambda x: int(x) + 1)
+        self.assertEqual(post.dtype, np.dtype(np.int64))
+
+    def test_map_object_tuple_result(self) -> None:
+        post = self._check(np.array([1, 2]), lambda x: (int(x), int(x)))
+        self.assertEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_list_result(self) -> None:
+        post = self._check(np.array([1, 2]), lambda x: [int(x)])
+        self.assertEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_mixed_str_nonstr(self) -> None:
+        post = self._check(np.array([1, 2, 3]), lambda x: str(x) if x > 1 else int(x))
+        self.assertEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_python_float(self) -> None:
+        self._check(np.array([1, 2, 3]), lambda x: 1.5)
+
+    def test_map_object_bigint_and_inexact(self) -> None:
+        # a large python int mixed with a python float -> object
+        post = self._check(np.array([1, 2]), lambda x: 10**18 if x == 1 else 1.5)
+        self.assertEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_bigint_only(self) -> None:
+        # big ints alone (no inexact) do not force object
+        post = self._check(np.array([1, 2]), lambda x: 10**18)
+        self.assertNotEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_enum_result(self) -> None:
+        post = self._check(np.array([1, 2]), lambda x: Color.R)
+        self.assertEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_object_input(self) -> None:
+        arr = np.array(['a', 'bb', 'ccc'], dtype=object)
+        post = self._check(arr, lambda x: len(x))
+        self.assertEqual(post.tolist(), [1, 2, 3])
+
+    def test_map_object_receives_numpy_scalar(self) -> None:
+        # elements are boxed as numpy scalars, matching Series/array iteration
+        seen = []
+        map_object(np.array([1.5, 2.5]), lambda x: seen.append(type(x)) or x)
+        self.assertTrue(all(t is np.float64 for t in seen))
+
+    def test_map_object_str_subclass_is_object(self) -> None:
+        # a str subclass is not an exact str -> sized object -> object array
+        class S(str):
+            pass
+
+        post = map_object(np.array([1, 2]), lambda x: S(str(x)))
+        self.assertEqual(post.dtype, np.dtype(object))
+
+    def test_map_object_empty(self) -> None:
+        post = self._check(np.array([], dtype=np.float64), lambda x: str(x))
+        self.assertEqual(len(post), 0)
+
+    def test_map_object_propagates_exception(self) -> None:
+        def bad(x):
+            raise ValueError('boom')
+
+        with self.assertRaises(ValueError):
+            map_object(np.array([1, 2]), bad)
+
+    def test_map_object_errors(self) -> None:
+        with self.assertRaises(ValueError):  # 2d
+            map_object(np.array([[1, 2]]), lambda x: x)
+        with self.assertRaises(TypeError):  # not callable
+            map_object(np.array([1, 2]), 3)
+        with self.assertRaises(TypeError):  # not an array
+            map_object([1, 2], lambda x: x)
+
+
+if __name__ == '__main__':
+    unittest.main()