Skip to content

Commit 3bc1474

Browse files
flexatoneflexatone
committed
first rendering
1 parent fbeefee commit 3bc1474

6 files changed

Lines changed: 608 additions & 0 deletions

File tree

‎src/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
from ._arraykit import split_after_count as split_after_count
2626
from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen
2727
from ._arraykit import write_array_to_file as write_array_to_file
28+
from ._arraykit import factorize as factorize
2829
from ._arraykit import count_iteration as count_iteration
2930
from ._arraykit import first_true_1d as first_true_1d
3031
from ._arraykit import first_true_2d as first_true_2d

‎src/__init__.pyi‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,9 @@ def write_array_to_file(
224224
fortran_order: bool = False,
225225
buffersize: int = 8192,
226226
) -> None: ...
227+
def factorize(
228+
array: np.ndarray, *, sort: bool = ...
229+
) -> tp.Tuple[np.ndarray, np.ndarray]: ...
227230
def first_true_1d(__array: np.ndarray, *, forward: bool) -> int: ...
228231
def first_true_2d(__array: np.ndarray, *, forward: bool, axis: int) -> np.ndarray: ...
229232
def nonzero_1d(__array: np.ndarray, /) -> np.ndarray: ...

‎src/_arraykit.c‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,10 @@ static PyMethodDef arraykit_methods[] = {
7070
(PyCFunction)write_array_to_file,
7171
METH_VARARGS | METH_KEYWORDS,
7272
NULL},
73+
{"factorize",
74+
(PyCFunction)factorize,
75+
METH_VARARGS | METH_KEYWORDS,
76+
NULL},
7377
{NULL},
7478
};
7579

‎src/auto_map.c‎

Lines changed: 324 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2520,6 +2520,330 @@ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
25202520
# undef INSERT_FLEXIBLE
25212521

25222522

2523+
//------------------------------------------------------------------------------
2524+
// factorize
2525+
2526+
// Detect a Python-level float NaN (a Python float or a NumPy floating scalar) so
2527+
// that object-dtype arrays collapse all NaN into a single code, matching the
2528+
// float-dtype behavior. Only real floats are treated as NaN here; None, NaT, and
2529+
// complex NaN remain ordinary distinct keys. Returns 1 for float NaN, else 0.
2530+
static inline int
2531+
factorize_obj_is_float_nan(PyObject* key)
2532+
{
2533+
if (PyFloat_Check(key)) {
2534+
return isnan(PyFloat_AS_DOUBLE(key));
2535+
}
2536+
if (PyArray_IsScalar(key, Half)) {
2537+
return npy_half_isnan(PyArrayScalar_VAL(key, Half));
2538+
}
2539+
if (PyArray_IsScalar(key, Float32)) {
2540+
return isnan(PyArrayScalar_VAL(key, Float32));
2541+
}
2542+
if (PyArray_IsScalar(key, Float64)) {
2543+
return isnan(PyArrayScalar_VAL(key, Float64));
2544+
}
2545+
return 0;
2546+
}
2547+
2548+
// Given a probe result `pos_expr` and its `hash_val` for the element at index
2549+
// `i`, either assign a new sequential code (empty slot) or reuse the code of the
2550+
// already-seen key (occupied slot). Stores the first-occurrence input index in
2551+
// the table so the reused lookup_hash_* probes compare against the right value.
2552+
// Requires `scratch`, `codes`, `code_of_index`, `first_index`, `k`, `i` in scope.
2553+
# define FACTORIZE_RECORD(pos_expr, hash_val) \
2554+
{ \
2555+
Py_ssize_t _pos = (pos_expr); \
2556+
if (_pos < 0) { \
2557+
goto fail; \
2558+
} \
2559+
if (scratch.table[_pos].hash == -1) { \
2560+
scratch.table[_pos].keys_pos = i; \
2561+
scratch.table[_pos].hash = (hash_val); \
2562+
first_index[k] = i; \
2563+
code_of_index[i] = k; \
2564+
codes[i] = k; \
2565+
k++; \
2566+
} \
2567+
else { \
2568+
codes[i] = code_of_index[scratch.table[_pos].keys_pos]; \
2569+
} \
2570+
} \
2571+
2572+
# define FACTORIZE_INT(npy_type, kat_lookup) \
2573+
{ \
2574+
for (npy_intp i = 0; i < n; i++) { \
2575+
npy_int64 v = (npy_int64)*(npy_type*)PyArray_GETPTR1(a, i); \
2576+
Py_hash_t hash = int_to_hash(v); \
2577+
FACTORIZE_RECORD(lookup_hash_int(&scratch, v, hash, kat_lookup), hash); \
2578+
} \
2579+
} \
2580+
2581+
# define FACTORIZE_UINT(npy_type, kat_lookup) \
2582+
{ \
2583+
for (npy_intp i = 0; i < n; i++) { \
2584+
npy_uint64 v = (npy_uint64)*(npy_type*)PyArray_GETPTR1(a, i); \
2585+
Py_hash_t hash = uint_to_hash(v); \
2586+
FACTORIZE_RECORD(lookup_hash_uint(&scratch, v, hash, kat_lookup), hash); \
2587+
} \
2588+
} \
2589+
2590+
// Floats need explicit NaN handling: all NaN collapse to one code and never
2591+
// enter the table (the map compares with `==`, so NaN would otherwise never
2592+
// match itself). +0.0/-0.0 and inf are handled correctly by the normal path.
2593+
# define FACTORIZE_FLOAT(npy_type, kat_lookup, post_deref) \
2594+
{ \
2595+
for (npy_intp i = 0; i < n; i++) { \
2596+
npy_double v = post_deref(*(npy_type*)PyArray_GETPTR1(a, i)); \
2597+
if (v != v) { \
2598+
if (nan_code < 0) { \
2599+
nan_code = k; \
2600+
first_index[k] = i; \
2601+
codes[i] = k; \
2602+
k++; \
2603+
} \
2604+
else { \
2605+
codes[i] = nan_code; \
2606+
} \
2607+
continue; \
2608+
} \
2609+
Py_hash_t hash = double_to_hash(v); \
2610+
FACTORIZE_RECORD(lookup_hash_double(&scratch, v, hash, kat_lookup), hash); \
2611+
} \
2612+
} \
2613+
2614+
# define FACTORIZE_FLEXIBLE(char_type, lookup_func, hash_func, get_end_func, dt_size_expr) \
2615+
{ \
2616+
Py_ssize_t dt_size = (dt_size_expr); \
2617+
for (npy_intp i = 0; i < n; i++) { \
2618+
char_type* v = (char_type*)PyArray_GETPTR1(a, i); \
2619+
Py_ssize_t ksize = get_end_func(v, dt_size) - v; \
2620+
Py_hash_t hash = hash_func(v, ksize); \
2621+
FACTORIZE_RECORD(lookup_func(&scratch, v, ksize, hash), hash); \
2622+
} \
2623+
} \
2624+
2625+
// Hash-based factorize: return (uniques, codes) such that
2626+
// array[i] == uniques[codes[i]], in O(n), reusing the AutoMap hash table.
2627+
PyObject *
2628+
factorize(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
2629+
{
2630+
static char *kwlist[] = {"array", "sort", NULL};
2631+
PyObject *array_obj = NULL;
2632+
int sort = 0;
2633+
if (!PyArg_ParseTupleAndKeywords(args, kwargs,
2634+
"O|$p:factorize", kwlist,
2635+
&array_obj,
2636+
&sort)) {
2637+
return NULL;
2638+
}
2639+
if (!PyArray_Check(array_obj)) {
2640+
PyErr_Format(PyExc_TypeError,
2641+
"Expected a NumPy array, not %s.", Py_TYPE(array_obj)->tp_name);
2642+
return NULL;
2643+
}
2644+
PyArrayObject *a = (PyArrayObject *)array_obj;
2645+
if (PyArray_NDIM(a) != 1) {
2646+
PyErr_SetString(PyExc_TypeError, "Array must be 1-dimensional");
2647+
return NULL;
2648+
}
2649+
2650+
npy_intp n = PyArray_SIZE(a);
2651+
int array_t = PyArray_TYPE(a);
2652+
KeysArrayType kat = at_to_kat(array_t, a);
2653+
2654+
PyObject *codes_arr = NULL;
2655+
PyObject *idx_arr = NULL;
2656+
PyObject *uniques = NULL;
2657+
PyObject *list = NULL; // object/KAT_LIST path only
2658+
npy_intp *code_of_index = NULL; // input index -> code
2659+
npy_intp *first_index = NULL; // code -> first-occurrence input index
2660+
npy_intp *rank = NULL; // sort remap
2661+
Py_ssize_t k = 0; // running distinct count
2662+
Py_ssize_t nan_code = -1; // shared code for all NaN, -1 until seen
2663+
2664+
// A private, stack-allocated scratch table; never exposed to Python. Only
2665+
// the fields read by grow_table / lookup_hash_* are set.
2666+
FAMObject scratch;
2667+
scratch.table = NULL;
2668+
scratch.table_size = 0;
2669+
scratch.keys = NULL;
2670+
scratch.keys_array_type = kat;
2671+
scratch.keys_size = n;
2672+
scratch.key_buffer = NULL; // never needed by the array lookups
2673+
2674+
codes_arr = PyArray_EMPTY(1, &n, NPY_INTP, 0);
2675+
if (!codes_arr) {
2676+
goto fail;
2677+
}
2678+
npy_intp *codes = (npy_intp*)PyArray_DATA((PyArrayObject*)codes_arr);
2679+
2680+
if (n > 0) {
2681+
code_of_index = PyMem_New(npy_intp, n);
2682+
first_index = PyMem_New(npy_intp, n);
2683+
if (!code_of_index || !first_index) {
2684+
PyErr_NoMemory();
2685+
goto fail;
2686+
}
2687+
}
2688+
2689+
// The lookups dereference scratch.keys at stored input indices. For a usable
2690+
// array KAT that is the (borrowed) input array; otherwise build a list.
2691+
if (kat) {
2692+
scratch.keys = array_obj;
2693+
}
2694+
else {
2695+
if (array_t == NPY_DATETIME || array_t == NPY_TIMEDELTA) {
2696+
list = PySequence_List(array_obj);
2697+
}
2698+
else {
2699+
list = PyArray_ToList(a);
2700+
}
2701+
if (!list) {
2702+
goto fail;
2703+
}
2704+
scratch.keys = list;
2705+
}
2706+
2707+
if (grow_table(&scratch, n)) {
2708+
goto fail;
2709+
}
2710+
2711+
switch (kat) {
2712+
case KAT_INT8: FACTORIZE_INT(npy_int8, KAT_INT8); break;
2713+
case KAT_INT16: FACTORIZE_INT(npy_int16, KAT_INT16); break;
2714+
case KAT_INT32: FACTORIZE_INT(npy_int32, KAT_INT32); break;
2715+
case KAT_INT64: FACTORIZE_INT(npy_int64, KAT_INT64); break;
2716+
case KAT_UINT8: FACTORIZE_UINT(npy_uint8, KAT_UINT8); break;
2717+
case KAT_UINT16: FACTORIZE_UINT(npy_uint16, KAT_UINT16); break;
2718+
case KAT_UINT32: FACTORIZE_UINT(npy_uint32, KAT_UINT32); break;
2719+
case KAT_UINT64: FACTORIZE_UINT(npy_uint64, KAT_UINT64); break;
2720+
case KAT_FLOAT16: FACTORIZE_FLOAT(npy_half, KAT_FLOAT16, npy_half_to_double); break;
2721+
case KAT_FLOAT32: FACTORIZE_FLOAT(npy_float, KAT_FLOAT32, ); break;
2722+
case KAT_FLOAT64: FACTORIZE_FLOAT(npy_double, KAT_FLOAT64, ); break;
2723+
case KAT_UNICODE:
2724+
FACTORIZE_FLEXIBLE(Py_UCS4, lookup_hash_unicode, unicode_to_hash,
2725+
ucs4_get_end_p, PyArray_ITEMSIZE(a) / UCS4_SIZE);
2726+
break;
2727+
case KAT_STRING:
2728+
FACTORIZE_FLEXIBLE(char, lookup_hash_string, string_to_hash,
2729+
char_get_end_p, PyArray_ITEMSIZE(a));
2730+
break;
2731+
case KAT_DTY:
2732+
case KAT_DTM:
2733+
case KAT_DTW:
2734+
case KAT_DTD:
2735+
case KAT_DTh:
2736+
case KAT_DTm:
2737+
case KAT_DTs:
2738+
case KAT_DTms:
2739+
case KAT_DTus:
2740+
case KAT_DTns:
2741+
case KAT_DTps:
2742+
case KAT_DTfs:
2743+
case KAT_DTas:
2744+
// datetime64/timedelta64 store an int64; NaT (INT64_MIN) compares
2745+
// equal to itself here, so all NaT collapse into one code naturally.
2746+
FACTORIZE_INT(npy_int64, KAT_INT64);
2747+
break;
2748+
default: { // KAT_LIST: object dtype, complex, dt64 without a unit
2749+
for (npy_intp i = 0; i < n; i++) {
2750+
PyObject* key = PyList_GET_ITEM(list, i); // borrowed
2751+
if (factorize_obj_is_float_nan(key)) {
2752+
if (nan_code < 0) {
2753+
nan_code = k;
2754+
first_index[k] = i;
2755+
codes[i] = k;
2756+
k++;
2757+
}
2758+
else {
2759+
codes[i] = nan_code;
2760+
}
2761+
continue;
2762+
}
2763+
Py_hash_t hash = PyObject_Hash(key);
2764+
if (hash == -1) {
2765+
goto fail;
2766+
}
2767+
FACTORIZE_RECORD(lookup_hash_obj(&scratch, key, hash), hash);
2768+
}
2769+
break;
2770+
}
2771+
}
2772+
2773+
// uniques: take the first-occurrence values out of the input, same dtype.
2774+
idx_arr = PyArray_EMPTY(1, &k, NPY_INTP, 0);
2775+
if (!idx_arr) {
2776+
goto fail;
2777+
}
2778+
if (k > 0) {
2779+
memcpy(PyArray_DATA((PyArrayObject*)idx_arr), first_index, k * sizeof(npy_intp));
2780+
}
2781+
uniques = PyArray_TakeFrom(a, idx_arr, 0, NULL, NPY_RAISE);
2782+
if (!uniques) {
2783+
goto fail;
2784+
}
2785+
2786+
if (sort && k > 1) {
2787+
PyObject* order = PyArray_ArgSort((PyArrayObject*)uniques, 0, NPY_QUICKSORT);
2788+
if (!order) {
2789+
goto fail;
2790+
}
2791+
npy_intp* order_data = (npy_intp*)PyArray_DATA((PyArrayObject*)order);
2792+
rank = PyMem_New(npy_intp, k);
2793+
if (!rank) {
2794+
Py_DECREF(order);
2795+
PyErr_NoMemory();
2796+
goto fail;
2797+
}
2798+
for (npy_intp j = 0; j < k; j++) {
2799+
rank[order_data[j]] = j;
2800+
}
2801+
for (npy_intp i = 0; i < n; i++) {
2802+
codes[i] = rank[codes[i]];
2803+
}
2804+
PyObject* uniques_sorted = PyArray_TakeFrom(
2805+
(PyArrayObject*)uniques, order, 0, NULL, NPY_RAISE);
2806+
Py_DECREF(order);
2807+
if (!uniques_sorted) {
2808+
goto fail;
2809+
}
2810+
Py_SETREF(uniques, uniques_sorted);
2811+
}
2812+
2813+
PyArray_CLEARFLAGS((PyArrayObject*)codes_arr, NPY_ARRAY_WRITEABLE);
2814+
PyArray_CLEARFLAGS((PyArrayObject*)uniques, NPY_ARRAY_WRITEABLE);
2815+
2816+
PyMem_Free(scratch.table);
2817+
PyMem_Free(code_of_index);
2818+
PyMem_Free(first_index);
2819+
PyMem_Free(rank);
2820+
Py_DECREF(idx_arr);
2821+
Py_XDECREF(list);
2822+
2823+
PyObject* result = PyTuple_Pack(2, uniques, codes_arr);
2824+
Py_DECREF(uniques);
2825+
Py_DECREF(codes_arr);
2826+
return result;
2827+
2828+
fail:
2829+
PyMem_Free(scratch.table);
2830+
PyMem_Free(code_of_index);
2831+
PyMem_Free(first_index);
2832+
PyMem_Free(rank);
2833+
Py_XDECREF(idx_arr);
2834+
Py_XDECREF(list);
2835+
Py_XDECREF(codes_arr);
2836+
Py_XDECREF(uniques);
2837+
return NULL;
2838+
}
2839+
2840+
# undef FACTORIZE_RECORD
2841+
# undef FACTORIZE_INT
2842+
# undef FACTORIZE_UINT
2843+
# undef FACTORIZE_FLOAT
2844+
# undef FACTORIZE_FLEXIBLE
2845+
2846+
25232847
static PyObject *
25242848
fam_repr(FAMObject *self)
25252849
{

‎src/auto_map.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,5 +10,7 @@ extern PyTypeObject FAMVType;
1010
extern PyTypeObject FAMType;
1111
extern PyObject *NonUniqueError;
1212

13+
PyObject *factorize(PyObject *m, PyObject *args, PyObject *kwargs);
14+
1315

1416
# endif /* ARRAYKIT_SRC_AUTO_MAP_H_ */

0 commit comments

Comments
 (0)