map_object()

README.rst

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

src/__init__.py

@@ -28,6 +28,8 @@
 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 prepare_iter_for_array as prepare_iter_for_array
 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,10 @@ 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 prepare_iter_for_array(
+    values: tp.Iterable[tp.Any], copy: bool = ...
+) -> tp.Tuple[tp.Optional[type], bool, tp.Sequence[tp.Any]]: ...
 def fill_directional(
     array: np.ndarray,
     target: np.ndarray,

src/_arraykit.c

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

src/methods.c

@@ -1403,6 +1403,294 @@ 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
+
+// running state of the prepare_iter_for_array dtype inference over a series of values
+typedef struct AK_InferState {
+    int has_str;
+    int has_non_str;
+    int has_inexact;
+    int has_big_int;
+    int has_tuple;    // a sized object was seen (tuple/list/array/...)
+    int needs_object; // the resolved dtype must be object rather than auto-detected
+} AK_InferState;
+
+// Classify one value per static_frame.core.util.prepare_iter_for_array's rules, updating
+// `s`. Once `s->needs_object` is set the caller can stop inspecting. `enum_type` is
+// enum.Enum (may be NULL to skip the Enum check).
+static inline void
+AK_infer_value(PyObject *v, PyObject *enum_type, AK_InferState *s)
+{
+    PyTypeObject *vt = Py_TYPE(v);
+    // exact str/bytes (Python or numpy scalar) -> string; subclasses fall through
+    if (PyUnicode_CheckExact(v) || PyBytes_CheckExact(v)
+            || PyArray_IsScalar(v, Unicode) || PyArray_IsScalar(v, String)) {
+        s->has_str = 1;
+    }
+    // a sized object (tuple, list, array, SF container, str subclass) -> object
+    else if ((vt->tp_as_sequence && vt->tp_as_sequence->sq_length)
+            || (vt->tp_as_mapping && vt->tp_as_mapping->mp_length)) {
+        s->has_tuple = 1;
+        s->needs_object = 1;
+    }
+    else {
+        s->has_non_str = 1;
+        if (vt == &PyFloat_Type || vt == &PyComplex_Type) {
+            s->has_inexact = 1;
+        }
+        else if (vt == &PyLong_Type) {
+            int overflow = 0;
+            long long lv = PyLong_AsLongLongAndOverflow(v, &overflow);
+            if (overflow || llabs(lv) > AK_INT_MAX_COERCIBLE_TO_FLOAT) {
+                s->has_big_int = 1;
+            }
+        }
+        else if (PyArray_IsScalar(v, Generic)) {
+            ; // any other numpy scalar: non-str, no inexact/big-int, not an Enum
+        }
+        else if (enum_type != NULL && PyObject_IsInstance(v, enum_type) == 1) {
+            s->needs_object = 1;
+        }
+    }
+    if ((s->has_str && s->has_non_str) || (s->has_big_int && s->has_inexact)) {
+        s->needs_object = 1;
+    }
+}
+
+// Import enum.Enum for the Enum inference case; returns a new reference or NULL
+// (clearing the error), in which case the Enum check is skipped.
+static PyObject *
+AK_import_enum(void)
+{
+    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();
+    }
+    return enum_type;
+}
+
+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 = AK_import_enum();
+    AK_InferState state = {0, 0, 0, 0, 0, 0};
+
+    // 1D: hoist the base pointer and element stride and walk a running pointer, rather
+    // than recomputing PyArray_GETPTR1 each iteration. For a contiguous array the stride
+    // is the itemsize (direct indexing into the flat buffer); a strided slice still works.
+    char *p = (char*)PyArray_DATA(array);
+    npy_intp stride = PyArray_STRIDES(array)[0];
+
+    for (npy_intp i = 0; i < n; i++, p += stride) {
+        PyObject *elem;
+        if (is_object) {
+            elem = *(PyObject**)p;
+            Py_INCREF(elem);
+        }
+        else {
+            elem = PyArray_ToScalar(p, 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 (!state.needs_object) {
+            AK_infer_value(r, enum_type, &state);
+        }
+    }
+    Py_XDECREF(enum_type);
+
+    PyObject *result;
+    if (state.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;
+}
+
+static char *prepare_iter_for_array_kwarg_names[] = {
+    "values",
+    "copy",
+    NULL
+};
+
+// Infer a dtype specifier for the elements of an iterable, matching
+// static_frame.core.util.prepare_iter_for_array: return ``(resolved, has_tuple, values)``
+// where `resolved` is None (let NumPy auto-detect) or the ``object`` type, `has_tuple`
+// marks that a sized object was seen, and `values` is a newly materialized list when
+// `copy` is true (e.g. a generator/dict/set) else the original iterable. The caller
+// decides `copy` (via is_gen_copy_values), keeping the SF-specific type policy out of C.
+PyObject *
+prepare_iter_for_array(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
+{
+    PyObject *values = NULL;
+    int copy = 0;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+            "O|p:prepare_iter_for_array",
+            prepare_iter_for_array_kwarg_names,
+            &values,
+            &copy
+            )) {
+        return NULL;
+    }
+    PyObject *enum_type = AK_import_enum();
+    AK_InferState state = {0, 0, 0, 0, 0, 0};
+    PyObject *values_out = NULL; // new list when copy, else a new ref to the original
+
+    if (copy) {
+        // materialize into a list while inspecting (generator/dict/set input). Pre-size
+        // the list from a length hint when one is available (sets, dicts, sized
+        // iterators) to avoid repeated reallocation; fall back to append growth for a
+        // bare generator (hint 0), and guard against an inexact hint over/under-shooting.
+        Py_ssize_t hint = PyObject_LengthHint(values, 0);
+        if (hint < 0) {
+            goto fail;
+        }
+        values_out = PyList_New(hint);
+        if (values_out == NULL) {
+            goto fail;
+        }
+        PyObject *iter = PyObject_GetIter(values);
+        if (iter == NULL) {
+            goto fail;
+        }
+        Py_ssize_t i = 0;
+        PyObject *item;
+        while ((item = PyIter_Next(iter)) != NULL) {
+            if (!state.needs_object) {
+                AK_infer_value(item, enum_type, &state);
+            }
+            if (i < hint) {
+                PyList_SET_ITEM(values_out, i, item); // steals reference
+            }
+            else { // hint underestimated the length
+                int rc = PyList_Append(values_out, item);
+                Py_DECREF(item);
+                if (rc != 0) {
+                    Py_DECREF(iter);
+                    goto fail;
+                }
+            }
+            i++;
+        }
+        Py_DECREF(iter);
+        if (PyErr_Occurred()) {
+            goto fail;
+        }
+        if (i < hint) { // hint overestimated: drop the trailing (NULL) slots
+            Py_SET_SIZE(values_out, i);
+        }
+    }
+    else {
+        // inspect only; direct indexing for list/tuple, else the iterator protocol
+        Py_INCREF(values);
+        values_out = values;
+        if (PyList_CheckExact(values) || PyTuple_CheckExact(values)) {
+            Py_ssize_t sz = PySequence_Fast_GET_SIZE(values);
+            for (Py_ssize_t i = 0; i < sz && !state.needs_object; i++) {
+                AK_infer_value(PySequence_Fast_GET_ITEM(values, i), enum_type, &state);
+            }
+        }
+        else {
+            PyObject *iter = PyObject_GetIter(values);
+            if (iter == NULL) {
+                goto fail;
+            }
+            PyObject *item;
+            while (!state.needs_object && (item = PyIter_Next(iter)) != NULL) {
+                AK_infer_value(item, enum_type, &state);
+                Py_DECREF(item);
+            }
+            Py_DECREF(iter);
+            if (PyErr_Occurred()) {
+                goto fail;
+            }
+        }
+    }
+    Py_XDECREF(enum_type);
+
+    PyObject *resolved = state.needs_object
+            ? (PyObject*)&PyBaseObject_Type // the ``object`` builtin -> object dtype
+            : Py_None;
+    PyObject *has_tuple = state.has_tuple ? Py_True : Py_False;
+    PyObject *result = PyTuple_Pack(3, resolved, has_tuple, values_out);
+    Py_DECREF(values_out);
+    return result;
+
+fail:
+    Py_XDECREF(enum_type);
+    Py_XDECREF(values_out);
+    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,12 @@ 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 *
+prepare_iter_for_array(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
+
 PyObject *
 fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);