arraykit 1.2.0__cp314-cp314-win32.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
arraykit/methods.c ADDED
@@ -0,0 +1,1135 @@
1
+ # include "Python.h"
2
+
3
+ # define NO_IMPORT_ARRAY
4
+ # define PY_ARRAY_UNIQUE_SYMBOL AK_ARRAY_API
5
+ # define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
6
+
7
+ # include "numpy/arrayobject.h"
8
+ # include "numpy/arrayscalars.h"
9
+ # include "numpy/halffloat.h"
10
+
11
+ # include "methods.h"
12
+ # include "utilities.h"
13
+
14
+ PyObject *
15
+ count_iteration(PyObject *Py_UNUSED(m), PyObject *iterable)
16
+ {
17
+ PyObject *iter = PyObject_GetIter(iterable);
18
+ if (iter == NULL) return NULL;
19
+
20
+ int count = 0;
21
+ PyObject *v;
22
+
23
+ while ((v = PyIter_Next(iter))) {
24
+ count++;
25
+ Py_DECREF(v);
26
+ }
27
+ Py_DECREF(iter);
28
+ if (PyErr_Occurred()) {
29
+ return NULL;
30
+ }
31
+ PyObject* result = PyLong_FromLong(count);
32
+ if (result == NULL) return NULL;
33
+ return result;
34
+ }
35
+
36
+ PyObject *
37
+ row_1d_filter(PyObject *Py_UNUSED(m), PyObject *a)
38
+ {
39
+ AK_CHECK_NUMPY_ARRAY_1D_2D(a);
40
+ PyArrayObject *array = (PyArrayObject *)a;
41
+
42
+ if (PyArray_NDIM(array) == 2) {
43
+ npy_intp dim[1] = {PyArray_DIM(array, 1)};
44
+ PyArray_Dims shape = {dim, 1};
45
+ // NOTE: this will set PyErr if shape is not compatible
46
+ return PyArray_Newshape(array, &shape, NPY_ANYORDER);
47
+ }
48
+ Py_INCREF(a);
49
+ return a;
50
+ }
51
+
52
+ PyObject *
53
+ slice_to_ascending_slice(PyObject *Py_UNUSED(m), PyObject *args) {
54
+
55
+ PyObject* slice;
56
+ PyObject* size;
57
+ if (!PyArg_ParseTuple(args,
58
+ "O!O!:slice_to_ascending_slice",
59
+ &PySlice_Type, &slice,
60
+ &PyLong_Type, &size)) {
61
+ return NULL;
62
+ }
63
+ // will delegate NULL on eroror
64
+ return AK_slice_to_ascending_slice(slice, PyLong_AsSsize_t(size));
65
+ }
66
+
67
+ PyObject *
68
+ column_2d_filter(PyObject *Py_UNUSED(m), PyObject *a)
69
+ {
70
+ AK_CHECK_NUMPY_ARRAY_1D_2D(a);
71
+ PyArrayObject *array = (PyArrayObject *)a;
72
+
73
+ if (PyArray_NDIM(array) == 1) {
74
+ // https://numpy.org/doc/stable/reference/c-api/types-and-structures.html#c.PyArray_Dims
75
+ npy_intp dim[2] = {PyArray_DIM(array, 0), 1};
76
+ PyArray_Dims shape = {dim, 2};
77
+ // PyArray_Newshape might return NULL and set PyErr, so no handling to do here
78
+ return PyArray_Newshape(array, &shape, NPY_ANYORDER); // already a PyObject*
79
+ }
80
+ Py_INCREF(a); // returning borrowed ref, must increment
81
+ return a;
82
+ }
83
+
84
+ PyObject *
85
+ column_1d_filter(PyObject *Py_UNUSED(m), PyObject *a)
86
+ {
87
+ AK_CHECK_NUMPY_ARRAY_1D_2D(a);
88
+ PyArrayObject *array = (PyArrayObject *)a;
89
+
90
+ if (PyArray_NDIM(array) == 2) {
91
+ npy_intp dim[1] = {PyArray_DIM(array, 0)};
92
+ PyArray_Dims shape = {dim, 1};
93
+ // NOTE: this will set PyErr if shape is not compatible
94
+ return PyArray_Newshape(array, &shape, NPY_ANYORDER);
95
+ }
96
+ Py_INCREF(a);
97
+ return a;
98
+ }
99
+
100
+ PyObject *
101
+ shape_filter(PyObject *Py_UNUSED(m), PyObject *a) {
102
+ AK_CHECK_NUMPY_ARRAY_1D_2D(a);
103
+ PyArrayObject *array = (PyArrayObject *)a;
104
+ npy_intp rows = PyArray_DIM(array, 0);
105
+ // If 1D array, set size for axis 1 at 1, else use 2D array to get the size of axis 1
106
+ npy_intp cols = PyArray_NDIM(array) == 1 ? 1 : PyArray_DIM(array, 1);
107
+ return AK_build_pair_ssize_t(rows, cols);
108
+ }
109
+
110
+ PyObject *
111
+ name_filter(PyObject *Py_UNUSED(m), PyObject *n) {
112
+ if (AK_UNLIKELY(PyObject_Hash(n) == -1)) {
113
+ return PyErr_Format(PyExc_TypeError,
114
+ "unhashable name (type '%s')",
115
+ Py_TYPE(n)->tp_name);
116
+ }
117
+ Py_INCREF(n);
118
+ return n;
119
+ }
120
+
121
+ PyObject *
122
+ mloc(PyObject *Py_UNUSED(m), PyObject *a)
123
+ {
124
+ AK_CHECK_NUMPY_ARRAY(a);
125
+ return PyLong_FromVoidPtr(PyArray_DATA((PyArrayObject *)a));
126
+ }
127
+
128
+ PyObject *
129
+ immutable_filter(PyObject *Py_UNUSED(m), PyObject *a) {
130
+ AK_CHECK_NUMPY_ARRAY(a);
131
+ return (PyObject *)AK_immutable_filter((PyArrayObject *)a);
132
+ }
133
+
134
+ PyObject *
135
+ resolve_dtype(PyObject *Py_UNUSED(m), PyObject *args)
136
+ {
137
+ PyArray_Descr *d1, *d2;
138
+ if (!PyArg_ParseTuple(args,
139
+ "O!O!:resolve_dtype",
140
+ &PyArrayDescr_Type, &d1,
141
+ &PyArrayDescr_Type, &d2)) {
142
+ return NULL;
143
+ }
144
+ return (PyObject *)AK_resolve_dtype(d1, d2);
145
+ }
146
+
147
+ PyObject *
148
+ resolve_dtype_iter(PyObject *Py_UNUSED(m), PyObject *arg) {
149
+ PyObject *iterator = PyObject_GetIter(arg);
150
+ if (iterator == NULL) {
151
+ // No need to set exception here. GetIter already sets TypeError
152
+ return NULL;
153
+ }
154
+ PyArray_Descr *resolved = NULL;
155
+ PyArray_Descr *dtype;
156
+ while ((dtype = (PyArray_Descr*) PyIter_Next(iterator))) {
157
+ if (!PyArray_DescrCheck(dtype)) {
158
+ PyErr_Format(
159
+ PyExc_TypeError, "argument must be an iterable over %s, not %s",
160
+ ((PyTypeObject *) &PyArrayDescr_Type)->tp_name,
161
+ Py_TYPE(dtype)->tp_name
162
+ );
163
+ Py_DECREF(iterator);
164
+ Py_DECREF(dtype);
165
+ Py_XDECREF(resolved);
166
+ return NULL;
167
+ }
168
+ if (!resolved) {
169
+ resolved = dtype;
170
+ continue;
171
+ }
172
+ Py_SETREF(resolved, AK_resolve_dtype(resolved, dtype));
173
+ Py_DECREF(dtype);
174
+ if (!resolved || PyDataType_ISOBJECT(resolved)) {
175
+ break;
176
+ }
177
+ }
178
+ Py_DECREF(iterator);
179
+ if (PyErr_Occurred()) {
180
+ return NULL;
181
+ }
182
+ if (!resolved) {
183
+ // this could happen if this function gets an empty tuple
184
+ PyErr_SetString(PyExc_ValueError, "iterable passed to resolve dtypes is empty");
185
+ }
186
+ return (PyObject *)resolved;
187
+ }
188
+
189
+ PyObject *
190
+ nonzero_1d(PyObject *Py_UNUSED(m), PyObject *a) {
191
+ AK_CHECK_NUMPY_ARRAY(a);
192
+ PyArrayObject* array = (PyArrayObject*)a;
193
+ if (PyArray_NDIM(array) != 1) {
194
+ PyErr_SetString(PyExc_ValueError, "Array must be 1-dimensional");
195
+ return NULL;
196
+ }
197
+ if (PyArray_TYPE(array) != NPY_BOOL) {
198
+ PyErr_SetString(PyExc_ValueError, "Array must be of type bool");
199
+ return NULL;
200
+ }
201
+ return AK_nonzero_1d(array);
202
+ }
203
+
204
+ PyObject*
205
+ is_objectable_dt64(PyObject *m, PyObject *a) {
206
+ AK_CHECK_NUMPY_ARRAY(a);
207
+ PyArrayObject* array = (PyArrayObject*)a;
208
+
209
+ // this returns a new reference
210
+ PyObject* dt_year = PyObject_GetAttrString(m, "dt_year");
211
+ int is_objectable = AK_is_objectable_dt64(array, dt_year);
212
+ Py_DECREF(dt_year);
213
+
214
+ switch (is_objectable) {
215
+ case -1:
216
+ return NULL;
217
+ case 0:
218
+ Py_RETURN_FALSE;
219
+ case 1:
220
+ Py_RETURN_TRUE;
221
+ }
222
+ return NULL;
223
+ }
224
+
225
+
226
+ PyObject*
227
+ is_objectable(PyObject *m, PyObject *a) {
228
+ AK_CHECK_NUMPY_ARRAY(a);
229
+ PyArrayObject* array = (PyArrayObject*)a;
230
+
231
+ char kind = PyArray_DESCR(array)->kind;
232
+ if ((kind == 'M' || kind == 'm')) {
233
+ // this returns a new reference
234
+ PyObject* dt_year = PyObject_GetAttrString(m, "dt_year");
235
+ int is_objectable = AK_is_objectable_dt64(array, dt_year);
236
+ Py_DECREF(dt_year);
237
+
238
+ switch (is_objectable) {
239
+ case -1:
240
+ return NULL;
241
+ case 0:
242
+ Py_RETURN_FALSE;
243
+ case 1:
244
+ Py_RETURN_TRUE;
245
+ }
246
+ }
247
+ Py_RETURN_TRUE;
248
+ }
249
+
250
+ // Convert array to the dtype provided. NOTE: mutable arrays will be returned unless the input array is immutable and no dtype change is needed
251
+ PyObject*
252
+ astype_array(PyObject* m, PyObject* args) {
253
+
254
+ PyObject* a = NULL;
255
+ PyObject* dtype_spec = Py_None;
256
+
257
+ if (!PyArg_ParseTuple(args, "O!|O:astype_array",
258
+ &PyArray_Type, &a,
259
+ &dtype_spec)) {
260
+ return NULL;
261
+ }
262
+ PyArrayObject* array = (PyArrayObject*)a;
263
+
264
+ PyArray_Descr* dtype = NULL;
265
+ if (dtype_spec == Py_None) {
266
+ dtype = PyArray_DescrFromType(NPY_DEFAULT_TYPE);
267
+ } else {
268
+ if (!PyArray_DescrConverter(dtype_spec, &dtype)) {
269
+ return NULL;
270
+ }
271
+ }
272
+
273
+ if (PyArray_EquivTypes(PyArray_DESCR(array), dtype)) {
274
+ Py_DECREF(dtype);
275
+
276
+ if (PyArray_ISWRITEABLE(array)) {
277
+ PyObject* result = PyArray_NewCopy(array, NPY_ANYORDER);
278
+ if (!result) {
279
+ return NULL;
280
+ }
281
+ return result;
282
+ }
283
+ else { // already immutable
284
+ Py_INCREF(a);
285
+ return a;
286
+ }
287
+ }
288
+ // if converting to an object
289
+ if (dtype->type_num == NPY_OBJECT) {
290
+ char kind = PyArray_DESCR(array)->kind;
291
+ if ((kind == 'M' || kind == 'm')) {
292
+ PyObject* dt_year = PyObject_GetAttrString(m, "dt_year");
293
+ int is_objectable = AK_is_objectable_dt64(array, dt_year);
294
+ Py_DECREF(dt_year);
295
+
296
+ if (!is_objectable) {
297
+ PyObject* result = PyArray_NewLikeArray(array, NPY_ANYORDER, dtype, 0);
298
+ if (!result) {
299
+ Py_DECREF(dtype);
300
+ return NULL;
301
+ }
302
+ PyObject** data = (PyObject**)PyArray_DATA((PyArrayObject*)result);
303
+
304
+ PyArrayIterObject* it = (PyArrayIterObject*)PyArray_IterNew(a);
305
+ if (!it) {
306
+ Py_DECREF(result);
307
+ return NULL;
308
+ }
309
+
310
+ npy_intp i = 0;
311
+ while (it->index < it->size) {
312
+ PyObject* item = PyArray_ToScalar(it->dataptr, array);
313
+ if (!item) {
314
+ Py_DECREF(result);
315
+ Py_DECREF(it);
316
+ return NULL;
317
+ }
318
+ data[i++] = item;
319
+ PyArray_ITER_NEXT(it);
320
+ }
321
+ Py_DECREF(it);
322
+ return result;
323
+ }
324
+ }
325
+ }
326
+ // all other cases: do a standard cast conversion
327
+ PyObject* result = PyArray_CastToType(array, dtype, 0);
328
+ if (!result) {
329
+ Py_DECREF(dtype);
330
+ return NULL;
331
+ }
332
+ return result;
333
+ }
334
+
335
+
336
+
337
+ static char *first_true_1d_kwarg_names[] = {
338
+ "array",
339
+ "forward",
340
+ NULL
341
+ };
342
+
343
+ PyObject *
344
+ first_true_1d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
345
+ {
346
+ PyArrayObject *array = NULL;
347
+ int forward = 1;
348
+
349
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
350
+ "O!|$p:first_true_1d",
351
+ first_true_1d_kwarg_names,
352
+ &PyArray_Type, &array,
353
+ &forward
354
+ )) {
355
+ return NULL;
356
+ }
357
+ if (PyArray_NDIM(array) != 1) {
358
+ PyErr_SetString(PyExc_ValueError, "Array must be 1-dimensional");
359
+ return NULL;
360
+ }
361
+ if (PyArray_TYPE(array) != NPY_BOOL) {
362
+ PyErr_SetString(PyExc_ValueError, "Array must be of type bool");
363
+ return NULL;
364
+ }
365
+ if (!PyArray_IS_C_CONTIGUOUS(array)) {
366
+ PyErr_SetString(PyExc_ValueError, "Array must be contiguous");
367
+ return NULL;
368
+ }
369
+
370
+ npy_intp lookahead = sizeof(npy_uint64);
371
+ npy_intp size = PyArray_SIZE(array);
372
+ lldiv_t size_div = lldiv((long long)size, lookahead); // quot, rem
373
+
374
+ npy_bool *array_buffer = (npy_bool*)PyArray_DATA(array);
375
+
376
+ NPY_BEGIN_THREADS_DEF;
377
+ NPY_BEGIN_THREADS;
378
+
379
+ Py_ssize_t position = -1;
380
+ npy_bool *p;
381
+ npy_bool *p_end;
382
+ npy_bool *p_end_roll;
383
+
384
+ if (forward) {
385
+ p = array_buffer;
386
+ p_end = p + size;
387
+ p_end_roll = p_end - size_div.rem;
388
+
389
+ while (p < p_end_roll) {
390
+ if (*(npy_uint64*)p != 0) {
391
+ break; // found a true within lookahead
392
+ }
393
+ p += lookahead;
394
+ }
395
+ while (p < p_end) {
396
+ if (*p) break;
397
+ p++;
398
+ }
399
+ }
400
+ else {
401
+ p = array_buffer + size - 1;
402
+ p_end = array_buffer - 1;
403
+ p_end_roll = p_end + size_div.rem;
404
+
405
+ while (p > p_end_roll) {
406
+ if (*(npy_uint64*)(p - lookahead + 1) != 0) {
407
+ break; // found a true within lookahead
408
+ }
409
+ p -= lookahead;
410
+ }
411
+ while (p > p_end) {
412
+ if (*p) break;
413
+ p--;
414
+ }
415
+ }
416
+ if (p != p_end) { // else, return -1
417
+ position = p - array_buffer;
418
+ }
419
+ NPY_END_THREADS;
420
+
421
+ PyObject* post = PyLong_FromSsize_t(position);
422
+ return post;
423
+ }
424
+
425
+ static char *first_true_2d_kwarg_names[] = {
426
+ "array",
427
+ "forward",
428
+ "axis",
429
+ NULL
430
+ };
431
+
432
+ PyObject *
433
+ first_true_2d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
434
+ {
435
+ PyArrayObject *array = NULL;
436
+ int forward = 1;
437
+ int axis = 0;
438
+
439
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
440
+ "O!|$pi:first_true_2d",
441
+ first_true_2d_kwarg_names,
442
+ &PyArray_Type,
443
+ &array,
444
+ &forward,
445
+ &axis
446
+ )) {
447
+ return NULL;
448
+ }
449
+ if (PyArray_NDIM(array) != 2) {
450
+ PyErr_SetString(PyExc_ValueError, "Array must be 2-dimensional");
451
+ return NULL;
452
+ }
453
+ if (PyArray_TYPE(array) != NPY_BOOL) {
454
+ PyErr_SetString(PyExc_ValueError, "Array must be of type bool");
455
+ return NULL;
456
+ }
457
+ if (axis < 0 || axis > 1) {
458
+ PyErr_SetString(PyExc_ValueError, "Axis must be 0 or 1");
459
+ return NULL;
460
+ }
461
+
462
+ // NOTE: we copy the entire array into contiguous memory when necessary.
463
+ // axis = 0 returns the pos per col
464
+ // axis = 1 returns the pos per row (as contiguous bytes)
465
+ // if c contiguous:
466
+ // axis == 0: transpose, copy to C
467
+ // axis == 1: keep
468
+ // if f contiguous:
469
+ // axis == 0: transpose, keep
470
+ // axis == 1: copy to C
471
+ // else
472
+ // axis == 0: transpose, copy to C
473
+ // axis == 1: copy to C
474
+
475
+ bool transpose = !axis; // if 1, false
476
+ bool corder = true;
477
+ if ((PyArray_IS_C_CONTIGUOUS(array) && axis == 1) ||
478
+ (PyArray_IS_F_CONTIGUOUS(array) && axis == 0)) {
479
+ corder = false;
480
+ }
481
+ // create pointer to "indicator" array; if newly allocated, it will need to be decrefed before function termination
482
+ PyArrayObject *array_ind = NULL;
483
+ bool decref_array_ind = false;
484
+
485
+ if (transpose && !corder) {
486
+ array_ind = (PyArrayObject *)PyArray_Transpose(array, NULL);
487
+ if (array_ind == NULL) return NULL;
488
+ decref_array_ind = true;
489
+ }
490
+ else if (!transpose && corder) {
491
+ array_ind = (PyArrayObject *)PyArray_NewCopy(array, NPY_CORDER);
492
+ if (array_ind == NULL) return NULL;
493
+ decref_array_ind = true;
494
+ }
495
+ else if (transpose && corder) {
496
+ PyArrayObject *tmp = (PyArrayObject *)PyArray_Transpose(array, NULL);
497
+ if (tmp == NULL) return NULL;
498
+
499
+ array_ind = (PyArrayObject *)PyArray_NewCopy(tmp, NPY_CORDER);
500
+ Py_DECREF((PyObject*)tmp);
501
+ if (array_ind == NULL) return NULL;
502
+ decref_array_ind = true;
503
+ }
504
+ else {
505
+ array_ind = array; // can use array, no decref needed
506
+ }
507
+
508
+ npy_intp lookahead = sizeof(npy_uint64);
509
+
510
+ // buffer of indicators
511
+ npy_bool *buffer_ind = (npy_bool*)PyArray_DATA(array_ind);
512
+
513
+ npy_intp count_row = PyArray_DIM(array_ind, 0);
514
+ npy_intp count_col = PyArray_DIM(array_ind, 1);
515
+
516
+ lldiv_t div_col = lldiv((long long)count_col, lookahead); // quot, rem
517
+
518
+ npy_intp dims_post = {count_row};
519
+ PyArrayObject *array_pos = (PyArrayObject*)PyArray_EMPTY(
520
+ 1, // ndim
521
+ &dims_post,// shape
522
+ NPY_INT64, // dtype
523
+ 0 // fortran
524
+ );
525
+ if (array_pos == NULL) {
526
+ return NULL;
527
+ }
528
+ npy_int64 *buffer_pos = (npy_int64*)PyArray_DATA(array_pos);
529
+
530
+ NPY_BEGIN_THREADS_DEF;
531
+ NPY_BEGIN_THREADS;
532
+
533
+ npy_intp position;
534
+ npy_bool *p;
535
+ npy_bool *p_start;
536
+ npy_bool *p_end;
537
+
538
+ // iterate one row at a time; short-circult when found
539
+ // for axis 1 rows are rows; for axis 0, rows are (post transpose) columns
540
+ for (npy_intp r = 0; r < count_row; r++) {
541
+ position = -1; // update for each row
542
+
543
+ if (forward) {
544
+ // get start of each row
545
+ p_start = buffer_ind + (count_col * r);
546
+ p = p_start;
547
+ p_end = p + count_col; // end of each row
548
+
549
+ // scan each row from the front and terminate when True
550
+ // remove from the end the remainder
551
+ while (p < p_end - div_col.rem) {
552
+ if (*(npy_uint64*)p != 0) {
553
+ break; // found a true
554
+ }
555
+ p += lookahead;
556
+ }
557
+ while (p < p_end) {
558
+ if (*p) {break;}
559
+ p++;
560
+ }
561
+ if (p != p_end) {
562
+ position = p - p_start;
563
+ }
564
+ }
565
+ else { // reverse
566
+ // start at the next row, then subtract one for last elem in previous row
567
+ p_start = buffer_ind + (count_col * (r + 1)) - 1;
568
+ p = p_start;
569
+ // end is 1 less than start of each row
570
+ p_end = buffer_ind + (count_col * r) - 1;
571
+
572
+ while (p > p_end + div_col.rem) {
573
+ // must go to start of lookahead
574
+ if (*(npy_uint64*)(p - lookahead + 1) != 0) {
575
+ break; // found a true
576
+ }
577
+ p -= lookahead;
578
+ }
579
+ while (p > p_end) {
580
+ if (*p) {break;}
581
+ p--;
582
+ }
583
+ if (p != p_end) {
584
+ position = p - (p_end + 1);
585
+ }
586
+ }
587
+ *buffer_pos++ = position;
588
+ }
589
+
590
+ NPY_END_THREADS;
591
+
592
+ if (decref_array_ind) {
593
+ Py_DECREF(array_ind); // created in this function
594
+ }
595
+ return (PyObject *)array_pos;
596
+ }
597
+
598
+ PyObject *
599
+ dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg)
600
+ {
601
+ // -------------------------------------------------------------------------
602
+ // 1. Handle fast, exact type checks first.
603
+ if (arg == Py_None) {
604
+ return (PyObject*)PyArray_DescrFromType(NPY_OBJECT);
605
+ }
606
+ if (PyFloat_CheckExact(arg)) {
607
+ return (PyObject*)PyArray_DescrFromType(NPY_FLOAT64);
608
+ }
609
+ if (PyLong_CheckExact(arg)) {
610
+ return (PyObject*)PyArray_DescrFromType(NPY_INT64);
611
+ }
612
+ if (PyBool_Check(arg)) {
613
+ return (PyObject*)PyArray_DescrFromType(NPY_BOOL);
614
+ }
615
+
616
+ PyObject* dtype = NULL;
617
+ // String
618
+ if (PyUnicode_CheckExact(arg)) {
619
+ PyArray_Descr* descr = PyArray_DescrFromType(NPY_UNICODE);
620
+ if (descr == NULL) return NULL;
621
+ dtype = (PyObject*)PyArray_DescrFromObject(arg, descr);
622
+ Py_DECREF(descr);
623
+ return dtype;
624
+ }
625
+ // Bytes
626
+ if (PyBytes_CheckExact(arg)) {
627
+ PyArray_Descr* descr = PyArray_DescrFromType(NPY_STRING);
628
+ if (descr == NULL) return NULL;
629
+ dtype = (PyObject*)PyArray_DescrFromObject(arg, descr);
630
+ Py_DECREF(descr);
631
+ return dtype;
632
+ }
633
+
634
+ // -------------------------------------------------------------------------
635
+ // 2. Construct dtype (slightly more complicated)
636
+ // Already known
637
+ dtype = PyObject_GetAttrString(arg, "dtype");
638
+ if (dtype) {
639
+ return dtype;
640
+ }
641
+ PyErr_Clear();
642
+ // -------------------------------------------------------------------------
643
+ // 3. Handles everything else.
644
+ return (PyObject*)PyArray_DescrFromType(NPY_OBJECT);
645
+ }
646
+
647
+ static char *isna_element_kwarg_names[] = {
648
+ "element",
649
+ "include_none",
650
+ NULL
651
+ };
652
+
653
+ PyObject *
654
+ isna_element(PyObject *m, PyObject *args, PyObject *kwargs)
655
+ {
656
+ PyObject *element;
657
+ int include_none = 1;
658
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
659
+ "O|p:isna_element", isna_element_kwarg_names,
660
+ &element,
661
+ &include_none)) {
662
+ return NULL;
663
+ }
664
+
665
+ // None
666
+ if (include_none && element == Py_None) {
667
+ Py_RETURN_TRUE;
668
+ }
669
+
670
+ // NaN
671
+ if (PyFloat_Check(element)) {
672
+ return PyBool_FromLong(isnan(PyFloat_AS_DOUBLE(element)));
673
+ }
674
+ if (PyArray_IsScalar(element, Half)) {
675
+ return PyBool_FromLong(npy_half_isnan(PyArrayScalar_VAL(element, Half)));
676
+ }
677
+ if (PyArray_IsScalar(element, Float32)) {
678
+ return PyBool_FromLong(isnan(PyArrayScalar_VAL(element, Float32)));
679
+ }
680
+ if (PyArray_IsScalar(element, Float64)) {
681
+ return PyBool_FromLong(isnan(PyArrayScalar_VAL(element, Float64)));
682
+ }
683
+ # ifdef PyFloat128ArrType_Type
684
+ if (PyArray_IsScalar(element, Float128)) {
685
+ return PyBool_FromLong(isnan(PyArrayScalar_VAL(element, Float128)));
686
+ }
687
+ # endif
688
+
689
+ // Complex NaN
690
+ if (PyComplex_Check(element)) {
691
+ Py_complex val = ((PyComplexObject*)element)->cval;
692
+ return PyBool_FromLong(isnan(val.real) || isnan(val.imag));
693
+ }
694
+ if (PyArray_IsScalar(element, Complex64)) {
695
+ npy_cfloat val = PyArrayScalar_VAL(element, Complex64);
696
+ return PyBool_FromLong(isnan(npy_crealf(val)) || isnan(npy_cimagf(val)));
697
+ }
698
+ if (PyArray_IsScalar(element, Complex128)) {
699
+ npy_cdouble val = PyArrayScalar_VAL(element, Complex128);
700
+ return PyBool_FromLong(isnan(npy_creal(val)) || isnan(npy_cimag(val)));
701
+ }
702
+ # ifdef PyComplex256ArrType_Type
703
+ if (PyArray_IsScalar(element, Complex256)) {
704
+ npy_clongdouble val = PyArrayScalar_VAL(element, Complex256);
705
+ return PyBool_FromLong(isnan(npy_creall(val)) || isnan(npy_cimagl(val)));
706
+ }
707
+ # endif
708
+
709
+ // NaT - Datetime
710
+ if (PyArray_IsScalar(element, Datetime)) {
711
+ return PyBool_FromLong(PyArrayScalar_VAL(element, Datetime) == NPY_DATETIME_NAT);
712
+ }
713
+ // NaT - Timedelta
714
+ if (PyArray_IsScalar(element, Timedelta)) {
715
+ return PyBool_FromLong(PyArrayScalar_VAL(element, Timedelta) == NPY_DATETIME_NAT);
716
+ }
717
+ // Try to identify Pandas Timestamp NATs
718
+ if (PyObject_HasAttrString(element, "to_numpy")) {
719
+ // strcmp returns 0 on match
720
+ return PyBool_FromLong(strcmp(element->ob_type->tp_name, "NaTType") == 0);
721
+ // the long way
722
+ // PyObject *to_numpy = PyObject_GetAttrString(element, "to_numpy");
723
+ // if (to_numpy == NULL) {
724
+ // return NULL;
725
+ // }
726
+ // if (!PyCallable_Check(to_numpy)) {
727
+ // Py_DECREF(to_numpy);
728
+ // Py_RETURN_FALSE;
729
+ // }
730
+ // PyObject* scalar = PyObject_CallFunction(to_numpy, NULL);
731
+ // Py_DECREF(to_numpy);
732
+ // if (scalar == NULL) {
733
+ // return NULL;
734
+ // }
735
+ // if (!PyArray_IsScalar(scalar, Datetime)) {
736
+ // Py_DECREF(scalar);
737
+ // Py_RETURN_FALSE;
738
+ // }
739
+ // PyObject* pb = PyBool_FromLong(PyArrayScalar_VAL(scalar, Datetime) == NPY_DATETIME_NAT);
740
+ // Py_DECREF(scalar);
741
+ // return pb;
742
+ }
743
+ Py_RETURN_FALSE;
744
+ }
745
+
746
+ PyObject *
747
+ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
748
+ {
749
+ /*
750
+ Used to determine the new indexers and index screen in an index hierarchy selection.
751
+
752
+ Example:
753
+
754
+ Context:
755
+ We are an index hierarchy, constructing a new index hierarchy from a
756
+ selection of ourself. We need to build this up for each depth. For
757
+ example:
758
+
759
+ index_at_depth: ["a", "b", "c", "d"]
760
+ indexer_at_depth: [1, 0, 0, 2, 3, 0, 3, 3, 2]
761
+ (i.e. our index_hierarchy has these labels at depth = ["b", "a", "a", "c", "d", "a", "d", "d", "c"])
762
+
763
+ Imagine we are choosing this selection:
764
+ index_hierarchy.iloc[1:4]
765
+ At our depth, this would result in these labels: ["a", "a", "c", "d"]
766
+
767
+ We need to output:
768
+ index_screen: [0, 2, 3]
769
+ - New index is created by: index_at_depth[[0, 2, 3]] (i.e. ["a", "c", "d"])
770
+ new_indexer: [0, 0, 1, 2]
771
+ - When applied to our new_index, results in ["a", "a", "c", "d"]
772
+
773
+ Function:
774
+ input:
775
+ indexers: [0, 0, 2, 3] (i.e. indexer_at_depth[1:4])
776
+ positions: [0, 1, 2, 3] (i.e. which ilocs from index that ``indexers`` maps to)
777
+
778
+ algorithm:
779
+ Loop through ``indexers``. Since we know that ``indexers`` only contains
780
+ integers from 0 -> ``num_unique`` - 1, we can use a new indexers called
781
+ ``element_locations`` to keep track of which elements have been found, and when.
782
+ (Use ``num_unique`` as the flag for which elements have not been
783
+ found since it's not possible for one of our inputs to equal that)
784
+
785
+ Using the above example, this would look like:
786
+
787
+ element_locations =
788
+ [4, 4, 4, 4] (starting)
789
+ [0, 4, 4, 4] (first loop) indexers[0] = 0, so mark it as the 0th element found
790
+ [0, 4, 4, 4] (second loop) indexers[1] = 0, already marked, move on
791
+ [0, 4, 1, 4] (third loop) indexers[2] = 2, so mark it as the 1th element found
792
+ [0, 4, 1, 2] (fourth loop) indexers[3] = 3, so mark it as the 2th element found
793
+
794
+ Now, if during this loop, we discover every single element, it means
795
+ we can exit early, and just return back the original inputs, since
796
+ those arrays contain all the information the caller needs! This is the
797
+ core optimization of this function.
798
+ Example:
799
+ indexers = [0, 3, 1, 2, 3, 1, 0, 0]
800
+ positions = [0, 1, 2, 3]
801
+
802
+ There is no remapping needed! Simple re-use everything!
803
+
804
+ Now, if we don't find all the elements, then we need to construct
805
+ ``new_indexers`` and ``index_screen``.
806
+
807
+ We can construct ``new_indexers`` during the loop, by using the
808
+ information we have placed into ``element_locations``.
809
+
810
+ Using the above example, this would look like:
811
+ [x, x, x, x] (starting)
812
+ [0, x, x, x] (first loop) element_locations[indexers[0]] = 0
813
+ [0, 0, x, x] (second loop) element_locations[indexers[1]] = 0
814
+ [0, 0, 1, x] (third loop) element_locations[indexers[2]] = 1
815
+ [0, 0, 1, 2] (fourth loop) element_locations[indexers[3]] = 2
816
+
817
+ Finally, all that's left is to construct ``index_screen``, which
818
+ is essentially a way to condense and remap ``element_locations``.
819
+ See ``AK_get_index_screen`` for more details.
820
+
821
+ output:
822
+ index_screen: [0, 2, 3]
823
+ new_indexer: [0, 0, 1, 2]
824
+
825
+ Equivalent Python code:
826
+
827
+ num_unique = len(positions)
828
+ element_locations = np.full(num_unique, num_unique, dtype=np.int64)
829
+ order_found = np.full(num_unique, num_unique, dtype=np.int64)
830
+ new_indexers = np.empty(len(indexers), dtype=np.int64)
831
+
832
+ num_found = 0
833
+
834
+ for i, element in enumerate(indexers):
835
+ if element_locations[element] == num_unique:
836
+ element_locations[element] = num_found
837
+ order_found[num_found] = element
838
+ num_found += 1
839
+
840
+ if num_found == num_unique:
841
+ return positions, indexers
842
+
843
+ new_indexers[i] = element_locations[element]
844
+
845
+ return order_found[:num_found], new_indexers
846
+ */
847
+ PyArrayObject *indexers;
848
+ PyArrayObject *positions;
849
+
850
+ static char *kwlist[] = {"indexers", "positions", NULL};
851
+
852
+ if (!PyArg_ParseTupleAndKeywords(args,
853
+ kwargs,
854
+ "O!O!:get_new_indexers_and_screen", kwlist,
855
+ &PyArray_Type, &indexers,
856
+ &PyArray_Type, &positions
857
+ ))
858
+ {
859
+ return NULL;
860
+ }
861
+
862
+ if (PyArray_NDIM(indexers) != 1) {
863
+ PyErr_SetString(PyExc_ValueError, "indexers must be 1-dimensional");
864
+ return NULL;
865
+ }
866
+
867
+ if (PyArray_NDIM(positions) != 1) {
868
+ PyErr_SetString(PyExc_ValueError, "positions must be 1-dimensional");
869
+ return NULL;
870
+ }
871
+
872
+ if (PyArray_TYPE(indexers) != NPY_INT64) {
873
+ PyErr_SetString(PyExc_ValueError, "Array must be of type np.int64");
874
+ return NULL;
875
+ }
876
+
877
+ npy_intp num_unique = PyArray_SIZE(positions);
878
+
879
+ if (num_unique > PyArray_SIZE(indexers)) {
880
+ // This algorithm is only optimal if the number of unique elements is
881
+ // less than the number of elements in the indexers.
882
+ // Otherwise, the most optimal code is ``np.unique(indexers, return_index=True)``
883
+ // and we don't want to re-implement that in C.
884
+ PyErr_SetString(
885
+ PyExc_ValueError,
886
+ "Number of unique elements must be less than or equal to the length of ``indexers``"
887
+ );
888
+ return NULL;
889
+ }
890
+
891
+ npy_intp dims = {num_unique};
892
+ PyArrayObject *element_locations = (PyArrayObject*)PyArray_EMPTY(
893
+ 1, // ndim
894
+ &dims, // shape
895
+ NPY_INT64, // dtype
896
+ 0 // fortran
897
+ );
898
+ if (element_locations == NULL) {
899
+ return NULL;
900
+ }
901
+
902
+ PyArrayObject *order_found = (PyArrayObject*)PyArray_EMPTY(
903
+ 1, // ndim
904
+ &dims, // shape
905
+ NPY_INT64, // dtype
906
+ 0 // fortran
907
+ );
908
+ if (order_found == NULL) {
909
+ Py_DECREF(element_locations);
910
+ return NULL;
911
+ }
912
+
913
+ PyObject *num_unique_pyint = PyLong_FromLong((long)num_unique);
914
+ if (num_unique_pyint == NULL) {
915
+ goto fail;
916
+ }
917
+
918
+ // We use ``num_unique`` here to signal that we haven't found the element yet
919
+ // This works, because each element must be 0 < num_unique.
920
+ int fill_success = PyArray_FillWithScalar(element_locations, num_unique_pyint);
921
+ if (fill_success != 0) {
922
+ Py_DECREF(num_unique_pyint);
923
+ goto fail;
924
+ }
925
+
926
+ fill_success = PyArray_FillWithScalar(order_found, num_unique_pyint);
927
+ Py_DECREF(num_unique_pyint);
928
+ if (fill_success != 0) {
929
+ goto fail;
930
+ }
931
+
932
+ PyArrayObject *new_indexers = (PyArrayObject*)PyArray_EMPTY(
933
+ 1, // ndim
934
+ PyArray_DIMS(indexers), // shape
935
+ NPY_INT64, // dtype
936
+ 0 // fortran
937
+ );
938
+ if (new_indexers == NULL) {
939
+ goto fail;
940
+ }
941
+
942
+ // We know that our incoming dtypes are all int64! This is a safe cast.
943
+ // Plus, it's easier (and less error prone) to work with native C-arrays
944
+ // over using numpy's iteration APIs.
945
+ npy_int64 *element_location_values = (npy_int64*)PyArray_DATA(element_locations);
946
+ npy_int64 *order_found_values = (npy_int64*)PyArray_DATA(order_found);
947
+ npy_int64 *new_indexers_values = (npy_int64*)PyArray_DATA(new_indexers);
948
+
949
+ // Now, implement the core algorithm by looping over the ``indexers``.
950
+ // We need to use numpy's iteration API, as the ``indexers`` could be
951
+ // C-contiguous, F-contiguous, both, or neither.
952
+ // See https://numpy.org/doc/stable/reference/c-api/iterator.html#simple-iteration-example
953
+ NpyIter *indexer_iter = NpyIter_New(
954
+ indexers, // array
955
+ NPY_ITER_READONLY | NPY_ITER_EXTERNAL_LOOP, // iter flags
956
+ NPY_KEEPORDER, // order
957
+ NPY_NO_CASTING, // casting
958
+ NULL // dtype
959
+ );
960
+ if (indexer_iter == NULL) {
961
+ Py_DECREF(new_indexers);
962
+ goto fail;
963
+ }
964
+
965
+ // The iternext function gets stored in a local variable so it can be called repeatedly in an efficient manner.
966
+ NpyIter_IterNextFunc *indexer_iternext = NpyIter_GetIterNext(indexer_iter, NULL);
967
+ if (indexer_iternext == NULL) {
968
+ NpyIter_Deallocate(indexer_iter);
969
+ Py_DECREF(new_indexers);
970
+ goto fail;
971
+ }
972
+
973
+ // All of these will be updated by the iterator
974
+ char **dataptr = NpyIter_GetDataPtrArray(indexer_iter);
975
+ npy_intp *strideptr = NpyIter_GetInnerStrideArray(indexer_iter);
976
+ npy_intp *innersizeptr = NpyIter_GetInnerLoopSizePtr(indexer_iter);
977
+
978
+ // No gil is required from here on!
979
+ NPY_BEGIN_THREADS_DEF;
980
+ NPY_BEGIN_THREADS;
981
+
982
+ size_t i = 0;
983
+ Py_ssize_t num_found = 0;
984
+ do {
985
+ // Get the inner loop data/stride/inner_size values
986
+ char* data = *dataptr;
987
+ npy_intp stride = *strideptr;
988
+ npy_intp inner_size = *innersizeptr;
989
+ npy_int64 element;
990
+
991
+ while (inner_size--) {
992
+ element = *((npy_int64 *)data);
993
+
994
+ if (element_location_values[element] == num_unique) {
995
+ element_location_values[element] = num_found;
996
+ order_found_values[num_found] = element;
997
+ ++num_found;
998
+
999
+ if (num_found == num_unique) {
1000
+ // This insight is core to the performance of the algorithm.
1001
+ // If we have found every possible indexer, we can simply return
1002
+ // back the inputs! Essentially, we can observe on <= single pass
1003
+ // that we have the opportunity for re-use
1004
+ goto finish_early;
1005
+ }
1006
+ }
1007
+
1008
+ new_indexers_values[i] = element_location_values[element];
1009
+
1010
+ data += stride;
1011
+ ++i;
1012
+ }
1013
+
1014
+ // Increment the iterator to the next inner loop
1015
+ } while(indexer_iternext(indexer_iter));
1016
+
1017
+ NPY_END_THREADS;
1018
+
1019
+ NpyIter_Deallocate(indexer_iter);
1020
+ Py_DECREF(element_locations);
1021
+
1022
+ // new_positions = order_found[:num_unique]
1023
+ PyObject *new_positions = PySequence_GetSlice(
1024
+ (PyObject*)order_found,
1025
+ 0,
1026
+ num_found);
1027
+ Py_DECREF(order_found);
1028
+ if (new_positions == NULL) {
1029
+ return NULL;
1030
+ }
1031
+
1032
+ // return new_positions, new_indexers
1033
+ PyObject *result = PyTuple_Pack(2, new_positions, new_indexers);
1034
+ Py_DECREF(new_indexers);
1035
+ Py_DECREF(new_positions);
1036
+ return result;
1037
+
1038
+ finish_early:
1039
+ NPY_END_THREADS;
1040
+
1041
+ NpyIter_Deallocate(indexer_iter);
1042
+ Py_DECREF(element_locations);
1043
+ Py_DECREF(order_found);
1044
+ Py_DECREF(new_indexers);
1045
+ return PyTuple_Pack(2, positions, indexers);
1046
+
1047
+ fail:
1048
+ Py_DECREF(element_locations);
1049
+ Py_DECREF(order_found);
1050
+ return NULL;
1051
+ }
1052
+
1053
+ static char *array_deepcopy_kwarg_names[] = {
1054
+ "array",
1055
+ "memo",
1056
+ NULL
1057
+ };
1058
+
1059
+ PyObject *
1060
+ array_deepcopy(PyObject *m, PyObject *args, PyObject *kwargs)
1061
+ {
1062
+ PyObject *array;
1063
+ PyObject *memo = NULL;
1064
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
1065
+ "O|O:array_deepcopy", array_deepcopy_kwarg_names,
1066
+ &array,
1067
+ &memo)) {
1068
+ return NULL;
1069
+ }
1070
+ if ((memo == NULL) || (memo == Py_None)) {
1071
+ memo = NULL;
1072
+ }
1073
+ else {
1074
+ if (!PyDict_Check(memo)) {
1075
+ PyErr_SetString(PyExc_TypeError, "memo must be a dict or None");
1076
+ return NULL;
1077
+ }
1078
+ }
1079
+ AK_CHECK_NUMPY_ARRAY(array);
1080
+
1081
+ // Perform a deepcopy on an array, using an optional memo dictionary, and specialized to depend on immutable arrays. This depends on the module object to get the deepcopy method. The `memo` object can be None.
1082
+ PyObject *id = PyLong_FromVoidPtr(array);
1083
+ if (!id) return NULL;
1084
+
1085
+ if (memo) {
1086
+ PyObject *found = PyDict_GetItemWithError(memo, id);
1087
+ if (found) { // found will be NULL if not in dict
1088
+ Py_INCREF(found); // got a borrowed ref, increment first
1089
+ Py_DECREF(id);
1090
+ return found;
1091
+ }
1092
+ else if (PyErr_Occurred()) {
1093
+ goto error;
1094
+ }
1095
+ }
1096
+
1097
+ // if dtype is object, call deepcopy with memo
1098
+ PyObject *array_new;
1099
+ PyArray_Descr *dtype = PyArray_DESCR((PyArrayObject*)array); // borrowed ref
1100
+
1101
+ if (PyDataType_ISOBJECT(dtype)) {
1102
+ // we store the deepcopy function on this module for faster lookup here
1103
+ PyObject *deepcopy = PyObject_GetAttrString(m, "deepcopy");
1104
+ if (!deepcopy) {
1105
+ goto error;
1106
+ }
1107
+ array_new = PyObject_CallFunctionObjArgs(deepcopy, array, memo, NULL);
1108
+ Py_DECREF(deepcopy);
1109
+ if (!array_new) {
1110
+ goto error;
1111
+ }
1112
+ }
1113
+ else {
1114
+ // if not a n object dtype, we will force a copy (even if this is an immutable array) so as to not hold on to any references
1115
+ Py_INCREF(dtype); // PyArray_FromArray steals a reference
1116
+ array_new = PyArray_FromArray(
1117
+ (PyArrayObject*)array,
1118
+ dtype,
1119
+ NPY_ARRAY_ENSURECOPY);
1120
+ if (!array_new) {
1121
+ goto error;
1122
+ }
1123
+ if (memo && PyDict_SetItem(memo, id, array_new)) {
1124
+ Py_DECREF(array_new);
1125
+ goto error;
1126
+ }
1127
+ }
1128
+ // set immutable
1129
+ PyArray_CLEARFLAGS((PyArrayObject *)array_new, NPY_ARRAY_WRITEABLE);
1130
+ Py_DECREF(id);
1131
+ return array_new;
1132
+ error:
1133
+ Py_DECREF(id);
1134
+ return NULL;
1135
+ }