arraykit 0.10.0__cp313-cp313-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,1002 @@
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
+ static char *first_true_1d_kwarg_names[] = {
205
+ "array",
206
+ "forward",
207
+ NULL
208
+ };
209
+
210
+ PyObject *
211
+ first_true_1d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
212
+ {
213
+ PyArrayObject *array = NULL;
214
+ int forward = 1;
215
+
216
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
217
+ "O!|$p:first_true_1d",
218
+ first_true_1d_kwarg_names,
219
+ &PyArray_Type, &array,
220
+ &forward
221
+ )) {
222
+ return NULL;
223
+ }
224
+ if (PyArray_NDIM(array) != 1) {
225
+ PyErr_SetString(PyExc_ValueError, "Array must be 1-dimensional");
226
+ return NULL;
227
+ }
228
+ if (PyArray_TYPE(array) != NPY_BOOL) {
229
+ PyErr_SetString(PyExc_ValueError, "Array must be of type bool");
230
+ return NULL;
231
+ }
232
+ if (!PyArray_IS_C_CONTIGUOUS(array)) {
233
+ PyErr_SetString(PyExc_ValueError, "Array must be contiguous");
234
+ return NULL;
235
+ }
236
+
237
+ npy_intp lookahead = sizeof(npy_uint64);
238
+ npy_intp size = PyArray_SIZE(array);
239
+ lldiv_t size_div = lldiv((long long)size, lookahead); // quot, rem
240
+
241
+ npy_bool *array_buffer = (npy_bool*)PyArray_DATA(array);
242
+
243
+ NPY_BEGIN_THREADS_DEF;
244
+ NPY_BEGIN_THREADS;
245
+
246
+ Py_ssize_t position = -1;
247
+ npy_bool *p;
248
+ npy_bool *p_end;
249
+ npy_bool *p_end_roll;
250
+
251
+ if (forward) {
252
+ p = array_buffer;
253
+ p_end = p + size;
254
+ p_end_roll = p_end - size_div.rem;
255
+
256
+ while (p < p_end_roll) {
257
+ if (*(npy_uint64*)p != 0) {
258
+ break; // found a true within lookahead
259
+ }
260
+ p += lookahead;
261
+ }
262
+ while (p < p_end) {
263
+ if (*p) break;
264
+ p++;
265
+ }
266
+ }
267
+ else {
268
+ p = array_buffer + size - 1;
269
+ p_end = array_buffer - 1;
270
+ p_end_roll = p_end + size_div.rem;
271
+
272
+ while (p > p_end_roll) {
273
+ if (*(npy_uint64*)(p - lookahead + 1) != 0) {
274
+ break; // found a true within lookahead
275
+ }
276
+ p -= lookahead;
277
+ }
278
+ while (p > p_end) {
279
+ if (*p) break;
280
+ p--;
281
+ }
282
+ }
283
+ if (p != p_end) { // else, return -1
284
+ position = p - array_buffer;
285
+ }
286
+ NPY_END_THREADS;
287
+
288
+ PyObject* post = PyLong_FromSsize_t(position);
289
+ return post;
290
+ }
291
+
292
+ static char *first_true_2d_kwarg_names[] = {
293
+ "array",
294
+ "forward",
295
+ "axis",
296
+ NULL
297
+ };
298
+
299
+ PyObject *
300
+ first_true_2d(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
301
+ {
302
+ PyArrayObject *array = NULL;
303
+ int forward = 1;
304
+ int axis = 0;
305
+
306
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
307
+ "O!|$pi:first_true_2d",
308
+ first_true_2d_kwarg_names,
309
+ &PyArray_Type,
310
+ &array,
311
+ &forward,
312
+ &axis
313
+ )) {
314
+ return NULL;
315
+ }
316
+ if (PyArray_NDIM(array) != 2) {
317
+ PyErr_SetString(PyExc_ValueError, "Array must be 2-dimensional");
318
+ return NULL;
319
+ }
320
+ if (PyArray_TYPE(array) != NPY_BOOL) {
321
+ PyErr_SetString(PyExc_ValueError, "Array must be of type bool");
322
+ return NULL;
323
+ }
324
+ if (axis < 0 || axis > 1) {
325
+ PyErr_SetString(PyExc_ValueError, "Axis must be 0 or 1");
326
+ return NULL;
327
+ }
328
+
329
+ // NOTE: we copy the entire array into contiguous memory when necessary.
330
+ // axis = 0 returns the pos per col
331
+ // axis = 1 returns the pos per row (as contiguous bytes)
332
+ // if c contiguous:
333
+ // axis == 0: transpose, copy to C
334
+ // axis == 1: keep
335
+ // if f contiguous:
336
+ // axis == 0: transpose, keep
337
+ // axis == 1: copy to C
338
+ // else
339
+ // axis == 0: transpose, copy to C
340
+ // axis == 1: copy to C
341
+
342
+ bool transpose = !axis; // if 1, false
343
+ bool corder = true;
344
+ if ((PyArray_IS_C_CONTIGUOUS(array) && axis == 1) ||
345
+ (PyArray_IS_F_CONTIGUOUS(array) && axis == 0)) {
346
+ corder = false;
347
+ }
348
+ // create pointer to "indicator" array; if newly allocated, it will need to be decrefed before function termination
349
+ PyArrayObject *array_ind = NULL;
350
+ bool decref_array_ind = false;
351
+
352
+ if (transpose && !corder) {
353
+ array_ind = (PyArrayObject *)PyArray_Transpose(array, NULL);
354
+ if (array_ind == NULL) return NULL;
355
+ decref_array_ind = true;
356
+ }
357
+ else if (!transpose && corder) {
358
+ array_ind = (PyArrayObject *)PyArray_NewCopy(array, NPY_CORDER);
359
+ if (array_ind == NULL) return NULL;
360
+ decref_array_ind = true;
361
+ }
362
+ else if (transpose && corder) {
363
+ PyArrayObject *tmp = (PyArrayObject *)PyArray_Transpose(array, NULL);
364
+ if (tmp == NULL) return NULL;
365
+
366
+ array_ind = (PyArrayObject *)PyArray_NewCopy(tmp, NPY_CORDER);
367
+ Py_DECREF((PyObject*)tmp);
368
+ if (array_ind == NULL) return NULL;
369
+ decref_array_ind = true;
370
+ }
371
+ else {
372
+ array_ind = array; // can use array, no decref needed
373
+ }
374
+
375
+ npy_intp lookahead = sizeof(npy_uint64);
376
+
377
+ // buffer of indicators
378
+ npy_bool *buffer_ind = (npy_bool*)PyArray_DATA(array_ind);
379
+
380
+ npy_intp count_row = PyArray_DIM(array_ind, 0);
381
+ npy_intp count_col = PyArray_DIM(array_ind, 1);
382
+
383
+ lldiv_t div_col = lldiv((long long)count_col, lookahead); // quot, rem
384
+
385
+ npy_intp dims_post = {count_row};
386
+ PyArrayObject *array_pos = (PyArrayObject*)PyArray_EMPTY(
387
+ 1, // ndim
388
+ &dims_post,// shape
389
+ NPY_INT64, // dtype
390
+ 0 // fortran
391
+ );
392
+ if (array_pos == NULL) {
393
+ return NULL;
394
+ }
395
+ npy_int64 *buffer_pos = (npy_int64*)PyArray_DATA(array_pos);
396
+
397
+ NPY_BEGIN_THREADS_DEF;
398
+ NPY_BEGIN_THREADS;
399
+
400
+ npy_intp position;
401
+ npy_bool *p;
402
+ npy_bool *p_start;
403
+ npy_bool *p_end;
404
+
405
+ // iterate one row at a time; short-circult when found
406
+ // for axis 1 rows are rows; for axis 0, rows are (post transpose) columns
407
+ for (npy_intp r = 0; r < count_row; r++) {
408
+ position = -1; // update for each row
409
+
410
+ if (forward) {
411
+ // get start of each row
412
+ p_start = buffer_ind + (count_col * r);
413
+ p = p_start;
414
+ p_end = p + count_col; // end of each row
415
+
416
+ // scan each row from the front and terminate when True
417
+ // remove from the end the remainder
418
+ while (p < p_end - div_col.rem) {
419
+ if (*(npy_uint64*)p != 0) {
420
+ break; // found a true
421
+ }
422
+ p += lookahead;
423
+ }
424
+ while (p < p_end) {
425
+ if (*p) {break;}
426
+ p++;
427
+ }
428
+ if (p != p_end) {
429
+ position = p - p_start;
430
+ }
431
+ }
432
+ else { // reverse
433
+ // start at the next row, then subtract one for last elem in previous row
434
+ p_start = buffer_ind + (count_col * (r + 1)) - 1;
435
+ p = p_start;
436
+ // end is 1 less than start of each row
437
+ p_end = buffer_ind + (count_col * r) - 1;
438
+
439
+ while (p > p_end + div_col.rem) {
440
+ // must go to start of lookahead
441
+ if (*(npy_uint64*)(p - lookahead + 1) != 0) {
442
+ break; // found a true
443
+ }
444
+ p -= lookahead;
445
+ }
446
+ while (p > p_end) {
447
+ if (*p) {break;}
448
+ p--;
449
+ }
450
+ if (p != p_end) {
451
+ position = p - (p_end + 1);
452
+ }
453
+ }
454
+ *buffer_pos++ = position;
455
+ }
456
+
457
+ NPY_END_THREADS;
458
+
459
+ if (decref_array_ind) {
460
+ Py_DECREF(array_ind); // created in this function
461
+ }
462
+ return (PyObject *)array_pos;
463
+ }
464
+
465
+ PyObject *
466
+ dtype_from_element(PyObject *Py_UNUSED(m), PyObject *arg)
467
+ {
468
+ // -------------------------------------------------------------------------
469
+ // 1. Handle fast, exact type checks first.
470
+ if (arg == Py_None) {
471
+ return (PyObject*)PyArray_DescrFromType(NPY_OBJECT);
472
+ }
473
+ if (PyFloat_CheckExact(arg)) {
474
+ return (PyObject*)PyArray_DescrFromType(NPY_FLOAT64);
475
+ }
476
+ if (PyLong_CheckExact(arg)) {
477
+ return (PyObject*)PyArray_DescrFromType(NPY_INT64);
478
+ }
479
+ if (PyBool_Check(arg)) {
480
+ return (PyObject*)PyArray_DescrFromType(NPY_BOOL);
481
+ }
482
+
483
+ PyObject* dtype = NULL;
484
+ // String
485
+ if (PyUnicode_CheckExact(arg)) {
486
+ PyArray_Descr* descr = PyArray_DescrFromType(NPY_UNICODE);
487
+ if (descr == NULL) return NULL;
488
+ dtype = (PyObject*)PyArray_DescrFromObject(arg, descr);
489
+ Py_DECREF(descr);
490
+ return dtype;
491
+ }
492
+ // Bytes
493
+ if (PyBytes_CheckExact(arg)) {
494
+ PyArray_Descr* descr = PyArray_DescrFromType(NPY_STRING);
495
+ if (descr == NULL) return NULL;
496
+ dtype = (PyObject*)PyArray_DescrFromObject(arg, descr);
497
+ Py_DECREF(descr);
498
+ return dtype;
499
+ }
500
+
501
+ // -------------------------------------------------------------------------
502
+ // 2. Construct dtype (slightly more complicated)
503
+ // Already known
504
+ dtype = PyObject_GetAttrString(arg, "dtype");
505
+ if (dtype) {
506
+ return dtype;
507
+ }
508
+ PyErr_Clear();
509
+ // -------------------------------------------------------------------------
510
+ // 3. Handles everything else.
511
+ return (PyObject*)PyArray_DescrFromType(NPY_OBJECT);
512
+ }
513
+
514
+ static char *isna_element_kwarg_names[] = {
515
+ "element",
516
+ "include_none",
517
+ NULL
518
+ };
519
+
520
+ PyObject *
521
+ isna_element(PyObject *m, PyObject *args, PyObject *kwargs)
522
+ {
523
+ PyObject *element;
524
+ int include_none = 1;
525
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
526
+ "O|p:isna_element", isna_element_kwarg_names,
527
+ &element,
528
+ &include_none)) {
529
+ return NULL;
530
+ }
531
+
532
+ // None
533
+ if (include_none && element == Py_None) {
534
+ Py_RETURN_TRUE;
535
+ }
536
+
537
+ // NaN
538
+ if (PyFloat_Check(element)) {
539
+ return PyBool_FromLong(isnan(PyFloat_AS_DOUBLE(element)));
540
+ }
541
+ if (PyArray_IsScalar(element, Half)) {
542
+ return PyBool_FromLong(npy_half_isnan(PyArrayScalar_VAL(element, Half)));
543
+ }
544
+ if (PyArray_IsScalar(element, Float32)) {
545
+ return PyBool_FromLong(isnan(PyArrayScalar_VAL(element, Float32)));
546
+ }
547
+ if (PyArray_IsScalar(element, Float64)) {
548
+ return PyBool_FromLong(isnan(PyArrayScalar_VAL(element, Float64)));
549
+ }
550
+ # ifdef PyFloat128ArrType_Type
551
+ if (PyArray_IsScalar(element, Float128)) {
552
+ return PyBool_FromLong(isnan(PyArrayScalar_VAL(element, Float128)));
553
+ }
554
+ # endif
555
+
556
+ // Complex NaN
557
+ if (PyComplex_Check(element)) {
558
+ Py_complex val = ((PyComplexObject*)element)->cval;
559
+ return PyBool_FromLong(isnan(val.real) || isnan(val.imag));
560
+ }
561
+ if (PyArray_IsScalar(element, Complex64)) {
562
+ npy_cfloat val = PyArrayScalar_VAL(element, Complex64);
563
+ return PyBool_FromLong(isnan(npy_crealf(val)) || isnan(npy_cimagf(val)));
564
+ }
565
+ if (PyArray_IsScalar(element, Complex128)) {
566
+ npy_cdouble val = PyArrayScalar_VAL(element, Complex128);
567
+ return PyBool_FromLong(isnan(npy_creal(val)) || isnan(npy_cimag(val)));
568
+ }
569
+ # ifdef PyComplex256ArrType_Type
570
+ if (PyArray_IsScalar(element, Complex256)) {
571
+ npy_clongdouble val = PyArrayScalar_VAL(element, Complex256);
572
+ return PyBool_FromLong(isnan(npy_creall(val)) || isnan(npy_cimagl(val)));
573
+ }
574
+ # endif
575
+
576
+ // NaT - Datetime
577
+ if (PyArray_IsScalar(element, Datetime)) {
578
+ return PyBool_FromLong(PyArrayScalar_VAL(element, Datetime) == NPY_DATETIME_NAT);
579
+ }
580
+ // NaT - Timedelta
581
+ if (PyArray_IsScalar(element, Timedelta)) {
582
+ return PyBool_FromLong(PyArrayScalar_VAL(element, Timedelta) == NPY_DATETIME_NAT);
583
+ }
584
+ // Try to identify Pandas Timestamp NATs
585
+ if (PyObject_HasAttrString(element, "to_numpy")) {
586
+ // strcmp returns 0 on match
587
+ return PyBool_FromLong(strcmp(element->ob_type->tp_name, "NaTType") == 0);
588
+ // the long way
589
+ // PyObject *to_numpy = PyObject_GetAttrString(element, "to_numpy");
590
+ // if (to_numpy == NULL) {
591
+ // return NULL;
592
+ // }
593
+ // if (!PyCallable_Check(to_numpy)) {
594
+ // Py_DECREF(to_numpy);
595
+ // Py_RETURN_FALSE;
596
+ // }
597
+ // PyObject* scalar = PyObject_CallFunction(to_numpy, NULL);
598
+ // Py_DECREF(to_numpy);
599
+ // if (scalar == NULL) {
600
+ // return NULL;
601
+ // }
602
+ // if (!PyArray_IsScalar(scalar, Datetime)) {
603
+ // Py_DECREF(scalar);
604
+ // Py_RETURN_FALSE;
605
+ // }
606
+ // PyObject* pb = PyBool_FromLong(PyArrayScalar_VAL(scalar, Datetime) == NPY_DATETIME_NAT);
607
+ // Py_DECREF(scalar);
608
+ // return pb;
609
+ }
610
+ Py_RETURN_FALSE;
611
+ }
612
+
613
+ PyObject *
614
+ get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
615
+ {
616
+ /*
617
+ Used to determine the new indexers and index screen in an index hierarchy selection.
618
+
619
+ Example:
620
+
621
+ Context:
622
+ We are an index hierarchy, constructing a new index hierarchy from a
623
+ selection of ourself. We need to build this up for each depth. For
624
+ example:
625
+
626
+ index_at_depth: ["a", "b", "c", "d"]
627
+ indexer_at_depth: [1, 0, 0, 2, 3, 0, 3, 3, 2]
628
+ (i.e. our index_hierarchy has these labels at depth = ["b", "a", "a", "c", "d", "a", "d", "d", "c"])
629
+
630
+ Imagine we are choosing this selection:
631
+ index_hierarchy.iloc[1:4]
632
+ At our depth, this would result in these labels: ["a", "a", "c", "d"]
633
+
634
+ We need to output:
635
+ index_screen: [0, 2, 3]
636
+ - New index is created by: index_at_depth[[0, 2, 3]] (i.e. ["a", "c", "d"])
637
+ new_indexer: [0, 0, 1, 2]
638
+ - When applied to our new_index, results in ["a", "a", "c", "d"]
639
+
640
+ Function:
641
+ input:
642
+ indexers: [0, 0, 2, 3] (i.e. indexer_at_depth[1:4])
643
+ positions: [0, 1, 2, 3] (i.e. which ilocs from index that ``indexers`` maps to)
644
+
645
+ algorithm:
646
+ Loop through ``indexers``. Since we know that ``indexers`` only contains
647
+ integers from 0 -> ``num_unique`` - 1, we can use a new indexers called
648
+ ``element_locations`` to keep track of which elements have been found, and when.
649
+ (Use ``num_unique`` as the flag for which elements have not been
650
+ found since it's not possible for one of our inputs to equal that)
651
+
652
+ Using the above example, this would look like:
653
+
654
+ element_locations =
655
+ [4, 4, 4, 4] (starting)
656
+ [0, 4, 4, 4] (first loop) indexers[0] = 0, so mark it as the 0th element found
657
+ [0, 4, 4, 4] (second loop) indexers[1] = 0, already marked, move on
658
+ [0, 4, 1, 4] (third loop) indexers[2] = 2, so mark it as the 1th element found
659
+ [0, 4, 1, 2] (fourth loop) indexers[3] = 3, so mark it as the 2th element found
660
+
661
+ Now, if during this loop, we discover every single element, it means
662
+ we can exit early, and just return back the original inputs, since
663
+ those arrays contain all the information the caller needs! This is the
664
+ core optimization of this function.
665
+ Example:
666
+ indexers = [0, 3, 1, 2, 3, 1, 0, 0]
667
+ positions = [0, 1, 2, 3]
668
+
669
+ There is no remapping needed! Simple re-use everything!
670
+
671
+ Now, if we don't find all the elements, then we need to construct
672
+ ``new_indexers`` and ``index_screen``.
673
+
674
+ We can construct ``new_indexers`` during the loop, by using the
675
+ information we have placed into ``element_locations``.
676
+
677
+ Using the above example, this would look like:
678
+ [x, x, x, x] (starting)
679
+ [0, x, x, x] (first loop) element_locations[indexers[0]] = 0
680
+ [0, 0, x, x] (second loop) element_locations[indexers[1]] = 0
681
+ [0, 0, 1, x] (third loop) element_locations[indexers[2]] = 1
682
+ [0, 0, 1, 2] (fourth loop) element_locations[indexers[3]] = 2
683
+
684
+ Finally, all that's left is to construct ``index_screen``, which
685
+ is essentially a way to condense and remap ``element_locations``.
686
+ See ``AK_get_index_screen`` for more details.
687
+
688
+ output:
689
+ index_screen: [0, 2, 3]
690
+ new_indexer: [0, 0, 1, 2]
691
+
692
+ Equivalent Python code:
693
+
694
+ num_unique = len(positions)
695
+ element_locations = np.full(num_unique, num_unique, dtype=np.int64)
696
+ order_found = np.full(num_unique, num_unique, dtype=np.int64)
697
+ new_indexers = np.empty(len(indexers), dtype=np.int64)
698
+
699
+ num_found = 0
700
+
701
+ for i, element in enumerate(indexers):
702
+ if element_locations[element] == num_unique:
703
+ element_locations[element] = num_found
704
+ order_found[num_found] = element
705
+ num_found += 1
706
+
707
+ if num_found == num_unique:
708
+ return positions, indexers
709
+
710
+ new_indexers[i] = element_locations[element]
711
+
712
+ return order_found[:num_found], new_indexers
713
+ */
714
+ PyArrayObject *indexers;
715
+ PyArrayObject *positions;
716
+
717
+ static char *kwlist[] = {"indexers", "positions", NULL};
718
+
719
+ if (!PyArg_ParseTupleAndKeywords(args,
720
+ kwargs,
721
+ "O!O!:get_new_indexers_and_screen", kwlist,
722
+ &PyArray_Type, &indexers,
723
+ &PyArray_Type, &positions
724
+ ))
725
+ {
726
+ return NULL;
727
+ }
728
+
729
+ if (PyArray_NDIM(indexers) != 1) {
730
+ PyErr_SetString(PyExc_ValueError, "indexers must be 1-dimensional");
731
+ return NULL;
732
+ }
733
+
734
+ if (PyArray_NDIM(positions) != 1) {
735
+ PyErr_SetString(PyExc_ValueError, "positions must be 1-dimensional");
736
+ return NULL;
737
+ }
738
+
739
+ if (PyArray_TYPE(indexers) != NPY_INT64) {
740
+ PyErr_SetString(PyExc_ValueError, "Array must be of type np.int64");
741
+ return NULL;
742
+ }
743
+
744
+ npy_intp num_unique = PyArray_SIZE(positions);
745
+
746
+ if (num_unique > PyArray_SIZE(indexers)) {
747
+ // This algorithm is only optimal if the number of unique elements is
748
+ // less than the number of elements in the indexers.
749
+ // Otherwise, the most optimal code is ``np.unique(indexers, return_index=True)``
750
+ // and we don't want to re-implement that in C.
751
+ PyErr_SetString(
752
+ PyExc_ValueError,
753
+ "Number of unique elements must be less than or equal to the length of ``indexers``"
754
+ );
755
+ return NULL;
756
+ }
757
+
758
+ npy_intp dims = {num_unique};
759
+ PyArrayObject *element_locations = (PyArrayObject*)PyArray_EMPTY(
760
+ 1, // ndim
761
+ &dims, // shape
762
+ NPY_INT64, // dtype
763
+ 0 // fortran
764
+ );
765
+ if (element_locations == NULL) {
766
+ return NULL;
767
+ }
768
+
769
+ PyArrayObject *order_found = (PyArrayObject*)PyArray_EMPTY(
770
+ 1, // ndim
771
+ &dims, // shape
772
+ NPY_INT64, // dtype
773
+ 0 // fortran
774
+ );
775
+ if (order_found == NULL) {
776
+ Py_DECREF(element_locations);
777
+ return NULL;
778
+ }
779
+
780
+ PyObject *num_unique_pyint = PyLong_FromLong((long)num_unique);
781
+ if (num_unique_pyint == NULL) {
782
+ goto fail;
783
+ }
784
+
785
+ // We use ``num_unique`` here to signal that we haven't found the element yet
786
+ // This works, because each element must be 0 < num_unique.
787
+ int fill_success = PyArray_FillWithScalar(element_locations, num_unique_pyint);
788
+ if (fill_success != 0) {
789
+ Py_DECREF(num_unique_pyint);
790
+ goto fail;
791
+ }
792
+
793
+ fill_success = PyArray_FillWithScalar(order_found, num_unique_pyint);
794
+ Py_DECREF(num_unique_pyint);
795
+ if (fill_success != 0) {
796
+ goto fail;
797
+ }
798
+
799
+ PyArrayObject *new_indexers = (PyArrayObject*)PyArray_EMPTY(
800
+ 1, // ndim
801
+ PyArray_DIMS(indexers), // shape
802
+ NPY_INT64, // dtype
803
+ 0 // fortran
804
+ );
805
+ if (new_indexers == NULL) {
806
+ goto fail;
807
+ }
808
+
809
+ // We know that our incoming dtypes are all int64! This is a safe cast.
810
+ // Plus, it's easier (and less error prone) to work with native C-arrays
811
+ // over using numpy's iteration APIs.
812
+ npy_int64 *element_location_values = (npy_int64*)PyArray_DATA(element_locations);
813
+ npy_int64 *order_found_values = (npy_int64*)PyArray_DATA(order_found);
814
+ npy_int64 *new_indexers_values = (npy_int64*)PyArray_DATA(new_indexers);
815
+
816
+ // Now, implement the core algorithm by looping over the ``indexers``.
817
+ // We need to use numpy's iteration API, as the ``indexers`` could be
818
+ // C-contiguous, F-contiguous, both, or neither.
819
+ // See https://numpy.org/doc/stable/reference/c-api/iterator.html#simple-iteration-example
820
+ NpyIter *indexer_iter = NpyIter_New(
821
+ indexers, // array
822
+ NPY_ITER_READONLY | NPY_ITER_EXTERNAL_LOOP, // iter flags
823
+ NPY_KEEPORDER, // order
824
+ NPY_NO_CASTING, // casting
825
+ NULL // dtype
826
+ );
827
+ if (indexer_iter == NULL) {
828
+ Py_DECREF(new_indexers);
829
+ goto fail;
830
+ }
831
+
832
+ // The iternext function gets stored in a local variable so it can be called repeatedly in an efficient manner.
833
+ NpyIter_IterNextFunc *indexer_iternext = NpyIter_GetIterNext(indexer_iter, NULL);
834
+ if (indexer_iternext == NULL) {
835
+ NpyIter_Deallocate(indexer_iter);
836
+ Py_DECREF(new_indexers);
837
+ goto fail;
838
+ }
839
+
840
+ // All of these will be updated by the iterator
841
+ char **dataptr = NpyIter_GetDataPtrArray(indexer_iter);
842
+ npy_intp *strideptr = NpyIter_GetInnerStrideArray(indexer_iter);
843
+ npy_intp *innersizeptr = NpyIter_GetInnerLoopSizePtr(indexer_iter);
844
+
845
+ // No gil is required from here on!
846
+ NPY_BEGIN_THREADS_DEF;
847
+ NPY_BEGIN_THREADS;
848
+
849
+ size_t i = 0;
850
+ Py_ssize_t num_found = 0;
851
+ do {
852
+ // Get the inner loop data/stride/inner_size values
853
+ char* data = *dataptr;
854
+ npy_intp stride = *strideptr;
855
+ npy_intp inner_size = *innersizeptr;
856
+ npy_int64 element;
857
+
858
+ while (inner_size--) {
859
+ element = *((npy_int64 *)data);
860
+
861
+ if (element_location_values[element] == num_unique) {
862
+ element_location_values[element] = num_found;
863
+ order_found_values[num_found] = element;
864
+ ++num_found;
865
+
866
+ if (num_found == num_unique) {
867
+ // This insight is core to the performance of the algorithm.
868
+ // If we have found every possible indexer, we can simply return
869
+ // back the inputs! Essentially, we can observe on <= single pass
870
+ // that we have the opportunity for re-use
871
+ goto finish_early;
872
+ }
873
+ }
874
+
875
+ new_indexers_values[i] = element_location_values[element];
876
+
877
+ data += stride;
878
+ ++i;
879
+ }
880
+
881
+ // Increment the iterator to the next inner loop
882
+ } while(indexer_iternext(indexer_iter));
883
+
884
+ NPY_END_THREADS;
885
+
886
+ NpyIter_Deallocate(indexer_iter);
887
+ Py_DECREF(element_locations);
888
+
889
+ // new_positions = order_found[:num_unique]
890
+ PyObject *new_positions = PySequence_GetSlice(
891
+ (PyObject*)order_found,
892
+ 0,
893
+ num_found);
894
+ Py_DECREF(order_found);
895
+ if (new_positions == NULL) {
896
+ return NULL;
897
+ }
898
+
899
+ // return new_positions, new_indexers
900
+ PyObject *result = PyTuple_Pack(2, new_positions, new_indexers);
901
+ Py_DECREF(new_indexers);
902
+ Py_DECREF(new_positions);
903
+ return result;
904
+
905
+ finish_early:
906
+ NPY_END_THREADS;
907
+
908
+ NpyIter_Deallocate(indexer_iter);
909
+ Py_DECREF(element_locations);
910
+ Py_DECREF(order_found);
911
+ Py_DECREF(new_indexers);
912
+ return PyTuple_Pack(2, positions, indexers);
913
+
914
+ fail:
915
+ Py_DECREF(element_locations);
916
+ Py_DECREF(order_found);
917
+ return NULL;
918
+ }
919
+
920
+ static char *array_deepcopy_kwarg_names[] = {
921
+ "array",
922
+ "memo",
923
+ NULL
924
+ };
925
+
926
+ PyObject *
927
+ array_deepcopy(PyObject *m, PyObject *args, PyObject *kwargs)
928
+ {
929
+ PyObject *array;
930
+ PyObject *memo = NULL;
931
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
932
+ "O|O:array_deepcopy", array_deepcopy_kwarg_names,
933
+ &array,
934
+ &memo)) {
935
+ return NULL;
936
+ }
937
+ if ((memo == NULL) || (memo == Py_None)) {
938
+ memo = NULL;
939
+ }
940
+ else {
941
+ if (!PyDict_Check(memo)) {
942
+ PyErr_SetString(PyExc_TypeError, "memo must be a dict or None");
943
+ return NULL;
944
+ }
945
+ }
946
+ AK_CHECK_NUMPY_ARRAY(array);
947
+
948
+ // 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.
949
+ PyObject *id = PyLong_FromVoidPtr(array);
950
+ if (!id) return NULL;
951
+
952
+ if (memo) {
953
+ PyObject *found = PyDict_GetItemWithError(memo, id);
954
+ if (found) { // found will be NULL if not in dict
955
+ Py_INCREF(found); // got a borrowed ref, increment first
956
+ Py_DECREF(id);
957
+ return found;
958
+ }
959
+ else if (PyErr_Occurred()) {
960
+ goto error;
961
+ }
962
+ }
963
+
964
+ // if dtype is object, call deepcopy with memo
965
+ PyObject *array_new;
966
+ PyArray_Descr *dtype = PyArray_DESCR((PyArrayObject*)array); // borrowed ref
967
+
968
+ if (PyDataType_ISOBJECT(dtype)) {
969
+ // we store the deepcopy function on this module for faster lookup here
970
+ PyObject *deepcopy = PyObject_GetAttrString(m, "deepcopy");
971
+ if (!deepcopy) {
972
+ goto error;
973
+ }
974
+ array_new = PyObject_CallFunctionObjArgs(deepcopy, array, memo, NULL);
975
+ Py_DECREF(deepcopy);
976
+ if (!array_new) {
977
+ goto error;
978
+ }
979
+ }
980
+ else {
981
+ // 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
982
+ Py_INCREF(dtype); // PyArray_FromArray steals a reference
983
+ array_new = PyArray_FromArray(
984
+ (PyArrayObject*)array,
985
+ dtype,
986
+ NPY_ARRAY_ENSURECOPY);
987
+ if (!array_new) {
988
+ goto error;
989
+ }
990
+ if (memo && PyDict_SetItem(memo, id, array_new)) {
991
+ Py_DECREF(array_new);
992
+ goto error;
993
+ }
994
+ }
995
+ // set immutable
996
+ PyArray_CLEARFLAGS((PyArrayObject *)array_new, NPY_ARRAY_WRITEABLE);
997
+ Py_DECREF(id);
998
+ return array_new;
999
+ error:
1000
+ Py_DECREF(id);
1001
+ return NULL;
1002
+ }