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/tri_map.c ADDED
@@ -0,0 +1,1386 @@
1
+ # include "Python.h"
2
+ # include "stdbool.h"
3
+
4
+ # define NO_IMPORT_ARRAY
5
+ # define PY_ARRAY_UNIQUE_SYMBOL AK_ARRAY_API
6
+ # define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
7
+
8
+ # include "numpy/arrayobject.h"
9
+ # include "numpy/arrayscalars.h"
10
+
11
+ # include "tri_map.h"
12
+ # include "utilities.h"
13
+
14
+ typedef struct TriMapOne {
15
+ Py_ssize_t from; // signed
16
+ Py_ssize_t to;
17
+ } TriMapOne;
18
+
19
+ typedef struct TriMapManyTo {
20
+ Py_ssize_t start;
21
+ Py_ssize_t stop;
22
+ } TriMapManyTo;
23
+
24
+ typedef struct TriMapManyFrom {
25
+ npy_intp src;
26
+ PyArrayObject* dst;
27
+ } TriMapManyFrom;
28
+
29
+ typedef struct TriMapObject {
30
+ PyObject_HEAD
31
+ Py_ssize_t src_len;
32
+ Py_ssize_t dst_len;
33
+ Py_ssize_t len;
34
+ bool is_many;
35
+ bool finalized;
36
+
37
+ PyObject* src_match; // array object
38
+ npy_bool* src_match_data; // contiguous C array
39
+ PyObject* dst_match; // array object
40
+ npy_bool* dst_match_data; // contiguous C array
41
+
42
+ PyObject* final_src_fill; // array object
43
+ PyObject* final_dst_fill; // array object
44
+
45
+ // register one
46
+ TriMapOne* src_one;
47
+ Py_ssize_t src_one_count;
48
+ Py_ssize_t src_one_capacity;
49
+
50
+ TriMapOne* dst_one;
51
+ Py_ssize_t dst_one_count;
52
+ Py_ssize_t dst_one_capacity;
53
+
54
+ // register_many
55
+ TriMapManyTo* many_to; // two integers for contiguous assignment region
56
+ TriMapManyFrom* many_from; // int and array, for src and dst (together)
57
+ Py_ssize_t many_count;
58
+ Py_ssize_t many_capacity;
59
+
60
+ } TriMapObject;
61
+
62
+ PyObject *
63
+ TriMap_new(PyTypeObject *cls, PyObject *args, PyObject *kwargs) {
64
+ TriMapObject *self = (TriMapObject *)cls->tp_alloc(cls, 0);
65
+ if (!self) {
66
+ return NULL;
67
+ }
68
+ return (PyObject *)self;
69
+ }
70
+
71
+ PyDoc_STRVAR(
72
+ TriMap_doc,
73
+ "\n"
74
+ "A utilty for three-way join mappings."
75
+ );
76
+
77
+ // Returns 0 on success, -1 on error.
78
+ int
79
+ TriMap_init(PyObject *self, PyObject *args, PyObject *kwargs) {
80
+ Py_ssize_t src_len;
81
+ Py_ssize_t dst_len;
82
+ if (!PyArg_ParseTuple(args,
83
+ "nn:__init__",
84
+ &src_len,
85
+ &dst_len)) {
86
+ return -1;
87
+ }
88
+ TriMapObject* tm = (TriMapObject*)self;
89
+ // handle all C types
90
+ tm->src_len = src_len;
91
+ tm->dst_len = dst_len;
92
+ tm->is_many = false;
93
+ tm->finalized = false;
94
+ tm->len = 0;
95
+
96
+ // we create arrays, and also pre-extract pointers to array data for fast insertion; we keep the array for optimal summing routines
97
+ npy_intp dims_src_len[] = {src_len};
98
+ tm->src_match = PyArray_ZEROS(1, dims_src_len, NPY_BOOL, 0);
99
+ if (tm->src_match == NULL) {
100
+ return -1;
101
+ }
102
+ tm->src_match_data = (npy_bool*)PyArray_DATA((PyArrayObject*)tm->src_match);
103
+
104
+ npy_intp dims_dst_len[] = {dst_len};
105
+ tm->dst_match = PyArray_ZEROS(1, dims_dst_len, NPY_BOOL, 0);
106
+ if (tm->dst_match == NULL) {
107
+ return -1;
108
+ }
109
+ tm->dst_match_data = (npy_bool*)PyArray_DATA((PyArrayObject*)tm->dst_match);
110
+
111
+ // register one
112
+ tm->src_one_count = 0;
113
+ tm->src_one_capacity = 16;
114
+ tm->src_one = (TriMapOne*)PyMem_Malloc(
115
+ sizeof(TriMapOne) * tm->src_one_capacity);
116
+ if (tm->src_one == NULL) {
117
+ PyErr_SetNone(PyExc_MemoryError);
118
+ return -1;
119
+ }
120
+ tm->dst_one_count = 0;
121
+ tm->dst_one_capacity = 16;
122
+ tm->dst_one = (TriMapOne*)PyMem_Malloc(
123
+ sizeof(TriMapOne) * tm->dst_one_capacity);
124
+ if (tm->dst_one == NULL) {
125
+ PyErr_SetNone(PyExc_MemoryError);
126
+ return -1;
127
+ }
128
+ // register many
129
+ tm->many_count = 0;
130
+ tm->many_capacity = 16;
131
+ tm->many_to = (TriMapManyTo*)PyMem_Malloc(
132
+ sizeof(TriMapManyTo) * tm->many_capacity);
133
+ if (tm->many_to == NULL) {
134
+ PyErr_SetNone(PyExc_MemoryError);
135
+ return -1;
136
+ }
137
+ tm->many_from = (TriMapManyFrom*)PyMem_Malloc(
138
+ sizeof(TriMapManyFrom) * tm->many_capacity);
139
+ if (tm->many_from == NULL) {
140
+ PyErr_SetNone(PyExc_MemoryError);
141
+ return -1;
142
+ }
143
+
144
+ return 0;
145
+ }
146
+
147
+ void
148
+ TriMap_dealloc(TriMapObject *self) {
149
+ // NOTE: we use XDECREF incase init fails before these objects get allocated
150
+ Py_XDECREF(self->src_match);
151
+ Py_XDECREF(self->dst_match);
152
+ Py_XDECREF(self->final_src_fill);
153
+ Py_XDECREF(self->final_dst_fill);
154
+
155
+ if (self->src_one != NULL) {
156
+ PyMem_Free(self->src_one);
157
+ }
158
+ if (self->dst_one != NULL) {
159
+ PyMem_Free(self->dst_one);
160
+ }
161
+ if (self->many_to != NULL) {
162
+ PyMem_Free(self->many_to);
163
+ }
164
+ if (self->many_from != NULL) {
165
+ // decref all arrays before freeing
166
+ for (Py_ssize_t i = 0; i < self->many_count; i++) {
167
+ // NOTE: using dot to get to pointer?
168
+ Py_DECREF((PyObject*)self->many_from[i].dst);
169
+ }
170
+ PyMem_Free(self->many_from);
171
+ }
172
+ Py_TYPE(self)->tp_free((PyObject *)self);
173
+ }
174
+
175
+ PyObject *
176
+ TriMap_repr(TriMapObject *self) {
177
+ const char *is_many = self->is_many ? "true" : "false";
178
+ const char *is_finalized = self->finalized ? "true" : "false";
179
+
180
+ npy_intp src_fill;
181
+ npy_intp dst_fill;
182
+ if (self->finalized) {
183
+ src_fill = PyArray_SIZE((PyArrayObject*)self->final_src_fill);
184
+ dst_fill = PyArray_SIZE((PyArrayObject*)self->final_dst_fill);
185
+ }
186
+ else {
187
+ src_fill = -1;
188
+ dst_fill = -1;
189
+ }
190
+
191
+ return PyUnicode_FromFormat("<%s(len: %i, src_fill: %i, dst_fill: %i, is_many: %s, is_finalized: %s)>",
192
+ Py_TYPE(self)->tp_name,
193
+ self->len,
194
+ src_fill,
195
+ dst_fill,
196
+ is_many,
197
+ is_finalized);
198
+ }
199
+
200
+ // Provide the integer positions connecting the `src` to the `dst`. If there is no match to `src` or `dst`, the unmatched position can be provided with -1. From each side, a connection is documented to the current `len`. Each time this is called `len` is incremented, indicating the inrease in position in the `final`. Return NULL on error.
201
+
202
+ // Inner function for calling from C; returns 0 on success, -1 on error. Exceptions will be set on error.
203
+ static inline int
204
+ AK_TM_register_one(TriMapObject* tm, Py_ssize_t src_from, Py_ssize_t dst_from) {
205
+ bool src_matched = src_from >= 0;
206
+ bool dst_matched = dst_from >= 0;
207
+ if (src_from >= tm->src_len || dst_from >= tm->dst_len) {
208
+ PyErr_SetString(PyExc_ValueError, "Out of bounds locator");
209
+ return -1;
210
+ }
211
+ if (src_matched) {
212
+ if (AK_UNLIKELY(tm->src_one_count == tm->src_one_capacity)) {
213
+ tm->src_one_capacity <<= 1; // get 2x the capacity
214
+ tm->src_one = PyMem_Realloc(tm->src_one,
215
+ sizeof(TriMapOne) * tm->src_one_capacity);
216
+ if (tm->src_one == NULL) {
217
+ PyErr_SetNone(PyExc_MemoryError);
218
+ return -1;
219
+ }
220
+ }
221
+ tm->src_one[tm->src_one_count] = (TriMapOne){src_from, tm->len};
222
+ tm->src_one_count += 1;
223
+ }
224
+ if (dst_matched) {
225
+ if (AK_UNLIKELY(tm->dst_one_count == tm->dst_one_capacity)) {
226
+ tm->dst_one_capacity <<= 1; // get 2x the capacity
227
+ tm->dst_one = PyMem_Realloc(tm->dst_one,
228
+ sizeof(TriMapOne) * tm->dst_one_capacity);
229
+ if (tm->dst_one == NULL) {
230
+ PyErr_SetNone(PyExc_MemoryError);
231
+ return -1;
232
+ }
233
+ }
234
+ tm->dst_one[tm->dst_one_count] = (TriMapOne){dst_from, tm->len};
235
+ tm->dst_one_count += 1;
236
+ }
237
+ if (src_matched && dst_matched) {
238
+ if (!tm->is_many) {
239
+ // if we have seen this connection before, we have a many
240
+ if (tm->src_match_data[src_from] || tm->dst_match_data[dst_from]) {
241
+ tm->is_many = true;
242
+ }
243
+ }
244
+ tm->src_match_data[src_from] = NPY_TRUE;
245
+ tm->dst_match_data[dst_from] = NPY_TRUE;
246
+ }
247
+ tm->len += 1;
248
+ return 0;
249
+ }
250
+
251
+ // Public function for calling from Python.
252
+ PyObject *
253
+ TriMap_register_one(TriMapObject *self, PyObject *args) {
254
+ Py_ssize_t src_from;
255
+ Py_ssize_t dst_from;
256
+ if (!PyArg_ParseTuple(args,
257
+ "nn:register_one",
258
+ &src_from,
259
+ &dst_from)) {
260
+ return NULL;
261
+ }
262
+ if (self->finalized) {
263
+ PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization");
264
+ return NULL;
265
+ }
266
+ if (AK_TM_register_one(self, src_from, dst_from)) {
267
+ return NULL;
268
+ }
269
+ Py_RETURN_NONE;
270
+ }
271
+
272
+ PyObject *
273
+ TriMap_register_unmatched_dst(TriMapObject *self) {
274
+ if (self->finalized) {
275
+ PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization");
276
+ return NULL;
277
+ }
278
+ PyArrayObject* dst_match_array = (PyArrayObject *)self->dst_match;
279
+ PyObject* sum_scalar = PyArray_Sum(
280
+ dst_match_array,
281
+ 0,
282
+ NPY_INT64, // this converts before sum; not sure this is necessary
283
+ NULL);
284
+ if (sum_scalar == NULL) {
285
+ return NULL;
286
+ }
287
+ // for a 1D array PyArray_SUM returns a scalar
288
+ npy_int64 sum = PyArrayScalar_VAL(sum_scalar, Int64);
289
+ Py_DECREF(sum_scalar);
290
+
291
+ if (sum < self->dst_len) {
292
+ PyArrayObject* dst_unmatched = (PyArrayObject *)PyObject_CallMethod(
293
+ self->dst_match, // PyObject
294
+ "__invert__",
295
+ NULL);
296
+ if (dst_unmatched == NULL) {
297
+ return NULL;
298
+ }
299
+ // derive indices for unmatched locations, call each with register_one
300
+ PyArrayObject* indices = (PyArrayObject*)AK_nonzero_1d(dst_unmatched);
301
+ if (indices == NULL) {
302
+ Py_DECREF((PyObject*)dst_unmatched);
303
+ return NULL;
304
+ }
305
+ // borrow ref to array in 1-element tuple
306
+ npy_int64 *index_data = (npy_int64 *)PyArray_DATA(indices);
307
+ npy_intp index_len = PyArray_SIZE(indices);
308
+
309
+ for (npy_intp i = 0; i < index_len; i++) {
310
+ if (AK_TM_register_one(self, -1, index_data[i])) {
311
+ Py_DECREF((PyObject*)dst_unmatched);
312
+ Py_DECREF((PyObject*)indices);
313
+ return NULL;
314
+ }
315
+ }
316
+ Py_DECREF((PyObject*)dst_unmatched);
317
+ Py_DECREF((PyObject*)indices);
318
+ }
319
+ Py_RETURN_NONE;
320
+ }
321
+
322
+ // Given an integer (for the src) and an array of integers (for the dst), store mappings from src to final and dst to final.
323
+ PyObject *
324
+ TriMap_register_many(TriMapObject *self, PyObject *args) {
325
+ Py_ssize_t src_from;
326
+ PyArrayObject* dst_from;
327
+ if (!PyArg_ParseTuple(args,
328
+ "nO!:register_many",
329
+ &src_from,
330
+ &PyArray_Type, &dst_from)) {
331
+ return NULL;
332
+ }
333
+ if (self->finalized) {
334
+ PyErr_SetString(PyExc_RuntimeError, "Cannot register post finalization");
335
+ return NULL;
336
+ }
337
+ int dst_from_type = PyArray_TYPE(dst_from);
338
+ if (dst_from_type != NPY_INT64) {
339
+ PyErr_SetString(PyExc_ValueError, "`dst_from` must be a 64 bit integer array");
340
+ return NULL;
341
+ }
342
+ npy_intp increment = PyArray_SIZE(dst_from);
343
+
344
+ if (AK_UNLIKELY(self->many_count == self->many_capacity)) {
345
+ self->many_capacity <<= 1; // get 2x the capacity
346
+ self->many_to = PyMem_Realloc(self->many_to,
347
+ sizeof(TriMapManyTo) * self->many_capacity);
348
+ if (self->many_to == NULL) {
349
+ PyErr_SetNone(PyExc_MemoryError);
350
+ return NULL;
351
+ }
352
+ self->many_from = PyMem_Realloc(self->many_from,
353
+ sizeof(TriMapManyFrom) * self->many_capacity);
354
+ if (self->many_from == NULL) {
355
+ PyErr_SetNone(PyExc_MemoryError);
356
+ return NULL;
357
+ }
358
+ }
359
+ // define contiguous region in final to map to
360
+ self->many_to[self->many_count] = (TriMapManyTo){self->len, self->len + increment};
361
+
362
+ Py_INCREF((PyObject*)dst_from); // decrefs on dealloc
363
+ self->many_from[self->many_count] = (TriMapManyFrom){src_from, dst_from};
364
+ self->many_count += 1;
365
+
366
+ self->src_match_data[src_from] = NPY_TRUE;
367
+ // iterate over dst_from and set values to True; cannot assume that dst_from is contiguous; dst_match_data is contiguous
368
+ for (Py_ssize_t i = 0; i < increment; i++){
369
+ npy_int64 pos = *(npy_int64*)PyArray_GETPTR1(dst_from, i); // always int64
370
+ self->dst_match_data[pos] = NPY_TRUE;
371
+ }
372
+ self->len += increment;
373
+ self->is_many = true;
374
+ Py_RETURN_NONE;
375
+ }
376
+
377
+ //------------------------------------------------------------------------------
378
+ // Determine, for src and dst, which indices will need fill values, and store those indices as an integer array in final_src_fill, final_dst_fill
379
+ PyObject *
380
+ TriMap_finalize(TriMapObject *self, PyObject *Py_UNUSED(unused)) {
381
+ TriMapObject* tm = (TriMapObject*)self;
382
+
383
+ if (self->finalized) {
384
+ PyErr_SetString(PyExc_RuntimeError, "Cannot call finalize twice");
385
+ return NULL;
386
+ }
387
+ // predefine all PyObjects to use goto error
388
+ PyObject* final_src_match = NULL;
389
+ PyObject* final_dst_match = NULL;
390
+ PyObject* final_src_unmatched = NULL;
391
+ PyObject* final_dst_unmatched = NULL;
392
+
393
+ npy_intp dims[] = {tm->len};
394
+
395
+ // initialize all to False
396
+ final_src_match = PyArray_ZEROS(1, dims, NPY_BOOL, 0);
397
+ if (final_src_match == NULL) {
398
+ goto error;
399
+ }
400
+ final_dst_match = PyArray_ZEROS(1, dims, NPY_BOOL, 0);
401
+ if (final_dst_match == NULL) {
402
+ goto error;
403
+ }
404
+
405
+ npy_bool* final_src_match_data = (npy_bool*)PyArray_DATA(
406
+ (PyArrayObject*)final_src_match);
407
+ npy_bool* final_dst_match_data = (npy_bool*)PyArray_DATA(
408
+ (PyArrayObject*)final_dst_match);
409
+
410
+ TriMapOne* o;
411
+ TriMapOne* o_end;
412
+ o = tm->src_one;
413
+ o_end = o + tm->src_one_count;
414
+ for (; o < o_end; o++) {
415
+ final_src_match_data[o->to] = NPY_TRUE;
416
+ }
417
+ o = tm->dst_one;
418
+ o_end = o + tm->dst_one_count;
419
+ for (; o < o_end; o++) {
420
+ final_dst_match_data[o->to] = NPY_TRUE;
421
+ }
422
+ // many assign from src and dst into the same final positions
423
+ npy_bool* s;
424
+ npy_bool* d;
425
+ npy_bool* end;
426
+ TriMapManyTo* m = tm->many_to;
427
+ TriMapManyTo* m_end = m + tm->many_count;
428
+
429
+ for (; m < m_end; m++) {
430
+ d = final_dst_match_data + m->start;
431
+ s = final_src_match_data + m->start;
432
+ end = final_src_match_data + m->stop;
433
+ while (s < end) {
434
+ *s++ = NPY_TRUE;
435
+ *d++ = NPY_TRUE;
436
+ }
437
+ }
438
+ // NOTE: could sum first to see if nonzero call is necessary; would skip invert and nonzero calls
439
+ final_src_unmatched = PyObject_CallMethod(
440
+ final_src_match, // PyObject
441
+ "__invert__",
442
+ NULL);
443
+ if (final_src_unmatched == NULL) {
444
+ goto error;
445
+ }
446
+
447
+ final_dst_unmatched = PyObject_CallMethod(
448
+ final_dst_match, // PyObject
449
+ "__invert__",
450
+ NULL);
451
+ if (final_dst_unmatched == NULL) {
452
+ goto error;
453
+ }
454
+ tm->final_src_fill = AK_nonzero_1d((PyArrayObject*)final_src_unmatched);
455
+ if (tm->final_src_fill == NULL) {
456
+ goto error;
457
+ }
458
+ tm->final_dst_fill = AK_nonzero_1d((PyArrayObject*)final_dst_unmatched);
459
+ if (tm->final_dst_fill == NULL) {
460
+ goto error;
461
+ }
462
+ Py_DECREF(final_src_match);
463
+ Py_DECREF(final_dst_match);
464
+ Py_DECREF(final_src_unmatched);
465
+ Py_DECREF(final_dst_unmatched);
466
+
467
+ tm->finalized = true;
468
+ Py_RETURN_NONE;
469
+ error: // all PyObject initialized to NULL, no more than 1 ref
470
+ Py_XDECREF(final_src_match);
471
+ Py_XDECREF(final_dst_match);
472
+ Py_XDECREF(final_src_unmatched);
473
+ Py_XDECREF(final_dst_unmatched);
474
+ return NULL;
475
+ }
476
+
477
+ PyObject *
478
+ TriMap_is_many(TriMapObject *self, PyObject *Py_UNUSED(unused)) {
479
+ if (!self->finalized) {
480
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
481
+ return NULL;
482
+ }
483
+ if (self->is_many) {
484
+ Py_RETURN_TRUE;
485
+ }
486
+ Py_RETURN_FALSE;
487
+ }
488
+
489
+ // Return True if the `src` will not need a fill. This is only correct of `src` is binding to a left join or an inner join.
490
+ PyObject *
491
+ TriMap_src_no_fill(TriMapObject *self, PyObject *Py_UNUSED(unused)) {
492
+ if (!self->finalized) {
493
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
494
+ return NULL;
495
+ }
496
+ if (PyArray_SIZE((PyArrayObject*)self->final_src_fill) == 0) {
497
+ Py_RETURN_TRUE;
498
+ }
499
+ Py_RETURN_FALSE;
500
+ }
501
+
502
+ // Return True if the `dst` will not need a fill. This is only correct of `dst` is binding to a left join or an inner join.
503
+ PyObject *
504
+ TriMap_dst_no_fill(TriMapObject *self, PyObject *Py_UNUSED(unused)) {
505
+ if (!self->finalized) {
506
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
507
+ return NULL;
508
+ }
509
+ if (PyArray_SIZE((PyArrayObject*)self->final_dst_fill) == 0) {
510
+ Py_RETURN_TRUE;
511
+ }
512
+ Py_RETURN_FALSE;
513
+ }
514
+
515
+ # define AK_TM_TRANSFER_SCALAR(npy_type_to, npy_type_from) do { \
516
+ npy_type_to* array_to_data = (npy_type_to*)PyArray_DATA(array_to); \
517
+ TriMapOne* o = one_pairs; \
518
+ TriMapOne* o_end = o + one_count; \
519
+ for (; o < o_end; o++) { \
520
+ array_to_data[o->to] = (npy_type_to) \
521
+ *(npy_type_from*)PyArray_GETPTR1( \
522
+ array_from, o->from); \
523
+ } \
524
+ npy_type_to* t; \
525
+ npy_type_to* t_end; \
526
+ npy_type_to f; \
527
+ npy_int64 f_pos; \
528
+ npy_intp dst_pos; \
529
+ PyArrayObject* dst; \
530
+ for (Py_ssize_t i = 0; i < tm->many_count; i++) { \
531
+ t = array_to_data + tm->many_to[i].start; \
532
+ t_end = array_to_data + tm->many_to[i].stop; \
533
+ if (from_src) { \
534
+ f = (npy_type_to)*(npy_type_from*)PyArray_GETPTR1( \
535
+ array_from, tm->many_from[i].src); \
536
+ while (t < t_end) { \
537
+ *t++ = f; \
538
+ } \
539
+ } \
540
+ else { \
541
+ dst_pos = 0; \
542
+ dst = tm->many_from[i].dst; \
543
+ while (t < t_end) { \
544
+ f_pos = *(npy_int64*)PyArray_GETPTR1(dst, dst_pos); \
545
+ *t++ = (npy_type_to) \
546
+ *(npy_type_from*)PyArray_GETPTR1( \
547
+ array_from, f_pos); \
548
+ dst_pos++; \
549
+ } \
550
+ } \
551
+ } \
552
+ } while (0) \
553
+
554
+ // Based on `tm` state, transfer from src or from dst (depending on `from_src`) to a `array_to`, a newly created contiguous array that is compatible with the values in `array_from`. Returns -1 on error. This only needs to match to / from type combinations that are possible from `resolve_dtype`, i.e., bool never goes to integer.
555
+ static inline int
556
+ AK_TM_transfer_scalar(TriMapObject* tm,
557
+ bool from_src,
558
+ PyArrayObject* array_from,
559
+ PyArrayObject* array_to) {
560
+ Py_ssize_t one_count = from_src ? tm->src_one_count : tm->dst_one_count;
561
+ TriMapOne* one_pairs = from_src ? tm->src_one : tm->dst_one;
562
+
563
+ switch(PyArray_TYPE(array_to)){
564
+ case NPY_BOOL:
565
+ AK_TM_TRANSFER_SCALAR(npy_bool, npy_bool);
566
+ return 0;
567
+ case NPY_INT64:
568
+ switch (PyArray_TYPE(array_from)) {
569
+ case NPY_INT64:
570
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_int64);
571
+ return 0;
572
+ case NPY_INT32:
573
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_int32);
574
+ return 0;
575
+ case NPY_INT16:
576
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_int16);
577
+ return 0;
578
+ case NPY_INT8:
579
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_int8);
580
+ return 0;
581
+ case NPY_UINT32:
582
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_uint32);
583
+ return 0;
584
+ case NPY_UINT16:
585
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_uint16);
586
+ return 0;
587
+ case NPY_UINT8:
588
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_uint8);
589
+ return 0;
590
+ }
591
+ break;
592
+ case NPY_INT32:
593
+ switch (PyArray_TYPE(array_from)) {
594
+ case NPY_INT32:
595
+ AK_TM_TRANSFER_SCALAR(npy_int32, npy_int32);
596
+ return 0;
597
+ case NPY_INT16:
598
+ AK_TM_TRANSFER_SCALAR(npy_int32, npy_int16);
599
+ return 0;
600
+ case NPY_INT8:
601
+ AK_TM_TRANSFER_SCALAR(npy_int32, npy_int8);
602
+ return 0;
603
+ case NPY_UINT16:
604
+ AK_TM_TRANSFER_SCALAR(npy_int32, npy_uint16);
605
+ return 0;
606
+ case NPY_UINT8:
607
+ AK_TM_TRANSFER_SCALAR(npy_int32, npy_uint8);
608
+ return 0;
609
+ }
610
+ break;
611
+ case NPY_INT16:
612
+ switch (PyArray_TYPE(array_from)) {
613
+ case NPY_INT16:
614
+ AK_TM_TRANSFER_SCALAR(npy_int16, npy_int16);
615
+ return 0;
616
+ case NPY_INT8:
617
+ AK_TM_TRANSFER_SCALAR(npy_int16, npy_int8);
618
+ return 0;
619
+ case NPY_UINT8:
620
+ AK_TM_TRANSFER_SCALAR(npy_int16, npy_uint8);
621
+ return 0;
622
+ }
623
+ break;
624
+ case NPY_INT8:
625
+ AK_TM_TRANSFER_SCALAR(npy_int8, npy_int8);
626
+ return 0;
627
+ case NPY_UINT64:
628
+ switch (PyArray_TYPE(array_from)) {
629
+ case NPY_UINT64:
630
+ AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint64);
631
+ return 0;
632
+ case NPY_UINT32:
633
+ AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint32);
634
+ return 0;
635
+ case NPY_UINT16:
636
+ AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint16);
637
+ return 0;
638
+ case NPY_UINT8:
639
+ AK_TM_TRANSFER_SCALAR(npy_uint64, npy_uint8);
640
+ return 0;
641
+ }
642
+ break;
643
+ case NPY_UINT32:
644
+ switch (PyArray_TYPE(array_from)) {
645
+ case NPY_UINT32:
646
+ AK_TM_TRANSFER_SCALAR(npy_uint32, npy_uint32);
647
+ return 0;
648
+ case NPY_UINT16:
649
+ AK_TM_TRANSFER_SCALAR(npy_uint32, npy_uint16);
650
+ return 0;
651
+ case NPY_UINT8:
652
+ AK_TM_TRANSFER_SCALAR(npy_uint32, npy_uint8);
653
+ return 0;
654
+ }
655
+ break;
656
+ case NPY_UINT16:
657
+ switch (PyArray_TYPE(array_from)) {
658
+ case NPY_UINT16:
659
+ AK_TM_TRANSFER_SCALAR(npy_uint16, npy_uint16);
660
+ return 0;
661
+ case NPY_UINT8:
662
+ AK_TM_TRANSFER_SCALAR(npy_uint16, npy_uint8);
663
+ return 0;
664
+ }
665
+ break;
666
+ case NPY_UINT8:
667
+ AK_TM_TRANSFER_SCALAR(npy_uint8, npy_uint8);
668
+ return 0;
669
+ case NPY_FLOAT64:
670
+ switch (PyArray_TYPE(array_from)) {
671
+ case NPY_FLOAT64:
672
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_float64);
673
+ return 0;
674
+ case NPY_FLOAT32:
675
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_float32);
676
+ return 0;
677
+ case NPY_FLOAT16:
678
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_float16);
679
+ return 0;
680
+ case NPY_INT64:
681
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_int64);
682
+ return 0;
683
+ case NPY_INT32:
684
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_int32);
685
+ return 0;
686
+ case NPY_INT16:
687
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_int16);
688
+ return 0;
689
+ case NPY_INT8:
690
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_int8);
691
+ return 0;
692
+ case NPY_UINT64:
693
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint64);
694
+ return 0;
695
+ case NPY_UINT32:
696
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint32);
697
+ return 0;
698
+ case NPY_UINT16:
699
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint16);
700
+ return 0;
701
+ case NPY_UINT8:
702
+ AK_TM_TRANSFER_SCALAR(npy_float64, npy_uint8);
703
+ return 0;
704
+ }
705
+ break;
706
+ case NPY_FLOAT32:
707
+ switch (PyArray_TYPE(array_from)) {
708
+ case NPY_FLOAT32:
709
+ AK_TM_TRANSFER_SCALAR(npy_float32, npy_float32);
710
+ return 0;
711
+ case NPY_FLOAT16:
712
+ AK_TM_TRANSFER_SCALAR(npy_float32, npy_float16);
713
+ return 0;
714
+ case NPY_INT16:
715
+ AK_TM_TRANSFER_SCALAR(npy_float32, npy_int16);
716
+ return 0;
717
+ case NPY_INT8:
718
+ AK_TM_TRANSFER_SCALAR(npy_float32, npy_int8);
719
+ return 0;
720
+ case NPY_UINT16:
721
+ AK_TM_TRANSFER_SCALAR(npy_float32, npy_uint16);
722
+ return 0;
723
+ case NPY_UINT8:
724
+ AK_TM_TRANSFER_SCALAR(npy_float32, npy_uint8);
725
+ return 0;
726
+ }
727
+ break;
728
+ case NPY_FLOAT16:
729
+ switch (PyArray_TYPE(array_from)) {
730
+ case NPY_FLOAT16:
731
+ AK_TM_TRANSFER_SCALAR(npy_float16, npy_float16);
732
+ return 0;
733
+ case NPY_INT8:
734
+ AK_TM_TRANSFER_SCALAR(npy_float16, npy_int8);
735
+ return 0;
736
+ case NPY_UINT16:
737
+ AK_TM_TRANSFER_SCALAR(npy_float16, npy_uint16);
738
+ return 0;
739
+ case NPY_UINT8:
740
+ AK_TM_TRANSFER_SCALAR(npy_float16, npy_uint8);
741
+ return 0;
742
+ }
743
+ break;
744
+ case NPY_DATETIME: {
745
+ AK_TM_TRANSFER_SCALAR(npy_int64, npy_int64);
746
+ return 0;
747
+ }
748
+ }
749
+ PyErr_SetString(PyExc_TypeError, "No handling for types");
750
+ return -1;
751
+ }
752
+ #undef AK_TM_TRANSFER_SCALAR
753
+
754
+ // Returns -1 on error. Specialized transfer from any type of an array to an object array.
755
+ static inline int
756
+ AK_TM_transfer_object(TriMapObject* tm,
757
+ bool from_src,
758
+ PyArrayObject* array_from,
759
+ PyArrayObject* array_to
760
+ ) {
761
+ Py_ssize_t one_count = from_src ? tm->src_one_count : tm->dst_one_count;
762
+ TriMapOne* one_pairs = from_src ? tm->src_one : tm->dst_one;
763
+
764
+ // NOTE: could use PyArray_Scalar instead of PyArray_GETITEM if we wanted to store scalars instead of Python objects; however, that is pretty uncommon for object arrays to store PyArray_Scalars
765
+ bool f_is_obj = PyArray_TYPE(array_from) == NPY_OBJECT;
766
+
767
+ // the passed in object array is contiguous and have NULL (not None) in each position
768
+ PyObject** array_to_data = (PyObject**)PyArray_DATA(array_to);
769
+ PyObject* pyo;
770
+ void* f;
771
+ TriMapOne* o = one_pairs;
772
+ TriMapOne* o_end = o + one_count;
773
+ for (; o < o_end; o++) {
774
+ f = PyArray_GETPTR1(array_from, o->from);
775
+ if (f_is_obj) {
776
+ pyo = *(PyObject**)f;
777
+ Py_INCREF(pyo);
778
+ }
779
+ else { // will convert any value to an object
780
+ pyo = PyArray_GETITEM(array_from, f);
781
+ }
782
+ array_to_data[o->to] = pyo;
783
+ }
784
+ PyObject** t;
785
+ PyObject** t_end;
786
+ npy_intp dst_pos;
787
+ npy_int64 f_pos;
788
+ PyArrayObject* dst;
789
+ for (Py_ssize_t i = 0; i < tm->many_count; i++) {
790
+ t = array_to_data + tm->many_to[i].start;
791
+ t_end = array_to_data + tm->many_to[i].stop;
792
+
793
+ if (from_src) {
794
+ f = PyArray_GETPTR1(array_from, tm->many_from[i].src);
795
+ if (f_is_obj) {
796
+ pyo = *(PyObject**)f;
797
+ Py_INCREF(pyo); // pre add new ref so equal to PyArray_GETITEM
798
+ }
799
+ else {
800
+ pyo = PyArray_GETITEM(array_from, f); // given a new ref
801
+ }
802
+ while (t < t_end) {
803
+ Py_INCREF(pyo); // one more than we need
804
+ *t++ = pyo;
805
+ }
806
+ Py_DECREF(pyo); // remove the extra ref
807
+ }
808
+ else { // from_dst, dst is an array
809
+ dst_pos = 0;
810
+ dst = tm->many_from[i].dst;
811
+ while (t < t_end) {
812
+ f_pos = *(npy_int64*)PyArray_GETPTR1(dst, dst_pos);
813
+ f = PyArray_GETPTR1(array_from, f_pos);
814
+ if (f_is_obj) {
815
+ pyo = *(PyObject**)f;
816
+ Py_INCREF(pyo);
817
+ }
818
+ else {
819
+ pyo = PyArray_GETITEM(array_from, f);
820
+ }
821
+ *t++ = pyo;
822
+ dst_pos++;
823
+ }
824
+ }
825
+ }
826
+ return 0;
827
+ }
828
+
829
+ // Returns -1 on error. Specialized transfer from any type of an array to an object array. For usage with merge, Will only transfer if the destination is not NULL.
830
+ static inline int
831
+ AK_TM_transfer_object_if_null(TriMapObject* tm,
832
+ bool from_src,
833
+ PyArrayObject* array_from,
834
+ PyArrayObject* array_to
835
+ ) {
836
+ Py_ssize_t one_count = from_src ? tm->src_one_count : tm->dst_one_count;
837
+ TriMapOne* one_pairs = from_src ? tm->src_one : tm->dst_one;
838
+
839
+ // NOTE: could use PyArray_Scalar instead of PyArray_GETITEM if we wanted to store scalars instead of Python objects; however, that is pretty uncommon for object arrays to store PyArray_Scalars
840
+ bool f_is_obj = PyArray_TYPE(array_from) == NPY_OBJECT;
841
+
842
+ // the passed in object array is contiguous and have NULL (not None) in each position
843
+ PyObject** array_to_data = (PyObject**)PyArray_DATA(array_to);
844
+ PyObject* pyo;
845
+ void* f;
846
+ TriMapOne* o = one_pairs;
847
+ TriMapOne* o_end = o + one_count;
848
+ for (; o < o_end; o++) {
849
+ if (array_to_data[o->to] == NULL) {
850
+ f = PyArray_GETPTR1(array_from, o->from);
851
+ if (f_is_obj) {
852
+ pyo = *(PyObject**)f;
853
+ Py_INCREF(pyo);
854
+ }
855
+ else { // will convert any value to an object
856
+ pyo = PyArray_GETITEM(array_from, f);
857
+ }
858
+ array_to_data[o->to] = pyo;
859
+ }
860
+ }
861
+ PyObject** t;
862
+ PyObject** t_end;
863
+ npy_intp dst_pos;
864
+ npy_int64 f_pos;
865
+ PyArrayObject* dst;
866
+ for (Py_ssize_t i = 0; i < tm->many_count; i++) {
867
+ t = array_to_data + tm->many_to[i].start;
868
+ t_end = array_to_data + tm->many_to[i].stop;
869
+
870
+ if (from_src) {
871
+ while (t < t_end) {
872
+ if (*t == NULL) {
873
+ f = PyArray_GETPTR1(array_from, tm->many_from[i].src);
874
+ if (f_is_obj) {
875
+ pyo = *(PyObject**)f;
876
+ Py_INCREF(pyo);
877
+ }
878
+ else {
879
+ pyo = PyArray_GETITEM(array_from, f); // given a new ref
880
+ }
881
+ *t++ = pyo;
882
+ }
883
+ else {
884
+ t++;
885
+ }
886
+ }
887
+ }
888
+ else { // from_dst, dst is an array
889
+ dst_pos = 0;
890
+ dst = tm->many_from[i].dst;
891
+ while (t < t_end) {
892
+ if (*t == NULL) {
893
+ f_pos = *(npy_int64*)PyArray_GETPTR1(dst, dst_pos);
894
+ f = PyArray_GETPTR1(array_from, f_pos);
895
+ if (f_is_obj) {
896
+ pyo = *(PyObject**)f;
897
+ Py_INCREF(pyo);
898
+ }
899
+ else {
900
+ pyo = PyArray_GETITEM(array_from, f);
901
+ }
902
+ *t++ = pyo;
903
+ dst_pos++;
904
+ }
905
+ else {
906
+ t++;
907
+ dst_pos++;
908
+ }
909
+ }
910
+ }
911
+ }
912
+ return 0;
913
+ }
914
+
915
+ // Returns -1 on error.
916
+ static inline int
917
+ AK_TM_fill_object(TriMapObject* tm,
918
+ bool from_src,
919
+ PyArrayObject* array_to,
920
+ PyObject* fill_value) {
921
+
922
+ PyArrayObject* final_fill = (PyArrayObject*)(from_src
923
+ ? tm->final_src_fill : tm->final_dst_fill);
924
+ PyObject** array_to_data = (PyObject**)PyArray_DATA(array_to);
925
+ npy_int64* p = (npy_int64*)PyArray_DATA(final_fill);
926
+ npy_int64* p_end = p + PyArray_SIZE(final_fill);
927
+ PyObject** target;
928
+ while (p < p_end) {
929
+ target = array_to_data + *p++;
930
+ Py_INCREF(fill_value);
931
+ *target = fill_value;
932
+ }
933
+ return 0;
934
+ }
935
+
936
+ #define AK_TM_TRANSFER_FLEXIBLE(c_type, from_src, array_from, array_to) do {\
937
+ Py_ssize_t one_count = from_src ? tm->src_one_count : tm->dst_one_count;\
938
+ TriMapOne* one_pairs = from_src ? tm->src_one : tm->dst_one; \
939
+ npy_intp t_element_size = PyArray_ITEMSIZE(array_to); \
940
+ npy_intp t_element_cp = t_element_size / sizeof(c_type); \
941
+ npy_intp f_element_size = PyArray_ITEMSIZE(array_from); \
942
+ c_type* array_to_data = (c_type*)PyArray_DATA(array_to); \
943
+ c_type* f; \
944
+ c_type* t; \
945
+ c_type* t_end; \
946
+ npy_intp dst_pos; \
947
+ npy_int64 f_pos; \
948
+ PyArrayObject* dst; \
949
+ TriMapOne* o = one_pairs; \
950
+ TriMapOne* o_end = o + one_count; \
951
+ for (; o < o_end; o++) { \
952
+ f = (c_type*)PyArray_GETPTR1(array_from, o->from); \
953
+ t = array_to_data + t_element_cp * o->to; \
954
+ memcpy(t, f, f_element_size); \
955
+ } \
956
+ for (Py_ssize_t i = 0; i < tm->many_count; i++) { \
957
+ t = array_to_data + t_element_cp * tm->many_to[i].start; \
958
+ t_end = array_to_data + t_element_cp * tm->many_to[i].stop; \
959
+ if (from_src) { \
960
+ f = (c_type*)PyArray_GETPTR1(array_from, tm->many_from[i].src);\
961
+ for (; t < t_end; t += t_element_cp) { \
962
+ memcpy(t, f, f_element_size); \
963
+ } \
964
+ } \
965
+ else { \
966
+ dst_pos = 0; \
967
+ dst = tm->many_from[i].dst; \
968
+ for (; t < t_end; t += t_element_cp) { \
969
+ f_pos = *(npy_int64*)PyArray_GETPTR1(dst, dst_pos); \
970
+ f = (c_type*)PyArray_GETPTR1(array_from, f_pos); \
971
+ memcpy(t, f, f_element_size); \
972
+ dst_pos++; \
973
+ } \
974
+ } \
975
+ } \
976
+ } while (0) \
977
+
978
+ // Returns -1 on error.
979
+ static inline int
980
+ AK_TM_fill_unicode(TriMapObject* tm,
981
+ bool from_src,
982
+ PyArrayObject* array_to,
983
+ PyObject* fill_value) {
984
+ PyArrayObject* final_fill = (PyArrayObject*)(from_src
985
+ ? tm->final_src_fill : tm->final_dst_fill);
986
+
987
+ Py_UCS4* array_to_data = (Py_UCS4*)PyArray_DATA(array_to);
988
+ // code points per element
989
+ npy_intp cp = PyArray_ITEMSIZE(array_to) / UCS4_SIZE;
990
+
991
+ bool decref_fill_value = false;
992
+ if (PyBytes_Check(fill_value)) {
993
+ fill_value = PyUnicode_FromEncodedObject(fill_value, "utf-8", NULL);
994
+ if (fill_value == NULL) {
995
+ return -1;
996
+ }
997
+ decref_fill_value = true;
998
+ }
999
+ else if (!PyUnicode_Check(fill_value)) {
1000
+ return -1;
1001
+ }
1002
+ Py_ssize_t fill_cp = PyUnicode_GET_LENGTH(fill_value) * UCS4_SIZE; // code points
1003
+ // p is the index position to fill
1004
+ npy_int64* p = (npy_int64*)PyArray_DATA(final_fill);
1005
+ npy_int64* p_end = p + PyArray_SIZE(final_fill);
1006
+ Py_UCS4* target;
1007
+ while (p < p_end) {
1008
+ target = array_to_data + (*p * cp);
1009
+ // disabling copying a null
1010
+ if (PyUnicode_AsUCS4(fill_value, target, fill_cp, 0) == NULL) {
1011
+ return -1;
1012
+ }
1013
+ p++;
1014
+ }
1015
+ if (decref_fill_value) {
1016
+ Py_DECREF(fill_value);
1017
+ }
1018
+ return 0;
1019
+ }
1020
+
1021
+ // Returns -1 on error.
1022
+ static inline int
1023
+ AK_TM_fill_string(TriMapObject* tm,
1024
+ bool from_src,
1025
+ PyArrayObject* array_to,
1026
+ PyObject* fill_value) {
1027
+ PyArrayObject* final_fill = (PyArrayObject*)(from_src
1028
+ ? tm->final_src_fill : tm->final_dst_fill);
1029
+
1030
+ char* array_to_data = (char*)PyArray_DATA(array_to);
1031
+ npy_intp cp = PyArray_ITEMSIZE(array_to);
1032
+ if (!PyBytes_Check(fill_value)) {
1033
+ return -1;
1034
+ }
1035
+ Py_ssize_t fill_cp = PyBytes_GET_SIZE(fill_value);
1036
+ const char* fill_data = PyBytes_AS_STRING(fill_value);
1037
+ // p is the index position to fill
1038
+ npy_int64* p = (npy_int64*)PyArray_DATA(final_fill);
1039
+ npy_int64* p_end = p + PyArray_SIZE(final_fill);
1040
+ char* target;
1041
+ while (p < p_end) {
1042
+ target = array_to_data + (*p++ * cp);
1043
+ memcpy(target, fill_data, fill_cp);
1044
+ }
1045
+ return 0;
1046
+ }
1047
+
1048
+ // Returns NULL on error.
1049
+ static inline PyObject *
1050
+ AK_TM_map_no_fill(TriMapObject* tm,
1051
+ bool from_src,
1052
+ PyArrayObject* array_from) {
1053
+ if (!(PyArray_NDIM(array_from) == 1)) {
1054
+ PyErr_SetString(PyExc_TypeError, "Array must be 1D");
1055
+ return NULL;
1056
+ }
1057
+ npy_intp dims[] = {tm->len};
1058
+ PyArrayObject* array_to;
1059
+ bool dtype_is_obj = PyArray_TYPE(array_from) == NPY_OBJECT;
1060
+ bool dtype_is_unicode = PyArray_TYPE(array_from) == NPY_UNICODE;
1061
+ bool dtype_is_string = PyArray_TYPE(array_from) == NPY_STRING;
1062
+
1063
+ // create to array
1064
+ if (dtype_is_obj) { // initializes values to NULL
1065
+ array_to = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_OBJECT);
1066
+ }
1067
+ else {
1068
+ PyArray_Descr* dtype = PyArray_DESCR(array_from); // borowed ref
1069
+ Py_INCREF(dtype);
1070
+ array_to = (PyArrayObject*)PyArray_Empty(1, dims, dtype, 0); // steals dtype ref
1071
+ }
1072
+ if (array_to == NULL) {
1073
+ PyErr_SetNone(PyExc_MemoryError);
1074
+ return NULL;
1075
+ }
1076
+ // transfer values
1077
+ if (dtype_is_obj) {
1078
+ if (AK_TM_transfer_object(tm, from_src, array_from, array_to)) {
1079
+ Py_DECREF((PyObject*)array_to);
1080
+ return NULL;
1081
+ }
1082
+ }
1083
+ else if (dtype_is_unicode) {
1084
+ AK_TM_TRANSFER_FLEXIBLE(Py_UCS4, from_src, array_from, array_to);
1085
+ }
1086
+ else if (dtype_is_string) {
1087
+ AK_TM_TRANSFER_FLEXIBLE(char, from_src, array_from, array_to);
1088
+ }
1089
+ else {
1090
+ if (AK_TM_transfer_scalar(tm, from_src, array_from, array_to)) {
1091
+ Py_DECREF((PyObject*)array_to);
1092
+ return NULL;
1093
+ }
1094
+ }
1095
+ PyArray_CLEARFLAGS(array_to, NPY_ARRAY_WRITEABLE);
1096
+ return (PyObject*)array_to;
1097
+ }
1098
+
1099
+ PyObject *
1100
+ TriMap_map_src_no_fill(TriMapObject *self, PyObject *arg) {
1101
+ if (!PyArray_Check(arg)) {
1102
+ PyErr_SetString(PyExc_TypeError, "Must provide an array");
1103
+ return NULL;
1104
+ }
1105
+ if (!self->finalized) {
1106
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
1107
+ return NULL;
1108
+ }
1109
+ PyArrayObject* array_from = (PyArrayObject*)arg;
1110
+ bool from_src = true;
1111
+ return AK_TM_map_no_fill(self, from_src, array_from);
1112
+ }
1113
+
1114
+ PyObject *
1115
+ TriMap_map_dst_no_fill(TriMapObject *self, PyObject *arg) {
1116
+ if (!PyArray_Check(arg)) {
1117
+ PyErr_SetString(PyExc_TypeError, "Must provide an array");
1118
+ return NULL;
1119
+ }
1120
+ if (!self->finalized) {
1121
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
1122
+ return NULL;
1123
+ }
1124
+ PyArrayObject* array_from = (PyArrayObject*)arg;
1125
+ bool from_src = false;
1126
+ return AK_TM_map_no_fill(self, from_src, array_from);
1127
+ }
1128
+
1129
+ static inline PyObject *
1130
+ TriMap_map_merge(TriMapObject *tm, PyObject *args)
1131
+ {
1132
+ // both are "from_" arrays
1133
+ PyArrayObject* array_src;
1134
+ PyArrayObject* array_dst;
1135
+
1136
+ if (!PyArg_ParseTuple(args,
1137
+ "O!O!:map_merge",
1138
+ &PyArray_Type, &array_src,
1139
+ &PyArray_Type, &array_dst
1140
+ )) {
1141
+ return NULL;
1142
+ }
1143
+ if (!tm->finalized) {
1144
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
1145
+ return NULL;
1146
+ }
1147
+ if (!(PyArray_NDIM(array_src) == 1)) {
1148
+ PyErr_SetString(PyExc_TypeError, "Array src must be 1D");
1149
+ return NULL;
1150
+ }
1151
+ if (!(PyArray_NDIM(array_dst) == 1)) {
1152
+ PyErr_SetString(PyExc_TypeError, "Array dst must be 1D");
1153
+ return NULL;
1154
+ }
1155
+ // passing a borrowed refs; returns a new ref
1156
+ PyArray_Descr* dtype = AK_resolve_dtype(
1157
+ PyArray_DESCR(array_src),
1158
+ PyArray_DESCR(array_dst));
1159
+ bool dtype_is_obj = dtype->type_num == NPY_OBJECT;
1160
+ bool dtype_is_unicode = dtype->type_num == NPY_UNICODE;
1161
+ bool dtype_is_string = dtype->type_num == NPY_STRING;
1162
+
1163
+ npy_intp dims[] = {tm->len};
1164
+
1165
+ // create to array_to
1166
+ PyArrayObject* array_to;
1167
+ if (dtype_is_obj) {
1168
+ Py_DECREF(dtype); // not needed
1169
+ // will initialize to NULL, not None
1170
+ array_to = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_OBJECT);
1171
+ }
1172
+ else if (dtype_is_unicode || dtype_is_string) {
1173
+ array_to = (PyArrayObject*)PyArray_Zeros(1, dims, dtype, 0); // steals dtype ref
1174
+ }
1175
+ else {
1176
+ array_to = (PyArrayObject*)PyArray_Empty(1, dims, dtype, 0); // steals dtype ref
1177
+ }
1178
+ if (array_to == NULL) {
1179
+ PyErr_SetNone(PyExc_MemoryError);
1180
+ return NULL;
1181
+ }
1182
+
1183
+ // if we have fill values in src, we need to transfer from dst
1184
+ bool transfer_from_dst = PyArray_SIZE((PyArrayObject*)tm->final_src_fill) != 0;
1185
+
1186
+ if (dtype_is_obj) {
1187
+ if (AK_TM_transfer_object(tm, true, array_src, array_to)) {
1188
+ Py_DECREF((PyObject*)array_to);
1189
+ return NULL;
1190
+ }
1191
+ if (transfer_from_dst) {
1192
+ if (AK_TM_transfer_object_if_null(tm, false, array_dst, array_to)) {
1193
+ Py_DECREF((PyObject*)array_to);
1194
+ return NULL;
1195
+ }
1196
+ }
1197
+ }
1198
+ else if (dtype_is_unicode) {
1199
+ AK_TM_TRANSFER_FLEXIBLE(Py_UCS4, true, array_src, array_to);
1200
+ if (transfer_from_dst) {
1201
+ AK_TM_TRANSFER_FLEXIBLE(Py_UCS4, false, array_dst, array_to);
1202
+ }
1203
+ }
1204
+ else if (dtype_is_string) {
1205
+ AK_TM_TRANSFER_FLEXIBLE(char, true, array_src, array_to);
1206
+ if (transfer_from_dst) {
1207
+ AK_TM_TRANSFER_FLEXIBLE(char, false, array_dst, array_to);
1208
+ }
1209
+ }
1210
+ else {
1211
+ if (AK_TM_transfer_scalar(tm, true, array_src, array_to)) {
1212
+ Py_DECREF((PyObject*)array_to);
1213
+ return NULL;
1214
+ }
1215
+ if (transfer_from_dst) {
1216
+ if (AK_TM_transfer_scalar(tm, false, array_dst, array_to)) {
1217
+ Py_DECREF((PyObject*)array_to);
1218
+ return NULL;
1219
+ }
1220
+ }
1221
+ }
1222
+ return (PyObject*)array_to;
1223
+ }
1224
+
1225
+ // Returns NULL on error.
1226
+ static inline PyObject *
1227
+ AK_TM_map_fill(TriMapObject* tm,
1228
+ bool from_src,
1229
+ PyArrayObject* array_from,
1230
+ PyObject* fill_value,
1231
+ PyArray_Descr* fill_value_dtype) {
1232
+ if (!(PyArray_NDIM(array_from) == 1)) {
1233
+ PyErr_SetString(PyExc_TypeError, "Array must be 1D");
1234
+ return NULL;
1235
+ }
1236
+ // passing a borrowed ref; returns a new ref
1237
+ PyArray_Descr* dtype = AK_resolve_dtype(PyArray_DESCR(array_from), fill_value_dtype);
1238
+ bool dtype_is_obj = dtype->type_num == NPY_OBJECT;
1239
+ bool dtype_is_unicode = dtype->type_num == NPY_UNICODE;
1240
+ bool dtype_is_string = dtype->type_num == NPY_STRING;
1241
+
1242
+ npy_intp dims[] = {tm->len};
1243
+ PyArrayObject* array_to;
1244
+
1245
+ if (dtype_is_obj) {
1246
+ Py_DECREF(dtype); // not needed
1247
+ // will initialize to NULL, not None
1248
+ array_to = (PyArrayObject*)PyArray_SimpleNew(1, dims, NPY_OBJECT);
1249
+ Py_INCREF(array_from); // normalize refs when casting
1250
+ }
1251
+ else if (dtype_is_unicode || dtype_is_string) {
1252
+ array_to = (PyArrayObject*)PyArray_Zeros(1, dims, dtype, 0); // steals dtype ref
1253
+ Py_INCREF(array_from); // normalize refs when casting
1254
+ }
1255
+ else {
1256
+ array_to = (PyArrayObject*)PyArray_Empty(1, dims, dtype, 0); // steals dtype ref
1257
+ if (PyArray_TYPE(array_from) == NPY_DATETIME &&
1258
+ PyArray_TYPE(array_to) == NPY_DATETIME &&
1259
+ AK_dt_unit_from_array(array_from) != AK_dt_unit_from_array(array_to)
1260
+ ) {
1261
+ // if trying to cast into a dt64 array, need to pre-convert; array_from is originally borrowed; calling cast sets it to a new ref
1262
+ dtype = PyArray_DESCR(array_to); // borrowed ref
1263
+ Py_INCREF(dtype);
1264
+ array_from = (PyArrayObject*)PyArray_CastToType(array_from, dtype, 0);
1265
+ }
1266
+ else {
1267
+ Py_INCREF(array_from); // normalize refs when casting
1268
+ }
1269
+ }
1270
+ if (array_to == NULL) {
1271
+ PyErr_SetNone(PyExc_MemoryError);
1272
+ Py_DECREF((PyObject*)array_from);
1273
+ return NULL;
1274
+ }
1275
+ // array_from, array_to inc refed and dec refed on error
1276
+ if (dtype_is_obj) {
1277
+ if (AK_TM_transfer_object(tm, from_src, array_from, array_to)) {
1278
+ goto error;
1279
+ }
1280
+ if (AK_TM_fill_object(tm, from_src, array_to, fill_value)) {
1281
+ goto error;
1282
+ }
1283
+ }
1284
+ else if (dtype_is_unicode) {
1285
+ AK_TM_TRANSFER_FLEXIBLE(Py_UCS4, from_src, array_from, array_to);
1286
+ if (AK_TM_fill_unicode(tm, from_src, array_to, fill_value)) {
1287
+ goto error;
1288
+ }
1289
+ }
1290
+ else if (dtype_is_string) {
1291
+ AK_TM_TRANSFER_FLEXIBLE(char, from_src, array_from, array_to);
1292
+ if (AK_TM_fill_string(tm, from_src, array_to, fill_value)) {
1293
+ goto error;
1294
+ }
1295
+ }
1296
+ else {
1297
+ // Most simple is to fill with scalar, then overwrite values as needed; for object and flexible dtypes this is not efficient; for object dtypes, this obligates us to decref the filled value when assigning
1298
+ if (PyArray_FillWithScalar(array_to, fill_value)) { // -1 on error
1299
+ goto error;
1300
+ }
1301
+ if (AK_TM_transfer_scalar(tm, from_src, array_from, array_to)) {
1302
+ goto error;
1303
+ }
1304
+ }
1305
+ Py_DECREF((PyObject*)array_from); // ref inc for this function
1306
+ PyArray_CLEARFLAGS(array_to, NPY_ARRAY_WRITEABLE);
1307
+ return (PyObject*)array_to;
1308
+ error:
1309
+ Py_DECREF((PyObject*)array_to);
1310
+ Py_DECREF((PyObject*)array_from);
1311
+ return NULL;
1312
+ }
1313
+ #undef AK_TM_TRANSFER_FLEXIBLE
1314
+
1315
+ PyObject *
1316
+ TriMap_map_src_fill(TriMapObject *self, PyObject *args) {
1317
+ PyArrayObject* array_from;
1318
+ PyObject* fill_value;
1319
+ PyArray_Descr* fill_value_dtype;
1320
+ if (!PyArg_ParseTuple(args,
1321
+ "O!OO!:map_src_fill",
1322
+ &PyArray_Type, &array_from,
1323
+ &fill_value,
1324
+ &PyArrayDescr_Type, &fill_value_dtype
1325
+ )) {
1326
+ return NULL;
1327
+ }
1328
+ if (!self->finalized) {
1329
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
1330
+ return NULL;
1331
+ }
1332
+ bool from_src = true;
1333
+ return AK_TM_map_fill(self, from_src, array_from, fill_value, fill_value_dtype);
1334
+ }
1335
+
1336
+ PyObject *
1337
+ TriMap_map_dst_fill(TriMapObject *self, PyObject *args) {
1338
+ PyArrayObject* array_from;
1339
+ PyObject* fill_value;
1340
+ PyArray_Descr* fill_value_dtype;
1341
+ if (!PyArg_ParseTuple(args,
1342
+ "O!OO!:map_dst_fill",
1343
+ &PyArray_Type, &array_from,
1344
+ &fill_value,
1345
+ &PyArrayDescr_Type, &fill_value_dtype
1346
+ )) {
1347
+ return NULL;
1348
+ }
1349
+ if (!self->finalized) {
1350
+ PyErr_SetString(PyExc_RuntimeError, "Finalization is required");
1351
+ return NULL;
1352
+ }
1353
+ bool from_src = false;
1354
+ return AK_TM_map_fill(self, from_src, array_from, fill_value, fill_value_dtype);
1355
+ }
1356
+
1357
+
1358
+
1359
+ static PyMethodDef TriMap_methods[] = {
1360
+ {"register_one", (PyCFunction)TriMap_register_one, METH_VARARGS, NULL},
1361
+ {"register_unmatched_dst", (PyCFunction)TriMap_register_unmatched_dst, METH_NOARGS, NULL},
1362
+ {"register_many", (PyCFunction)TriMap_register_many, METH_VARARGS, NULL},
1363
+ {"finalize", (PyCFunction)TriMap_finalize, METH_NOARGS, NULL},
1364
+ {"is_many", (PyCFunction)TriMap_is_many, METH_NOARGS, NULL},
1365
+ {"src_no_fill", (PyCFunction)TriMap_src_no_fill, METH_NOARGS, NULL},
1366
+ {"dst_no_fill", (PyCFunction)TriMap_dst_no_fill, METH_NOARGS, NULL},
1367
+ {"map_src_no_fill", (PyCFunction)TriMap_map_src_no_fill, METH_O, NULL},
1368
+ {"map_dst_no_fill", (PyCFunction)TriMap_map_dst_no_fill, METH_O, NULL},
1369
+ {"map_src_fill", (PyCFunction)TriMap_map_src_fill, METH_VARARGS, NULL},
1370
+ {"map_dst_fill", (PyCFunction)TriMap_map_dst_fill, METH_VARARGS, NULL},
1371
+ {"map_merge", (PyCFunction)TriMap_map_merge, METH_VARARGS, NULL},
1372
+ {NULL},
1373
+ };
1374
+
1375
+ PyTypeObject TriMapType = {
1376
+ PyVarObject_HEAD_INIT(NULL, 0)
1377
+ .tp_basicsize = sizeof(TriMapObject), // this does not get size of struct
1378
+ .tp_dealloc = (destructor)TriMap_dealloc,
1379
+ .tp_doc = TriMap_doc,
1380
+ .tp_flags = Py_TPFLAGS_DEFAULT,
1381
+ .tp_methods = TriMap_methods,
1382
+ .tp_name = "arraykit.TriMap",
1383
+ .tp_new = TriMap_new,
1384
+ .tp_init = TriMap_init,
1385
+ .tp_repr = (reprfunc)TriMap_repr,
1386
+ };