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/auto_map.c ADDED
@@ -0,0 +1,2661 @@
1
+ // For background on the hashtable design first implemented in AutoMap, see the following:
2
+ // https://github.com/brandtbucher/automap/blob/b787199d38d6bfa1b55484e5ea1e89b31cc1fa72/automap.c#L12
3
+ # include <math.h>
4
+ # include "Python.h"
5
+ # include "stdbool.h"
6
+
7
+ # define PY_SSIZE_T_CLEAN
8
+
9
+ # define NO_IMPORT_ARRAY
10
+ # define PY_ARRAY_UNIQUE_SYMBOL AK_ARRAY_API
11
+ # define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
12
+
13
+ # include "numpy/arrayobject.h"
14
+ # include "numpy/arrayscalars.h"
15
+ # include "numpy/halffloat.h"
16
+ # include "auto_map.h"
17
+ # include "utilities.h"
18
+
19
+ //------------------------------------------------------------------------------
20
+ // Common
21
+
22
+ PyObject *NonUniqueError;
23
+
24
+ // The main storage "table" is an array of TableElement
25
+ typedef struct TableElement{
26
+ Py_ssize_t keys_pos;
27
+ Py_hash_t hash;
28
+ } TableElement;
29
+
30
+ // Table configuration; experimentation shows that these values work well:
31
+ # define LOAD 0.9
32
+ # define SCAN 16
33
+
34
+
35
+ // Partial, two-argument version of PyUnicode_FromKindAndData for consistent templating with bytes version.
36
+ static inline PyObject*
37
+ PyUnicode_FromUCS4AndData(const void *buffer, Py_ssize_t size) {
38
+ return PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, buffer, size);
39
+ }
40
+
41
+ typedef enum KeysArrayType{
42
+ KAT_LIST = 0, // must be falsy
43
+
44
+ KAT_INT8, // order matters as ranges of size are used in selection
45
+ KAT_INT16,
46
+ KAT_INT32,
47
+ KAT_INT64,
48
+
49
+ KAT_UINT8,
50
+ KAT_UINT16,
51
+ KAT_UINT32,
52
+ KAT_UINT64,
53
+
54
+ KAT_FLOAT16,
55
+ KAT_FLOAT32,
56
+ KAT_FLOAT64,
57
+
58
+ KAT_UNICODE,
59
+ KAT_STRING,
60
+
61
+ KAT_DTY,
62
+ KAT_DTM,
63
+ KAT_DTW,
64
+ KAT_DTD,
65
+
66
+ KAT_DTh,
67
+ KAT_DTm,
68
+ KAT_DTs,
69
+ KAT_DTms,
70
+ KAT_DTus,
71
+ KAT_DTns,
72
+ KAT_DTps,
73
+ KAT_DTfs,
74
+ KAT_DTas,
75
+ } KeysArrayType;
76
+
77
+
78
+ KeysArrayType
79
+ at_to_kat(int array_t, PyArrayObject* a) {
80
+ switch (array_t) {
81
+ case NPY_INT64:
82
+ return KAT_INT64;
83
+ case NPY_INT32:
84
+ return KAT_INT32;
85
+ case NPY_INT16:
86
+ return KAT_INT16;
87
+ case NPY_INT8:
88
+ return KAT_INT8;
89
+
90
+ case NPY_UINT64:
91
+ return KAT_UINT64;
92
+ case NPY_UINT32:
93
+ return KAT_UINT32;
94
+ case NPY_UINT16:
95
+ return KAT_UINT16;
96
+ case NPY_UINT8:
97
+ return KAT_UINT8;
98
+
99
+ case NPY_FLOAT64:
100
+ return KAT_FLOAT64;
101
+ case NPY_FLOAT32:
102
+ return KAT_FLOAT32;
103
+ case NPY_FLOAT16:
104
+ return KAT_FLOAT16;
105
+
106
+ case NPY_UNICODE:
107
+ return KAT_UNICODE;
108
+ case NPY_STRING:
109
+ return KAT_STRING;
110
+
111
+ case NPY_DATETIME: {
112
+ NPY_DATETIMEUNIT dtu = AK_dt_unit_from_array(a);
113
+ switch (dtu) {
114
+ case NPY_FR_Y:
115
+ return KAT_DTY;
116
+ case NPY_FR_M:
117
+ return KAT_DTM;
118
+ case NPY_FR_W:
119
+ return KAT_DTW;
120
+ case NPY_FR_D:
121
+ return KAT_DTD;
122
+ case NPY_FR_h:
123
+ return KAT_DTh;
124
+ case NPY_FR_m:
125
+ return KAT_DTm;
126
+ case NPY_FR_s:
127
+ return KAT_DTs;
128
+ case NPY_FR_ms:
129
+ return KAT_DTms;
130
+ case NPY_FR_us:
131
+ return KAT_DTus;
132
+ case NPY_FR_ns:
133
+ return KAT_DTns;
134
+ case NPY_FR_ps:
135
+ return KAT_DTps;
136
+ case NPY_FR_fs:
137
+ return KAT_DTfs;
138
+ case NPY_FR_as:
139
+ return KAT_DTas;
140
+ case NPY_FR_ERROR:
141
+ case NPY_FR_GENERIC:
142
+ return KAT_LIST; // fall back to list
143
+ }
144
+ }
145
+ default:
146
+ return KAT_LIST;
147
+ }
148
+ }
149
+
150
+ // To determine when we can use direct array lookups, this function return 1 if we match, 0 if we do not match. Given a keys array type and the kind of lookup key, return true only for the largest KAT types.
151
+ int
152
+ kat_is_kind(KeysArrayType kat, char kind) {
153
+ switch (kat) {
154
+ case KAT_INT64:
155
+ // case KAT_INT32:
156
+ // case KAT_INT16:
157
+ // case KAT_INT8:
158
+ return kind == 'i';
159
+
160
+ case KAT_UINT64:
161
+ // case KAT_UINT32:
162
+ // case KAT_UINT16:
163
+ // case KAT_UINT8:
164
+ return kind == 'u';
165
+
166
+ case KAT_FLOAT64:
167
+ // case KAT_FLOAT32:
168
+ // case KAT_FLOAT16:
169
+ return kind == 'f';
170
+
171
+ case KAT_UNICODE:
172
+ return kind == 'U';
173
+ case KAT_STRING:
174
+ return kind == 'S';
175
+
176
+ case KAT_DTY:
177
+ case KAT_DTM:
178
+ case KAT_DTW:
179
+ case KAT_DTD:
180
+ case KAT_DTh:
181
+ case KAT_DTm:
182
+ case KAT_DTs:
183
+ case KAT_DTms:
184
+ case KAT_DTus:
185
+ case KAT_DTns:
186
+ case KAT_DTps:
187
+ case KAT_DTfs:
188
+ case KAT_DTas:
189
+ return kind == 'M';
190
+
191
+ default:
192
+ return 0;
193
+ }
194
+ }
195
+
196
+ // Given a KAT, determine if it matches a NumPy dt64 unit.
197
+ bool
198
+ kat_is_datetime_unit(KeysArrayType kat, NPY_DATETIMEUNIT unit) {
199
+ switch (kat) {
200
+ case KAT_DTY:
201
+ if (unit == NPY_FR_Y ) {return true;}
202
+ break;
203
+ case KAT_DTM:
204
+ if (unit == NPY_FR_M ) {return true;}
205
+ break;
206
+ case KAT_DTW:
207
+ if (unit == NPY_FR_W ) {return true;}
208
+ break;
209
+ case KAT_DTD:
210
+ if (unit == NPY_FR_D ) {return true;}
211
+ break;
212
+ case KAT_DTh:
213
+ if (unit == NPY_FR_h ) {return true;}
214
+ break;
215
+ case KAT_DTm:
216
+ if (unit == NPY_FR_m ) {return true;}
217
+ break;
218
+ case KAT_DTs:
219
+ if (unit == NPY_FR_s ) {return true;}
220
+ break;
221
+ case KAT_DTms:
222
+ if (unit == NPY_FR_ms) {return true;}
223
+ break;
224
+ case KAT_DTus:
225
+ if (unit == NPY_FR_us) {return true;}
226
+ break;
227
+ case KAT_DTns:
228
+ if (unit == NPY_FR_ns) {return true;}
229
+ break;
230
+ case KAT_DTps:
231
+ if (unit == NPY_FR_ps) {return true;}
232
+ break;
233
+ case KAT_DTfs:
234
+ if (unit == NPY_FR_fs) {return true;}
235
+ break;
236
+ case KAT_DTas:
237
+ if (unit == NPY_FR_as) {return true;}
238
+ break;
239
+ default: // non dt64 KATs
240
+ return false;
241
+ }
242
+ return false;
243
+ }
244
+
245
+ typedef struct FAMObject{
246
+ PyObject_HEAD
247
+ Py_ssize_t table_size;
248
+ TableElement *table; // an array of TableElement structs
249
+ PyObject *keys;
250
+ KeysArrayType keys_array_type;
251
+ Py_ssize_t keys_size;
252
+ Py_UCS4* key_buffer;
253
+ } FAMObject;
254
+
255
+ typedef enum ViewKind{
256
+ ITEMS,
257
+ KEYS,
258
+ VALUES,
259
+ } ViewKind;
260
+
261
+ // Return the end pointer, or the pointer to the location after the last valid character. The end pointer minus the start pointer is the number of characters. For an empty string, all characters are NULL, and the start pointer and end pointer should be equal. NOTE: would like to use strchr(str, '\0') instead of this routine, but some buffers might not have a null terminator and stread by full to the the dt_size.
262
+ static inline Py_UCS4*
263
+ ucs4_get_end_p(Py_UCS4* p_start, Py_ssize_t dt_size) {
264
+ for (Py_UCS4* p = p_start + dt_size - 1; p >= p_start; p--) {
265
+ if (*p != '\0') {
266
+ return p + 1; // 1 after first non-null
267
+ }
268
+ }
269
+ return p_start;
270
+ }
271
+
272
+ static inline char*
273
+ char_get_end_p(char* p_start, Py_ssize_t dt_size) {
274
+ for (char* p = p_start + dt_size - 1; p >= p_start; p--) {
275
+ if (*p != '\0') {
276
+ return p + 1; // 1 after first non-null
277
+ }
278
+ }
279
+ return p_start;
280
+ }
281
+
282
+ // This masks the input with INT64_MAX, which removes the MSB; we then cast to an int64; the range is now between 0 and INT64_MAX. We then use the MSB of the original value; if set, we negate the number, producing negative values for the upper half of the uint64 range. Note that we only need to check for hash -1 in this branch.
283
+ static inline Py_hash_t
284
+ uint_to_hash(npy_uint64 v) {
285
+ Py_hash_t hash = (Py_hash_t)(v & INT64_MAX);
286
+ if (v >> 63) {
287
+ hash = -hash;
288
+ }
289
+ if (hash == -1) { // might happen due to overflow on 32 bit systems
290
+ return -2;
291
+ }
292
+ return hash;
293
+ }
294
+
295
+ static inline Py_hash_t
296
+ int_to_hash(npy_int64 v) {
297
+ Py_hash_t hash = (Py_hash_t)v;
298
+ if (hash == -1) {
299
+ return -2;
300
+ }
301
+ return hash;
302
+ }
303
+
304
+ // This is a adapted from https://github.com/python/cpython/blob/ba65a065cf07a7a9f53be61057a090f7311a5ad7/Python/pyhash.c#L92
305
+ #define HASH_MODULUS (((size_t)1 << 61) - 1)
306
+ #define HASH_BITS 61
307
+ static inline Py_hash_t
308
+ double_to_hash(double v)
309
+ {
310
+ int e, sign;
311
+ double m;
312
+ Py_uhash_t x, y;
313
+
314
+ if (isinf(v)) {
315
+ return v > 0 ? 314159 : -314159;
316
+ }
317
+ if (isnan(v)) {
318
+ return 0;
319
+ }
320
+ m = frexp(v, &e);
321
+ sign = 1;
322
+ if (m < 0) {
323
+ sign = -1;
324
+ m = -m;
325
+ }
326
+ x = 0;
327
+ while (m) {
328
+ x = ((x << 28) & HASH_MODULUS) | x >> (HASH_BITS - 28);
329
+ m *= 268435456.0; /* 2**28 */
330
+ e -= 28;
331
+ y = (Py_uhash_t)m; /* pull out integer part */
332
+ m -= y;
333
+ x += y;
334
+ if (x >= HASH_MODULUS)
335
+ x -= HASH_MODULUS;
336
+ }
337
+ e = e >= 0 ? e % HASH_BITS : HASH_BITS-1-((-1-e) % HASH_BITS);
338
+ x = ((x << e) & HASH_MODULUS) | x >> (HASH_BITS - e);
339
+ x = x * sign;
340
+ if (x == (Py_uhash_t)-1)
341
+ x = (Py_uhash_t)-2;
342
+ return (Py_hash_t)x;
343
+ }
344
+
345
+ // The `str` arg is a pointer to a C-array of Py_UCS4; we will only read `len` characters from this. This is a "djb2" hash algorithm.
346
+ static inline Py_hash_t
347
+ unicode_to_hash(Py_UCS4 *str, Py_ssize_t len) {
348
+ Py_UCS4* p = str;
349
+ Py_UCS4* p_end = str + len;
350
+ Py_hash_t hash = 5381;
351
+ while (p < p_end) {
352
+ hash = ((hash << 5) + hash) + *p++;
353
+ }
354
+ if (hash == -1) {
355
+ return -2;
356
+ }
357
+ return hash;
358
+ }
359
+
360
+ static inline Py_hash_t
361
+ string_to_hash(char *str, Py_ssize_t len) {
362
+ char* p = str;
363
+ char* p_end = str + len;
364
+ Py_hash_t hash = 5381;
365
+ while (p < p_end) {
366
+ hash = ((hash << 5) + hash) + *p++;
367
+ }
368
+ if (hash == -1) {
369
+ return -2;
370
+ }
371
+ return hash;
372
+ }
373
+
374
+ //------------------------------------------------------------------------------
375
+ // FrozenAutoMapIterator functions
376
+
377
+ typedef struct FAMIObject {
378
+ PyObject_HEAD
379
+ FAMObject *fam;
380
+ PyArrayObject* keys_array;
381
+ ViewKind kind;
382
+ bool reversed;
383
+ Py_ssize_t index; // current index state, mutated in-place
384
+ } FAMIObject;
385
+
386
+ static void
387
+ fami_dealloc(FAMIObject *self)
388
+ {
389
+ Py_DECREF(self->fam);
390
+ PyObject_Del((PyObject *)self);
391
+ }
392
+
393
+ static FAMIObject *
394
+ fami_iter(FAMIObject *self)
395
+ {
396
+ Py_INCREF(self);
397
+ return self;
398
+ }
399
+
400
+ // For a FAMI, Return appropriate PyObject for items, keys, and values. For consistency with NumPy array iteration, arrays use PyArray_ToScalar instead of PyArray_GETITEM.
401
+ static PyObject *
402
+ fami_iternext(FAMIObject *self)
403
+ {
404
+ Py_ssize_t index;
405
+ if (self->reversed) {
406
+ index = self->fam->keys_size - ++self->index;
407
+ if (index < 0) {
408
+ return NULL;
409
+ }
410
+ }
411
+ else {
412
+ index = self->index++;
413
+ }
414
+ if (self->fam->keys_size <= index) {
415
+ return NULL;
416
+ }
417
+ switch (self->kind) {
418
+ case ITEMS: {
419
+ if (self->fam->keys_array_type) {
420
+ return Py_BuildValue(
421
+ "NN",
422
+ PyArray_ToScalar(PyArray_GETPTR1(self->keys_array, index), self->keys_array),
423
+ PyLong_FromSsize_t(index)
424
+ );
425
+ }
426
+ else {
427
+ PyObject* t = PyTuple_New(2);
428
+ if (!t) { return NULL; }
429
+ #if PY_VERSION_HEX >= 0x030D0000 // Python 3.13+
430
+ PyObject* k = PyList_GetItemRef(self->fam->keys, index);
431
+ #else
432
+ PyObject* k = PyList_GET_ITEM(self->fam->keys, index);
433
+ Py_XINCREF(k);
434
+ #endif
435
+ if (!k) { return NULL; }
436
+ PyTuple_SET_ITEM(t, 0, k);
437
+ PyTuple_SET_ITEM(t, 1, PyLong_FromSsize_t(index));
438
+ return t;
439
+ }
440
+ }
441
+ case KEYS: {
442
+ if (self->fam->keys_array_type) {
443
+ return PyArray_ToScalar(PyArray_GETPTR1(self->keys_array, index), self->keys_array);
444
+ }
445
+ else {
446
+ #if PY_VERSION_HEX >= 0x030D0000 // Python 3.13+
447
+ PyObject* yield = PyList_GetItemRef(self->fam->keys, index);
448
+ #else
449
+ PyObject* yield = PyList_GET_ITEM(self->fam->keys, index);
450
+ Py_XINCREF(yield);
451
+ #endif
452
+ if (!yield) { return NULL; }
453
+ return yield;
454
+ }
455
+ }
456
+ case VALUES: {
457
+ return PyLong_FromSsize_t(index);
458
+ }
459
+ }
460
+ Py_UNREACHABLE();
461
+ }
462
+
463
+ static PyObject *
464
+ fami_length_hint(FAMIObject *self)
465
+ {
466
+ Py_ssize_t len = Py_MAX(0, self->fam->keys_size - self->index);
467
+ return PyLong_FromSsize_t(len);
468
+ }
469
+
470
+ static PyObject *fami_new(FAMObject *, ViewKind, bool);
471
+
472
+ static PyObject *
473
+ fami_reversed(FAMIObject *self)
474
+ {
475
+ return fami_new(self->fam, self->kind, !self->reversed);
476
+ }
477
+
478
+ static PyMethodDef fami_methods[] = {
479
+ {"__length_hint__", (PyCFunction)fami_length_hint, METH_NOARGS, NULL},
480
+ {"__reversed__", (PyCFunction)fami_reversed, METH_NOARGS, NULL},
481
+ {NULL},
482
+ };
483
+
484
+ PyTypeObject FAMIType = {
485
+ PyVarObject_HEAD_INIT(NULL, 0)
486
+ .tp_basicsize = sizeof(FAMIObject),
487
+ .tp_dealloc = (destructor) fami_dealloc,
488
+ .tp_iter = (getiterfunc) fami_iter,
489
+ .tp_iternext = (iternextfunc) fami_iternext,
490
+ .tp_methods = fami_methods,
491
+ .tp_name = "arraykit.FrozenAutoMapIterator",
492
+ };
493
+
494
+ static PyObject *
495
+ fami_new(FAMObject *fam, ViewKind kind, bool reversed)
496
+ {
497
+ FAMIObject *fami = PyObject_New(FAMIObject, &FAMIType);
498
+ if (!fami) {
499
+ return NULL;
500
+ }
501
+ Py_INCREF(fam);
502
+ fami->fam = fam;
503
+ if (fam->keys_array_type) {
504
+ fami->keys_array = (PyArrayObject *)fam->keys;
505
+ }
506
+ else {
507
+ fami->keys_array = NULL;
508
+ }
509
+ fami->kind = kind;
510
+ fami->reversed = reversed;
511
+ fami->index = 0;
512
+ return (PyObject *)fami;
513
+ }
514
+
515
+ //------------------------------------------------------------------------------
516
+ // FrozenAutoMapView functions
517
+
518
+ // A FAMVObject contains a reference to the FAM from which it was derived
519
+ typedef struct FAMVObject{
520
+ PyObject_HEAD
521
+ FAMObject *fam;
522
+ ViewKind kind;
523
+ } FAMVObject;
524
+
525
+ # define FAMV_SET_OP(name, op) \
526
+ static PyObject * \
527
+ name(PyObject *left, PyObject *right) \
528
+ { \
529
+ left = PySet_New(left); \
530
+ if (!left) { \
531
+ return NULL; \
532
+ } \
533
+ right = PySet_New(right); \
534
+ if (!right) { \
535
+ Py_DECREF(left); \
536
+ return NULL; \
537
+ } \
538
+ PyObject *result = PyNumber_InPlace##op(left, right); \
539
+ Py_DECREF(left); \
540
+ Py_DECREF(right); \
541
+ return result; \
542
+ }
543
+
544
+ FAMV_SET_OP(famv_and, And)
545
+ FAMV_SET_OP(famv_or, Or)
546
+ FAMV_SET_OP(famv_subtract, Subtract)
547
+ FAMV_SET_OP(famv_xor, Xor)
548
+
549
+ # undef FAMV_SET_OP
550
+
551
+ static PyNumberMethods famv_as_number = {
552
+ .nb_and = (binaryfunc) famv_and,
553
+ .nb_or = (binaryfunc) famv_or,
554
+ .nb_subtract = (binaryfunc) famv_subtract,
555
+ .nb_xor = (binaryfunc) famv_xor,
556
+ };
557
+
558
+ static int fam_contains(FAMObject *, PyObject *);
559
+ static PyObject *famv_fami_new(FAMVObject *);
560
+
561
+ static int
562
+ famv_contains(FAMVObject *self, PyObject *other)
563
+ {
564
+ if (self->kind == KEYS) {
565
+ return fam_contains(self->fam, other);
566
+ }
567
+ PyObject *iterator = famv_fami_new(self);
568
+ if (!iterator) {
569
+ return -1;
570
+ }
571
+ int result = PySequence_Contains(iterator, other);
572
+ Py_DECREF(iterator);
573
+ return result;
574
+ }
575
+
576
+ static PySequenceMethods famv_as_sequence = {
577
+ .sq_contains = (objobjproc) famv_contains,
578
+ };
579
+
580
+ static void
581
+ famv_dealloc(FAMVObject *self)
582
+ {
583
+ Py_DECREF(self->fam);
584
+ PyObject_Del((PyObject *)self);
585
+ }
586
+
587
+ static PyObject *
588
+ famv_fami_new(FAMVObject *self)
589
+ {
590
+ return fami_new(self->fam, self->kind, false);
591
+ }
592
+
593
+ static PyObject *
594
+ famv_length_hint(FAMVObject *self)
595
+ {
596
+ return PyLong_FromSsize_t(self->fam->keys_size);
597
+ }
598
+
599
+ static PyObject *
600
+ famv_reversed(FAMVObject *self)
601
+ {
602
+ return fami_new(self->fam, self->kind, true);
603
+ }
604
+
605
+ static PyObject *
606
+ famv_isdisjoint(FAMVObject *self, PyObject *other)
607
+ {
608
+ PyObject *intersection = famv_and((PyObject *)self, other);
609
+ if (!intersection) {
610
+ return NULL;
611
+ }
612
+ Py_ssize_t result = PySet_GET_SIZE(intersection);
613
+ Py_DECREF(intersection);
614
+ return PyBool_FromLong(result);
615
+ }
616
+
617
+ static PyObject *
618
+ famv_richcompare(FAMVObject *self, PyObject *other, int op)
619
+ {
620
+ PyObject *left = PySet_New((PyObject *)self);
621
+ if (!left) {
622
+ return NULL;
623
+ }
624
+ PyObject *right = PySet_New(other);
625
+ if (!right) {
626
+ Py_DECREF(left);
627
+ return NULL;
628
+ }
629
+ PyObject *result = PyObject_RichCompare(left, right, op);
630
+ Py_DECREF(left);
631
+ Py_DECREF(right);
632
+ return result;
633
+ }
634
+
635
+ static PyMethodDef famv_methods[] = {
636
+ {"__length_hint__", (PyCFunction) famv_length_hint, METH_NOARGS, NULL},
637
+ {"__reversed__", (PyCFunction) famv_reversed, METH_NOARGS, NULL},
638
+ {"isdisjoint", (PyCFunction) famv_isdisjoint, METH_O, NULL},
639
+ {NULL},
640
+ };
641
+
642
+ PyTypeObject FAMVType = {
643
+ PyVarObject_HEAD_INIT(NULL, 0)
644
+ .tp_as_number = &famv_as_number,
645
+ .tp_as_sequence = &famv_as_sequence,
646
+ .tp_basicsize = sizeof(FAMVObject),
647
+ .tp_dealloc = (destructor) famv_dealloc,
648
+ .tp_iter = (getiterfunc) famv_fami_new,
649
+ .tp_methods = famv_methods,
650
+ .tp_name = "arraykit.FrozenAutoMapView",
651
+ .tp_richcompare = (richcmpfunc) famv_richcompare,
652
+ };
653
+
654
+ static PyObject *
655
+ famv_new(FAMObject *fam, ViewKind kind)
656
+ {
657
+ FAMVObject *famv = (FAMVObject *)PyObject_New(FAMVObject, &FAMVType);
658
+ if (!famv) {
659
+ return NULL;
660
+ }
661
+ famv->kind = kind;
662
+ famv->fam = fam;
663
+ Py_INCREF(fam);
664
+ return (PyObject *)famv;
665
+ }
666
+
667
+ //------------------------------------------------------------------------------
668
+ // FrozenAutoMap functions
669
+
670
+ // Given a key and a computed hash, return the table_pos if that hash and key are found, or if not, the first table position that has not been assigned. Return -1 on error.
671
+ static Py_ssize_t
672
+ lookup_hash_obj(FAMObject *self, PyObject *key, Py_hash_t hash)
673
+ {
674
+ TableElement *table = self->table;
675
+ Py_ssize_t mask = self->table_size - 1;
676
+ Py_hash_t mixin = Py_ABS(hash);
677
+ Py_ssize_t table_pos = hash & mask;
678
+
679
+ PyObject *guess = NULL;
680
+ PyObject *keys = self->keys;
681
+ int result = -1;
682
+ Py_hash_t h = 0;
683
+
684
+ while (1) {
685
+ for (Py_ssize_t i = 0; i < SCAN; i++) {
686
+ h = table[table_pos].hash;
687
+ if (h == -1) { // Miss. Found a position that can be used for insertion.
688
+ return table_pos;
689
+ }
690
+ if (h != hash) { // Collision.
691
+ table_pos++;
692
+ continue;
693
+ }
694
+ guess = PyList_GET_ITEM(keys, table[table_pos].keys_pos);
695
+ if (guess == key) { // Hit. Object ID comparison
696
+ return table_pos;
697
+ }
698
+
699
+ // if key is a dt64, only do PyObject_RichCompareBool if units match
700
+ if (PyArray_IsScalar(key, Datetime) && PyArray_IsScalar(guess, Datetime)) {
701
+ if (AK_dt_unit_from_scalar((PyDatetimeScalarObject *)key)
702
+ != AK_dt_unit_from_scalar((PyDatetimeScalarObject *)guess)) {
703
+ table_pos++;
704
+ continue;
705
+ }
706
+ }
707
+
708
+ result = PyObject_RichCompareBool(guess, key, Py_EQ);
709
+ if (result < 0) { // Error.
710
+ return -1;
711
+ }
712
+ if (result) { // Hit.
713
+ return table_pos;
714
+ }
715
+ table_pos++;
716
+ }
717
+ table_pos = (5 * (table_pos - SCAN) + (mixin >>= 1) + 1) & mask;
718
+ }
719
+ }
720
+
721
+
722
+ // Used for both integer and datetime types; for this reason kat is passed in separately.
723
+ static Py_ssize_t
724
+ lookup_hash_int(FAMObject *self, npy_int64 key, Py_hash_t hash, KeysArrayType kat)
725
+ {
726
+ TableElement *table = self->table;
727
+ Py_ssize_t mask = self->table_size - 1;
728
+ Py_hash_t mixin = Py_ABS(hash);
729
+ Py_ssize_t table_pos = hash & mask; // taking the modulo
730
+
731
+ PyArrayObject *a = (PyArrayObject *)self->keys;
732
+ npy_int64 k = 0;
733
+ Py_hash_t h = 0;
734
+
735
+ while (1) {
736
+ for (Py_ssize_t i = 0; i < SCAN; i++) {
737
+ h = table[table_pos].hash;
738
+ if (h == -1) { // Miss. Position that can be used for insertion.
739
+ return table_pos;
740
+ }
741
+ if (h != hash) {
742
+ table_pos++;
743
+ continue;
744
+ }
745
+ switch (kat) {
746
+ case KAT_INT64:
747
+ k = *(npy_int64*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
748
+ break;
749
+ case KAT_INT32:
750
+ k = *(npy_int32*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
751
+ break;
752
+ case KAT_INT16:
753
+ k = *(npy_int16*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
754
+ break;
755
+ case KAT_INT8:
756
+ k = *(npy_int8*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
757
+ break;
758
+ default:
759
+ return -1;
760
+ }
761
+ if (key == k) {
762
+ return table_pos;
763
+ }
764
+ table_pos++;
765
+ }
766
+ table_pos = (5 * (table_pos - SCAN) + (mixin >>= 1) + 1) & mask;
767
+ }
768
+ }
769
+
770
+
771
+ // NOTE: kat is passed in separately to match the interface of lookup_hash_int.
772
+ static Py_ssize_t
773
+ lookup_hash_uint(FAMObject *self, npy_uint64 key, Py_hash_t hash, KeysArrayType kat)
774
+ {
775
+ TableElement *table = self->table;
776
+ Py_ssize_t mask = self->table_size - 1;
777
+ Py_hash_t mixin = Py_ABS(hash);
778
+ Py_ssize_t table_pos = hash & mask;
779
+
780
+ PyArrayObject *a = (PyArrayObject *)self->keys;
781
+ npy_uint64 k = 0;
782
+ Py_hash_t h = 0;
783
+
784
+ while (1) {
785
+ for (Py_ssize_t i = 0; i < SCAN; i++) {
786
+ h = table[table_pos].hash;
787
+ if (h == -1) {
788
+ return table_pos;
789
+ }
790
+ if (h != hash) {
791
+ table_pos++;
792
+ continue;
793
+ }
794
+ switch (kat) {
795
+ case KAT_UINT64:
796
+ k = *(npy_uint64*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
797
+ break;
798
+ case KAT_UINT32:
799
+ k = *(npy_uint32*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
800
+ break;
801
+ case KAT_UINT16:
802
+ k = *(npy_uint16*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
803
+ break;
804
+ case KAT_UINT8:
805
+ k = *(npy_uint8*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
806
+ break;
807
+ default:
808
+ return -1;
809
+ }
810
+ if (key == k) {
811
+ return table_pos;
812
+ }
813
+ table_pos++;
814
+ }
815
+ table_pos = (5 * (table_pos - SCAN) + (mixin >>= 1) + 1) & mask;
816
+ }
817
+ }
818
+
819
+
820
+ // NOTE: kat is passed in separately to match the interface of lookup_hash_int
821
+ static Py_ssize_t
822
+ lookup_hash_double(FAMObject *self, npy_double key, Py_hash_t hash, KeysArrayType kat)
823
+ {
824
+ TableElement *table = self->table;
825
+ Py_ssize_t mask = self->table_size - 1;
826
+ Py_hash_t mixin = Py_ABS(hash);
827
+ Py_ssize_t table_pos = hash & mask;
828
+
829
+ PyArrayObject *a = (PyArrayObject *)self->keys;
830
+ npy_double k = 0;
831
+ Py_hash_t h = 0;
832
+
833
+ while (1) {
834
+ for (Py_ssize_t i = 0; i < SCAN; i++) {
835
+ h = table[table_pos].hash;
836
+ if (h == -1) {
837
+ return table_pos;
838
+ }
839
+ if (h != hash) {
840
+ table_pos++;
841
+ continue;
842
+ }
843
+ switch (kat) {
844
+ case KAT_FLOAT64:
845
+ k = *(npy_double*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
846
+ break;
847
+ case KAT_FLOAT32:
848
+ k = *(npy_float*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
849
+ break;
850
+ case KAT_FLOAT16:
851
+ k = npy_half_to_double(*(npy_half*)PyArray_GETPTR1(a, table[table_pos].keys_pos));
852
+ break;
853
+ default:
854
+ return -1;
855
+ }
856
+ if (key == k) {
857
+ return table_pos;
858
+ }
859
+ table_pos++;
860
+ }
861
+ table_pos = (5 * (table_pos - SCAN) + (mixin >>= 1) + 1) & mask;
862
+ }
863
+ }
864
+
865
+
866
+ // Compare a passed Py_UCS4 array to stored keys. This does not use any dynamic memory. Returns -1 on error.
867
+ static Py_ssize_t
868
+ lookup_hash_unicode(
869
+ FAMObject *self,
870
+ Py_UCS4* key,
871
+ Py_ssize_t key_size,
872
+ Py_hash_t hash)
873
+ {
874
+ TableElement *table = self->table;
875
+ Py_ssize_t mask = self->table_size - 1;
876
+ Py_hash_t mixin = Py_ABS(hash);
877
+ Py_ssize_t table_pos = hash & mask;
878
+
879
+ PyArrayObject *a = (PyArrayObject *)self->keys;
880
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
881
+ Py_ssize_t cmp_bytes = Py_MIN(key_size, dt_size) * UCS4_SIZE;
882
+
883
+ Py_hash_t h = 0;
884
+ Py_UCS4* p_start = NULL;
885
+
886
+ while (1) {
887
+ for (Py_ssize_t i = 0; i < SCAN; i++) {
888
+ h = table[table_pos].hash;
889
+ if (h == -1) {
890
+ return table_pos;
891
+ }
892
+ if (h != hash) {
893
+ table_pos++;
894
+ continue;
895
+ }
896
+ p_start = (Py_UCS4*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
897
+ // memcmp returns 0 on match
898
+ if (!memcmp(p_start, key, cmp_bytes)) {
899
+ return table_pos;
900
+ }
901
+ table_pos++;
902
+ }
903
+ table_pos = (5 * (table_pos - SCAN) + (mixin >>= 1) + 1) & mask;
904
+ }
905
+ }
906
+
907
+
908
+ // Compare a passed char array to stored keys. This does not use any dynamic memory. Returns -1 on error.
909
+ static Py_ssize_t
910
+ lookup_hash_string(
911
+ FAMObject *self,
912
+ char* key,
913
+ Py_ssize_t key_size,
914
+ Py_hash_t hash)
915
+ {
916
+ TableElement *table = self->table;
917
+ Py_ssize_t mask = self->table_size - 1;
918
+ Py_hash_t mixin = Py_ABS(hash);
919
+ Py_ssize_t table_pos = hash & mask;
920
+
921
+ PyArrayObject *a = (PyArrayObject *)self->keys;
922
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a);
923
+ Py_ssize_t cmp_bytes = Py_MIN(key_size, dt_size);
924
+
925
+ Py_hash_t h = 0;
926
+ char* p_start = NULL;
927
+
928
+ while (1) {
929
+ for (Py_ssize_t i = 0; i < SCAN; i++) {
930
+ h = table[table_pos].hash;
931
+ if (h == -1) {
932
+ return table_pos;
933
+ }
934
+ if (h != hash) {
935
+ table_pos++;
936
+ continue;
937
+ }
938
+ p_start = (char*)PyArray_GETPTR1(a, table[table_pos].keys_pos);
939
+ // memcmp returns 0 on match
940
+ if (!memcmp(p_start, key, cmp_bytes)) {
941
+ return table_pos;
942
+ }
943
+ table_pos++;
944
+ }
945
+ table_pos = (5 * (table_pos - SCAN) + (mixin >>= 1) + 1) & mask;
946
+ }
947
+ }
948
+
949
+
950
+ static Py_ssize_t
951
+ lookup_int(FAMObject *self, PyObject* key) {
952
+ npy_int64 v = 0;
953
+ // NOTE: we handle PyArray Scalar Byte, Short, UByte, UShort with PyNumber_Check, below, saving four branches here
954
+ if (PyArray_IsScalar(key, LongLong)) {
955
+ v = (npy_int64)PyArrayScalar_VAL(key, LongLong);
956
+ }
957
+ else if (PyArray_IsScalar(key, Long)) {
958
+ v = (npy_int64)PyArrayScalar_VAL(key, Long);
959
+ }
960
+ else if (PyLong_Check(key)) {
961
+ v = PyLong_AsLongLong(key);
962
+ if (v == -1 && PyErr_Occurred()) {
963
+ PyErr_Clear();
964
+ return -1;
965
+ }
966
+ }
967
+ else if (PyArray_IsScalar(key, Double)) {
968
+ double dv = PyArrayScalar_VAL(key, Double);
969
+ if (floor(dv) != dv) {
970
+ return -1;
971
+ }
972
+ v = (npy_int64)dv;
973
+ }
974
+ else if (PyFloat_Check(key)) {
975
+ double dv = PyFloat_AsDouble(key);
976
+ if (dv == -1.0 && PyErr_Occurred()) {
977
+ PyErr_Clear();
978
+ return -1;
979
+ }
980
+ v = (npy_int64)dv; // truncate to integer
981
+ if (v != dv) {
982
+ return -1;
983
+ }
984
+ }
985
+ else if (PyArray_IsScalar(key, ULongLong)) {
986
+ v = (npy_int64)PyArrayScalar_VAL(key, ULongLong);
987
+ }
988
+ else if (PyArray_IsScalar(key, ULong)) {
989
+ v = (npy_int64)PyArrayScalar_VAL(key, ULong);
990
+ }
991
+ else if (PyArray_IsScalar(key, Int)) {
992
+ v = (npy_int64)PyArrayScalar_VAL(key, Int);
993
+ }
994
+ else if (PyArray_IsScalar(key, UInt)) {
995
+ v = (npy_int64)PyArrayScalar_VAL(key, UInt);
996
+ }
997
+ else if (PyArray_IsScalar(key, Float)) {
998
+ double dv = (double)PyArrayScalar_VAL(key, Float);
999
+ if (floor(dv) != dv) {
1000
+ return -1;
1001
+ }
1002
+ v = (npy_int64)dv;
1003
+ }
1004
+ else if (PyArray_IsScalar(key, Half)) {
1005
+ double dv = npy_half_to_double(PyArrayScalar_VAL(key, Half));
1006
+ if (floor(dv) != dv) {
1007
+ return -1;
1008
+ }
1009
+ v = (npy_int64)dv;
1010
+ }
1011
+ else if (PyBool_Check(key)) {
1012
+ v = PyObject_IsTrue(key);
1013
+ }
1014
+ else if (PyNumber_Check(key)) {
1015
+ // NOTE: this returns a Py_ssize_t, which might be 32 bit. This can be used for PyArray_Scalars <= ssize_t.
1016
+ v = (npy_int64)PyNumber_AsSsize_t(key, PyExc_OverflowError);
1017
+ if (v == -1 && PyErr_Occurred()) {
1018
+ return -1;
1019
+ }
1020
+ }
1021
+ else {
1022
+ return -1;
1023
+ }
1024
+ Py_hash_t hash = int_to_hash(v);
1025
+ return lookup_hash_int(self, v, hash, self->keys_array_type);
1026
+ }
1027
+
1028
+
1029
+ // In current usage as an AM, np.datetime64 will match to any numpy Scalar that is at or greater than the resolution of the values stored here. No matches are made to other numerics or Python datetime objects. For AM to be consistent with FAM, we will do the same for now.
1030
+ static Py_ssize_t
1031
+ lookup_datetime(FAMObject *self, PyObject* key) {
1032
+ npy_int64 v = 0; // int64
1033
+ if (PyArray_IsScalar(key, Datetime)) {
1034
+ v = (npy_int64)PyArrayScalar_VAL(key, Datetime);
1035
+ // if we observe a NAT, we skip unit checks
1036
+
1037
+ if (v != NPY_DATETIME_NAT) {
1038
+ NPY_DATETIMEUNIT key_unit = AK_dt_unit_from_scalar(
1039
+ (PyDatetimeScalarObject *)key);
1040
+ if (!kat_is_datetime_unit(self->keys_array_type, key_unit)) {
1041
+ return -1;
1042
+ }
1043
+ }
1044
+ }
1045
+ else {
1046
+ return -1;
1047
+ }
1048
+ Py_hash_t hash = int_to_hash(v);
1049
+ return lookup_hash_int(self, v, hash, KAT_INT64);
1050
+ }
1051
+
1052
+
1053
+ static Py_ssize_t
1054
+ lookup_uint(FAMObject *self, PyObject* key) {
1055
+ npy_uint64 v = 0;
1056
+
1057
+ // NOTE: we handle PyArray Scalar Byte, Short, UByte, UShort with PyNumber_Check, below, saving four branches here
1058
+ if (PyArray_IsScalar(key, ULongLong)) {
1059
+ v = (npy_uint64)PyArrayScalar_VAL(key, ULongLong);
1060
+ }
1061
+ else if (PyArray_IsScalar(key, ULong)) {
1062
+ v = (npy_uint64)PyArrayScalar_VAL(key, ULong);
1063
+ }
1064
+ else if (PyArray_IsScalar(key, LongLong)) {
1065
+ npy_int64 si = (npy_int64)PyArrayScalar_VAL(key, LongLong);
1066
+ if (si < 0) {
1067
+ return -1;
1068
+ }
1069
+ v = (npy_uint64)si;
1070
+ }
1071
+ else if (PyArray_IsScalar(key, Long)) {
1072
+ npy_int64 si = (npy_int64)PyArrayScalar_VAL(key, Long);
1073
+ if (si < 0) {
1074
+ return -1;
1075
+ }
1076
+ v = (npy_uint64)si;
1077
+ }
1078
+ else if (PyLong_Check(key)) {
1079
+ v = PyLong_AsUnsignedLongLong(key);
1080
+ if (v == (unsigned long long)-1 && PyErr_Occurred()) {
1081
+ PyErr_Clear();
1082
+ return -1;
1083
+ }
1084
+ }
1085
+ else if (PyArray_IsScalar(key, Double)) {
1086
+ double dv = PyArrayScalar_VAL(key, Double);
1087
+ if (dv < 0 || floor(dv) != dv) {
1088
+ return -1;
1089
+ }
1090
+ v = (npy_uint64)dv;
1091
+ }
1092
+ else if (PyFloat_Check(key)) {
1093
+ double dv = PyFloat_AsDouble(key);
1094
+ if (dv == -1.0 && PyErr_Occurred()) {
1095
+ PyErr_Clear();
1096
+ return -1;
1097
+ }
1098
+ if (dv < 0) {
1099
+ return -1;
1100
+ }
1101
+ v = (npy_uint64)dv; // truncate to integer
1102
+ if (v != dv) {
1103
+ return -1;
1104
+ }
1105
+ }
1106
+ else if (PyArray_IsScalar(key, Int)) {
1107
+ npy_int64 si = (npy_int64)PyArrayScalar_VAL(key, Int);
1108
+ if (si < 0) {
1109
+ return -1;
1110
+ }
1111
+ v = (npy_uint64)si;
1112
+ }
1113
+ else if (PyArray_IsScalar(key, UInt)) {
1114
+ v = (npy_uint64)PyArrayScalar_VAL(key, UInt);
1115
+ }
1116
+ else if (PyArray_IsScalar(key, Float)) {
1117
+ double dv = (double)PyArrayScalar_VAL(key, Float);
1118
+ if (dv < 0 || floor(dv) != dv) {
1119
+ return -1;
1120
+ }
1121
+ v = (npy_uint64)dv;
1122
+ }
1123
+ else if (PyArray_IsScalar(key, Half)) {
1124
+ double dv = npy_half_to_double(PyArrayScalar_VAL(key, Half));
1125
+ if (dv < 0 || floor(dv) != dv) {
1126
+ return -1;
1127
+ }
1128
+ v = (npy_uint64)dv;
1129
+ }
1130
+ else if (PyBool_Check(key)) {
1131
+ v = PyObject_IsTrue(key);
1132
+ }
1133
+ else if (PyNumber_Check(key)) {
1134
+ // NOTE: this returns a Py_ssize_t, which might be 32 bit. This can be used for PyArray_Scalars <= ssize_t.
1135
+ npy_int64 si = PyNumber_AsSsize_t(key, PyExc_OverflowError);
1136
+ if (si == -1 && PyErr_Occurred()) {
1137
+ PyErr_Clear();
1138
+ return -1;
1139
+ }
1140
+ if (si < 0) {
1141
+ return -1;
1142
+ }
1143
+ v = (npy_uint64)si;
1144
+ }
1145
+ else {
1146
+ return -1;
1147
+ }
1148
+ return lookup_hash_uint(self, v, uint_to_hash(v), self->keys_array_type);
1149
+ }
1150
+
1151
+
1152
+ static Py_ssize_t
1153
+ lookup_double(FAMObject *self, PyObject* key) {
1154
+ double v = 0;
1155
+ if (PyArray_IsScalar(key, Double)) {
1156
+ v = PyArrayScalar_VAL(key, Double);
1157
+ }
1158
+ else if (PyFloat_Check(key)) {
1159
+ v = PyFloat_AsDouble(key);
1160
+ if (v == -1.0 && PyErr_Occurred()) {
1161
+ PyErr_Clear();
1162
+ return -1;
1163
+ }
1164
+ }
1165
+ else if (PyLong_Check(key)) {
1166
+ v = (double)PyLong_AsLongLong(key);
1167
+ if (v == -1 && PyErr_Occurred()) {
1168
+ PyErr_Clear();
1169
+ return -1;
1170
+ }
1171
+ }
1172
+ // NOTE: we handle PyArray Scalar Byte, Short with PyNumber_Check, below, saving four branches here
1173
+ else if (PyArray_IsScalar(key, LongLong)) {
1174
+ v = (double)PyArrayScalar_VAL(key, LongLong);
1175
+ }
1176
+ else if (PyArray_IsScalar(key, Long)) {
1177
+ v = (double)PyArrayScalar_VAL(key, Long);
1178
+ }
1179
+ else if (PyArray_IsScalar(key, Int)) {
1180
+ v = (double)PyArrayScalar_VAL(key, Int);
1181
+ }
1182
+ else if (PyArray_IsScalar(key, ULongLong)) {
1183
+ v = (double)PyArrayScalar_VAL(key, ULongLong);
1184
+ }
1185
+ else if (PyArray_IsScalar(key, ULong)) {
1186
+ v = (double)PyArrayScalar_VAL(key, ULong);
1187
+ }
1188
+ else if (PyArray_IsScalar(key, UInt)) {
1189
+ v = (double)PyArrayScalar_VAL(key, UInt);
1190
+ }
1191
+ else if (PyArray_IsScalar(key, Float)) {
1192
+ v = (double)PyArrayScalar_VAL(key, Float);
1193
+ }
1194
+ else if (PyArray_IsScalar(key, Half)) {
1195
+ v = npy_half_to_double(PyArrayScalar_VAL(key, Half));
1196
+ }
1197
+ else if (PyBool_Check(key)) {
1198
+ v = PyObject_IsTrue(key);
1199
+ }
1200
+ else if (PyNumber_Check(key)) {
1201
+ // NOTE: this returns a Py_ssize_t, which might be 32 bit. This can be used for PyArray_Scalars <= ssize_t.
1202
+ npy_int64 si = PyNumber_AsSsize_t(key, PyExc_OverflowError);
1203
+ if (si == -1 && PyErr_Occurred()) {
1204
+ PyErr_Clear();
1205
+ return -1;
1206
+ }
1207
+ v = (double)si;
1208
+ }
1209
+ else {
1210
+ return -1;
1211
+ }
1212
+ return lookup_hash_double(self, v, double_to_hash(v), self->keys_array_type);
1213
+ }
1214
+
1215
+
1216
+ static Py_ssize_t
1217
+ lookup_unicode(FAMObject *self, PyObject* key) {
1218
+ // NOTE: while we can identify and use PyArray_IsScalar(key, Unicode), this did not improve performance and fails on Windows.
1219
+ if (!PyUnicode_Check(key)) {
1220
+ return -1;
1221
+ }
1222
+ PyArrayObject *a = (PyArrayObject *)self->keys;
1223
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
1224
+ // if the key_size is greater than the dtype size of the array, we know there cannot be a match
1225
+ Py_ssize_t k_size = PyUnicode_GetLength(key);
1226
+ if (k_size > dt_size) {
1227
+ return -1;
1228
+ }
1229
+ // The buffer will have dt_size + 1 storage. We copy a NULL character so do not have to clear the buffer, but instead can reuse it and still discover the lookup
1230
+ if (!PyUnicode_AsUCS4(key, self->key_buffer, dt_size+1, 1)) {
1231
+ return -1; // exception will be set
1232
+ }
1233
+ Py_hash_t hash = unicode_to_hash(self->key_buffer, k_size);
1234
+ return lookup_hash_unicode(self, self->key_buffer, k_size, hash);
1235
+ }
1236
+
1237
+
1238
+ static Py_ssize_t
1239
+ lookup_string(FAMObject *self, PyObject* key) {
1240
+ if (!PyBytes_Check(key)) {
1241
+ return -1;
1242
+ }
1243
+ PyArrayObject *a = (PyArrayObject *)self->keys;
1244
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a);
1245
+ Py_ssize_t k_size = PyBytes_GET_SIZE(key);
1246
+ if (k_size > dt_size) {
1247
+ return -1;
1248
+ }
1249
+ char* k = PyBytes_AS_STRING(key);
1250
+ Py_hash_t hash = string_to_hash(k, k_size);
1251
+ return lookup_hash_string(self, k, k_size, hash);
1252
+ }
1253
+
1254
+
1255
+ // Given a key as a PyObject, return the Py_ssize_t keys_pos value stored in the TableElement. Return -1 on key not found (without setting an exception) and -1 on error (with setting an exception).
1256
+ static Py_ssize_t
1257
+ lookup(FAMObject *self, PyObject *key) {
1258
+ Py_ssize_t table_pos = -1;
1259
+
1260
+ switch (self->keys_array_type) {
1261
+ case KAT_INT64:
1262
+ case KAT_INT32:
1263
+ case KAT_INT16:
1264
+ case KAT_INT8:
1265
+ table_pos = lookup_int(self, key);
1266
+ break;
1267
+ case KAT_UINT64:
1268
+ case KAT_UINT32:
1269
+ case KAT_UINT16:
1270
+ case KAT_UINT8:
1271
+ table_pos = lookup_uint(self, key);
1272
+ break;
1273
+ case KAT_FLOAT64:
1274
+ case KAT_FLOAT32:
1275
+ case KAT_FLOAT16:
1276
+ table_pos = lookup_double(self, key);
1277
+ break;
1278
+ case KAT_UNICODE:
1279
+ table_pos = lookup_unicode(self, key);
1280
+ break;
1281
+ case KAT_STRING:
1282
+ table_pos = lookup_string(self, key);
1283
+ break;
1284
+ case KAT_DTY:
1285
+ case KAT_DTM:
1286
+ case KAT_DTW:
1287
+ case KAT_DTD:
1288
+ case KAT_DTh:
1289
+ case KAT_DTm:
1290
+ case KAT_DTs:
1291
+ case KAT_DTms:
1292
+ case KAT_DTus:
1293
+ case KAT_DTns:
1294
+ case KAT_DTps:
1295
+ case KAT_DTfs:
1296
+ case KAT_DTas:
1297
+ table_pos = lookup_datetime(self, key);
1298
+ break;
1299
+ case KAT_LIST: {
1300
+ Py_hash_t hash = PyObject_Hash(key);
1301
+ if (hash == -1) {
1302
+ return -1;
1303
+ }
1304
+ table_pos = lookup_hash_obj(self, key, hash);
1305
+ break;
1306
+ }
1307
+ }
1308
+ // A -1 hash is an unused storage location
1309
+ if ((table_pos < 0) || (self->table[table_pos].hash == -1)) {
1310
+ return -1;
1311
+ }
1312
+ return self->table[table_pos].keys_pos;
1313
+ }
1314
+
1315
+ // Insert a key_pos, hash pair into the table. Assumes table already has appropriate size. When inserting a new item, `hash` is -1, forcing a fresh hash to be computed here. Return 0 on success, -1 on error.
1316
+ static int
1317
+ insert_obj(
1318
+ FAMObject *self,
1319
+ PyObject *key, // NOTE: a borrowed reference
1320
+ Py_ssize_t keys_pos,
1321
+ Py_hash_t hash)
1322
+ {
1323
+ if (hash == -1) {
1324
+ hash = PyObject_Hash(key);
1325
+ if (hash == -1) {
1326
+ return -1;
1327
+ }
1328
+ }
1329
+ // table position is not dependent on keys_pos
1330
+ Py_ssize_t table_pos = lookup_hash_obj(self, key, hash);
1331
+
1332
+ if (table_pos < 0) {
1333
+ return -1;
1334
+ }
1335
+ // We expect, on insertion, to get back a table_pos that points to an unassigned hash value (-1); if we get anything else, we have found a match to an already-existing key, and thus raise a NonUniqueError error.
1336
+ if (self->table[table_pos].hash != -1) {
1337
+ PyErr_SetObject(NonUniqueError, key);
1338
+ return -1;
1339
+ }
1340
+ self->table[table_pos].keys_pos = keys_pos;
1341
+ self->table[table_pos].hash = hash;
1342
+ return 0;
1343
+ }
1344
+
1345
+
1346
+ static int
1347
+ insert_int(
1348
+ FAMObject *self,
1349
+ npy_int64 key,
1350
+ Py_ssize_t keys_pos,
1351
+ Py_hash_t hash,
1352
+ KeysArrayType kat)
1353
+ {
1354
+ if (hash == -1) {
1355
+ hash = int_to_hash(key);
1356
+ }
1357
+ // table position is not dependent on keys_pos
1358
+ Py_ssize_t table_pos = lookup_hash_int(self, key, hash, kat);
1359
+ if (table_pos < 0) {
1360
+ return -1;
1361
+ }
1362
+ if (self->table[table_pos].hash != -1) {
1363
+ PyObject* er = PyLong_FromLongLong(key); // for error reporting
1364
+ if (er == NULL) {
1365
+ return -1;
1366
+ }
1367
+ PyErr_SetObject(NonUniqueError, er);
1368
+ Py_DECREF(er);
1369
+ return -1;
1370
+ }
1371
+ self->table[table_pos].keys_pos = keys_pos;
1372
+ self->table[table_pos].hash = hash; // key is the hash
1373
+ return 0;
1374
+ }
1375
+
1376
+
1377
+ static int
1378
+ insert_uint(
1379
+ FAMObject *self,
1380
+ npy_uint64 key,
1381
+ Py_ssize_t keys_pos,
1382
+ Py_hash_t hash,
1383
+ KeysArrayType kat)
1384
+ {
1385
+ if (hash == -1) {
1386
+ hash = uint_to_hash(key);
1387
+ }
1388
+ Py_ssize_t table_pos = lookup_hash_uint(self, key, hash, kat);
1389
+
1390
+ if (table_pos < 0) {
1391
+ return -1;
1392
+ }
1393
+ if (self->table[table_pos].hash != -1) {
1394
+ PyObject* er = PyLong_FromUnsignedLongLong(key);
1395
+ if (er == NULL) {
1396
+ return -1;
1397
+ }
1398
+ PyErr_SetObject(NonUniqueError, er);
1399
+ Py_DECREF(er);
1400
+ return -1;
1401
+ }
1402
+ self->table[table_pos].keys_pos = keys_pos;
1403
+ self->table[table_pos].hash = hash;
1404
+ return 0;
1405
+ }
1406
+
1407
+
1408
+ static int
1409
+ insert_double(
1410
+ FAMObject *self,
1411
+ npy_double key,
1412
+ Py_ssize_t keys_pos,
1413
+ Py_hash_t hash,
1414
+ KeysArrayType kat)
1415
+ {
1416
+ if (hash == -1) {
1417
+ hash = double_to_hash(key);
1418
+ }
1419
+ // table position is not dependent on keys_pos
1420
+ Py_ssize_t table_pos = lookup_hash_double(self, key, hash, kat);
1421
+
1422
+ if (table_pos < 0) {
1423
+ return -1;
1424
+ }
1425
+ if (self->table[table_pos].hash != -1) {
1426
+ PyObject* er = PyFloat_FromDouble(key);
1427
+ if (er == NULL) {
1428
+ return -1;
1429
+ }
1430
+ PyErr_SetObject(NonUniqueError, er);
1431
+ Py_DECREF(er);
1432
+ return -1;
1433
+ }
1434
+ self->table[table_pos].keys_pos = keys_pos;
1435
+ self->table[table_pos].hash = hash;
1436
+ return 0;
1437
+ }
1438
+
1439
+
1440
+ static int
1441
+ insert_unicode(
1442
+ FAMObject *self,
1443
+ Py_UCS4* key,
1444
+ Py_ssize_t key_size,
1445
+ Py_ssize_t keys_pos,
1446
+ Py_hash_t hash)
1447
+ {
1448
+ if (hash == -1) {
1449
+ hash = unicode_to_hash(key, key_size);
1450
+ }
1451
+ // table position is not dependent on keys_pos
1452
+ Py_ssize_t table_pos = lookup_hash_unicode(self, key, key_size, hash);
1453
+ if (table_pos < 0) {
1454
+ return -1;
1455
+ }
1456
+ if (self->table[table_pos].hash != -1) {
1457
+ PyObject* er = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, key, key_size);
1458
+ if (er == NULL) {
1459
+ return -1;
1460
+ }
1461
+ PyErr_SetObject(NonUniqueError, er);
1462
+ Py_DECREF(er);
1463
+ return -1;
1464
+ }
1465
+ self->table[table_pos].keys_pos = keys_pos;
1466
+ self->table[table_pos].hash = hash;
1467
+ return 0;
1468
+ }
1469
+
1470
+
1471
+ static int
1472
+ insert_string(
1473
+ FAMObject *self,
1474
+ char* key,
1475
+ Py_ssize_t key_size,
1476
+ Py_ssize_t keys_pos,
1477
+ Py_hash_t hash)
1478
+ {
1479
+ if (hash == -1) {
1480
+ hash = string_to_hash(key, key_size);
1481
+ }
1482
+ // table position is not dependent on keys_pos
1483
+ Py_ssize_t table_pos = lookup_hash_string(self, key, key_size, hash);
1484
+ if (table_pos < 0) {
1485
+ return -1;
1486
+ }
1487
+ if (self->table[table_pos].hash != -1) {
1488
+ PyObject* er = PyBytes_FromStringAndSize(key, key_size);
1489
+ if (er == NULL) {
1490
+ return -1;
1491
+ }
1492
+ PyErr_SetObject(NonUniqueError, er);
1493
+ Py_DECREF(er);
1494
+ return -1;
1495
+ }
1496
+ self->table[table_pos].keys_pos = keys_pos;
1497
+ self->table[table_pos].hash = hash;
1498
+ return 0;
1499
+ }
1500
+
1501
+
1502
+ //------------------------------------------------------------------------------
1503
+
1504
+ // Called in fam_new(), extend(), append(), with the size of observed keys. This table is updated only when append or extending. Only if there is an old table will keys be accessed Returns 0 on success, -1 on failure.
1505
+ static int
1506
+ grow_table(FAMObject *self, Py_ssize_t keys_size)
1507
+ {
1508
+ Py_ssize_t keys_load = keys_size / LOAD;
1509
+ Py_ssize_t size_old = self->table_size;
1510
+ if (keys_load < size_old) {
1511
+ return 0;
1512
+ }
1513
+
1514
+ // get the next power of 2 greater than current keys_load
1515
+ Py_ssize_t size_new = 1;
1516
+ while (size_new <= keys_load) {
1517
+ size_new <<= 1;
1518
+ }
1519
+ // size_new > keys_load; we know that keys_load >= size_old, so size_new must be > size_old
1520
+ TableElement *table_old = self->table;
1521
+ TableElement *table_new = PyMem_New(TableElement, size_new + SCAN - 1);
1522
+ if (!table_new) {
1523
+ return -1;
1524
+ }
1525
+
1526
+ // initialize all hash and keys_pos values to -1
1527
+ Py_ssize_t table_pos;
1528
+ for (table_pos = 0; table_pos < size_new + SCAN - 1; table_pos++) {
1529
+ table_new[table_pos].hash = -1;
1530
+ table_new[table_pos].keys_pos = -1;
1531
+ }
1532
+ self->table = table_new;
1533
+ self->table_size = size_new;
1534
+
1535
+ // if we have an old table, move them into the new table
1536
+ if (size_old) {
1537
+ if (self->keys_array_type) {
1538
+ PyErr_SetString(PyExc_NotImplementedError, "Cannot grow table for array keys");
1539
+ goto restore;
1540
+ }
1541
+ Py_ssize_t i;
1542
+ Py_hash_t h;
1543
+ for (table_pos = 0; table_pos < size_old + SCAN - 1; table_pos++) {
1544
+ i = table_old[table_pos].keys_pos;
1545
+ h = table_old[table_pos].hash;
1546
+ if ((h != -1) && insert_obj(self, PyList_GET_ITEM(self->keys, i), i, h))
1547
+ {
1548
+ goto restore;
1549
+ }
1550
+ }
1551
+ }
1552
+ PyMem_Del(table_old);
1553
+ return 0;
1554
+ restore:
1555
+ PyMem_Del(self->table);
1556
+ self->table = table_old;
1557
+ self->table_size = size_old;
1558
+ return -1;
1559
+ }
1560
+
1561
+
1562
+ // Given a new, possibly un-initialized FAMObject, copy attrs from self to new. Return 0 on success, -1 on error.
1563
+ int
1564
+ copy_to_new(PyTypeObject *cls, FAMObject *self, FAMObject *new)
1565
+ {
1566
+ if (self->keys_array_type) {
1567
+ new->keys = self->keys;
1568
+ Py_INCREF(new->keys);
1569
+ }
1570
+ else {
1571
+ new->keys = PySequence_List(self->keys);
1572
+ if (!new->keys) {
1573
+ return -1;
1574
+ }
1575
+ }
1576
+ new->table_size = self->table_size;
1577
+ new->keys_array_type = self->keys_array_type;
1578
+ new->keys_size = self->keys_size;
1579
+
1580
+ new->key_buffer = NULL;
1581
+ if (new->keys_array_type == KAT_UNICODE) {
1582
+ PyArrayObject *a = (PyArrayObject *)new->keys;
1583
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
1584
+ new->key_buffer = (Py_UCS4*)PyMem_Malloc((dt_size+1) * UCS4_SIZE);
1585
+ }
1586
+
1587
+ Py_ssize_t table_size_alloc = new->table_size + SCAN - 1;
1588
+ new->table = PyMem_New(TableElement, table_size_alloc);
1589
+ if (!new->table) {
1590
+ // Py_DECREF(new->keys); // assume this will get cleaned up
1591
+ return -1;
1592
+ }
1593
+ memcpy(new->table, self->table, table_size_alloc * sizeof(TableElement));
1594
+ return 0;
1595
+ }
1596
+
1597
+
1598
+ static PyObject *
1599
+ fam_new(PyTypeObject *cls, PyObject *args, PyObject *kwargs);
1600
+
1601
+
1602
+ // Create a copy of self. Used in `fam_or()`. Returns a new FAMObject on success, NULL on error.
1603
+ static FAMObject *
1604
+ copy(PyTypeObject *cls, FAMObject *self)
1605
+ {
1606
+ if (!PyType_IsSubtype(cls, &AMType) && !PyObject_TypeCheck(self, &AMType)) {
1607
+ Py_INCREF(self);
1608
+ return self;
1609
+ }
1610
+ // fam_new to allocate and full struct attrs
1611
+ FAMObject *new = (FAMObject*)fam_new(cls, NULL, NULL);
1612
+ if (!new) {
1613
+ return NULL;
1614
+ }
1615
+ if (copy_to_new(cls, self, new)) {
1616
+ Py_DECREF(new); // assume this will decref any partially set attrs of new
1617
+ }
1618
+ return new;
1619
+ }
1620
+
1621
+
1622
+
1623
+ // Returns -1 on error, 0 on success.
1624
+ static int
1625
+ extend(FAMObject *self, PyObject *keys)
1626
+ {
1627
+ if (self->keys_array_type) {
1628
+ PyErr_SetString(PyExc_NotImplementedError, "Not supported for array keys");
1629
+ return -1;
1630
+ }
1631
+ // this should fail for self->keys types that are not a list
1632
+ keys = PySequence_Fast(keys, "expected an iterable of keys");
1633
+ if (!keys) {
1634
+ return -1;
1635
+ }
1636
+ Py_ssize_t size_extend = PySequence_Fast_GET_SIZE(keys);
1637
+ self->keys_size += size_extend;
1638
+
1639
+ if (grow_table(self, self->keys_size)) {
1640
+ Py_DECREF(keys);
1641
+ return -1;
1642
+ }
1643
+
1644
+ PyObject **keys_fi = PySequence_Fast_ITEMS(keys);
1645
+
1646
+ for (Py_ssize_t index = 0; index < size_extend; index++) {
1647
+ // get the new keys_size after each append
1648
+ if (insert_obj(self, keys_fi[index], PyList_GET_SIZE(self->keys), -1) ||
1649
+ PyList_Append(self->keys, keys_fi[index]))
1650
+ {
1651
+ Py_DECREF(keys);
1652
+ return -1;
1653
+ }
1654
+ }
1655
+ Py_DECREF(keys);
1656
+ return 0;
1657
+ }
1658
+
1659
+
1660
+ // Returns -1 on error, 0 on success.
1661
+ static int
1662
+ append(FAMObject *self, PyObject *key)
1663
+ {
1664
+ if (self->keys_array_type) {
1665
+ PyErr_SetString(PyExc_NotImplementedError, "Not supported for array keys");
1666
+ return -1;
1667
+ }
1668
+ self->keys_size++;
1669
+
1670
+ if (grow_table(self, self->keys_size)) {
1671
+ return -1;
1672
+ }
1673
+ // keys_size is already incremented; provide last index
1674
+ if (insert_obj(self, key, self->keys_size - 1, -1) ||
1675
+ PyList_Append(self->keys, key))
1676
+ {
1677
+ return -1;
1678
+ }
1679
+ return 0;
1680
+ }
1681
+
1682
+
1683
+ static Py_ssize_t
1684
+ fam_length(FAMObject *self)
1685
+ {
1686
+ return self->keys_size;
1687
+ }
1688
+
1689
+
1690
+ // Given a key for a FAM, return the Python integer associated with that key. Utility function used in both fam_subscript() and fam_get()
1691
+ static PyObject *
1692
+ get(FAMObject *self, PyObject *key, PyObject *missing) {
1693
+ Py_ssize_t keys_pos = lookup(self, key);
1694
+ if (keys_pos < 0) {
1695
+ if (PyErr_Occurred()) {
1696
+ return NULL;
1697
+ }
1698
+ if (missing) {
1699
+ Py_INCREF(missing);
1700
+ return missing;
1701
+ }
1702
+ PyErr_SetObject(PyExc_KeyError, key);
1703
+ return NULL;
1704
+ }
1705
+ return PyLong_FromSsize_t(keys_pos);
1706
+ }
1707
+
1708
+
1709
+ // Give an array of the same kind as KAT, lookup and load all keys_pos. Depends on self, key_size, key_array, table_pos, i, k, b
1710
+ # define GET_ALL_SCALARS(npy_type_src, npy_type_dst, kat, lookup_func, hash_func, to_obj_func, post_deref) \
1711
+ { \
1712
+ npy_type_dst v; \
1713
+ for (; i < key_size; i++) { \
1714
+ v = post_deref(*(npy_type_src*)PyArray_GETPTR1(key_array, i)); \
1715
+ table_pos = lookup_func(self, v, hash_func(v), kat); \
1716
+ if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1717
+ Py_DECREF(array); \
1718
+ if (PyErr_Occurred()) { \
1719
+ return NULL; \
1720
+ } \
1721
+ k = to_obj_func(v); \
1722
+ if (k == NULL) { \
1723
+ return NULL; \
1724
+ } \
1725
+ PyErr_SetObject(PyExc_KeyError, k); \
1726
+ Py_DECREF(k); \
1727
+ return NULL; \
1728
+ } \
1729
+ b[i] = (npy_int64)self->table[table_pos].keys_pos; \
1730
+ } \
1731
+ } \
1732
+
1733
+ # define GET_ALL_DT64(npy_type_src, npy_type_dst, kat, lookup_func, hash_func) \
1734
+ { \
1735
+ npy_type_dst v; \
1736
+ for (; i < key_size; i++) { \
1737
+ v = *(npy_type_src*)PyArray_GETPTR1(key_array, i); \
1738
+ table_pos = lookup_func(self, v, hash_func(v), kat); \
1739
+ if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1740
+ Py_DECREF(array); \
1741
+ if (PyErr_Occurred()) { \
1742
+ return NULL; \
1743
+ } \
1744
+ k = PyArray_ToScalar(&v, key_array); \
1745
+ if (k == NULL) { \
1746
+ return NULL; \
1747
+ } \
1748
+ PyErr_SetObject(PyExc_KeyError, k); \
1749
+ Py_DECREF(k); \
1750
+ return NULL; \
1751
+ } \
1752
+ b[i] = (npy_int64)self->table[table_pos].keys_pos; \
1753
+ } \
1754
+ } \
1755
+
1756
+ # define GET_ALL_FLEXIBLE(char_type, get_end_func, lookup_func, hash_func, to_obj_func) \
1757
+ { \
1758
+ char_type* v; \
1759
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(key_array) / sizeof(char_type);\
1760
+ Py_ssize_t k_size; \
1761
+ for (; i < key_size; i++) { \
1762
+ v = (char_type*)PyArray_GETPTR1(key_array, i); \
1763
+ k_size = get_end_func(v, dt_size) - v; \
1764
+ table_pos = lookup_func(self, v, k_size, hash_func(v, k_size)); \
1765
+ if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1766
+ Py_DECREF(array); \
1767
+ if (PyErr_Occurred()) { \
1768
+ return NULL; \
1769
+ } \
1770
+ k = to_obj_func(v, k_size); \
1771
+ if (k == NULL) { \
1772
+ return NULL; \
1773
+ } \
1774
+ PyErr_SetObject(PyExc_KeyError, k); \
1775
+ Py_DECREF(k); \
1776
+ return NULL; \
1777
+ } \
1778
+ b[i] = (npy_int64)self->table[table_pos].keys_pos; \
1779
+ } \
1780
+ } \
1781
+
1782
+ // Given a list or array of keys, return an array of the lookup-up integer values. If any unmatched keys are found, a KeyError will raise. An immutable array is always returned.
1783
+ static PyObject *
1784
+ fam_get_all(FAMObject *self, PyObject *key) {
1785
+ Py_ssize_t key_size = 0;
1786
+ Py_ssize_t keys_pos = -1;
1787
+ PyObject* k = NULL;
1788
+ PyObject *array = NULL;
1789
+ Py_ssize_t i = 0;
1790
+
1791
+ int key_is_list;
1792
+ if (PyList_CheckExact(key)) {
1793
+ key_is_list = 1;
1794
+ key_size = PyList_GET_SIZE(key);
1795
+ }
1796
+ else if (PyArray_Check(key)) {
1797
+ key_is_list = 0;
1798
+ key_size = PyArray_SIZE((PyArrayObject *)key);
1799
+ }
1800
+ else {
1801
+ PyErr_SetString(PyExc_TypeError, "Must provide a list or array.");
1802
+ return NULL;
1803
+ }
1804
+
1805
+ // construct array to be returned; this is a little expensive if we do not yet know if we can use it
1806
+ npy_intp dims[] = {key_size};
1807
+ array = PyArray_EMPTY(1, dims, NPY_INT64, 0);
1808
+ if (array == NULL) {
1809
+ return NULL;
1810
+ }
1811
+ npy_int64* b = (npy_int64*)PyArray_DATA((PyArrayObject*)array);
1812
+
1813
+ if (key_is_list) {
1814
+ for (; i < key_size; i++) {
1815
+ k = PyList_GET_ITEM(key, i); // borrow
1816
+ keys_pos = lookup(self, k);
1817
+ if (keys_pos < 0) {
1818
+ Py_DECREF(array);
1819
+ if (PyErr_Occurred()) {
1820
+ return NULL;
1821
+ }
1822
+ PyErr_SetObject(PyExc_KeyError, k);
1823
+ return NULL;
1824
+ }
1825
+ b[i] = (npy_int64)keys_pos;
1826
+ }
1827
+ }
1828
+ else { // key is an array
1829
+ PyArrayObject* key_array = (PyArrayObject *)key;
1830
+ // if key is an np array of the same kind as this FAMs keys, we can do optimized lookups; otherwise, we have to go through scalar to do full branching and coercion into lookup
1831
+ int key_array_t = PyArray_TYPE(key_array);
1832
+
1833
+ // NOTE: we only match numeric kinds of the KAT is 64 bit; we could support, for each key_array_t, a switch for every KAT, but the size of that code is huge and the performance benefit is not massive
1834
+ if (kat_is_kind(self->keys_array_type, PyArray_DESCR(key_array)->kind)) {
1835
+ Py_ssize_t table_pos;
1836
+ switch (key_array_t) { // type of passed in array
1837
+ case NPY_INT64:
1838
+ GET_ALL_SCALARS(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash, PyLong_FromLongLong,);
1839
+ break;
1840
+ case NPY_INT32:
1841
+ GET_ALL_SCALARS(npy_int32, npy_int64, KAT_INT32, lookup_hash_int, int_to_hash, PyLong_FromLongLong,);
1842
+ break;
1843
+ case NPY_INT16:
1844
+ GET_ALL_SCALARS(npy_int16, npy_int64, KAT_INT16, lookup_hash_int, int_to_hash, PyLong_FromLongLong,);
1845
+ break;
1846
+ case NPY_INT8:
1847
+ GET_ALL_SCALARS(npy_int8, npy_int64, KAT_INT8, lookup_hash_int, int_to_hash, PyLong_FromLongLong,);
1848
+ break;
1849
+ case NPY_UINT64:
1850
+ GET_ALL_SCALARS(npy_uint64, npy_uint64, KAT_UINT64, lookup_hash_uint, uint_to_hash, PyLong_FromUnsignedLongLong,);
1851
+ break;
1852
+ case NPY_UINT32:
1853
+ GET_ALL_SCALARS(npy_uint32, npy_uint64, KAT_UINT32, lookup_hash_uint, uint_to_hash, PyLong_FromUnsignedLongLong,);
1854
+ break;
1855
+ case NPY_UINT16:
1856
+ GET_ALL_SCALARS(npy_uint16, npy_uint64, KAT_UINT16, lookup_hash_uint, uint_to_hash, PyLong_FromUnsignedLongLong,);
1857
+ break;
1858
+ case NPY_UINT8:
1859
+ GET_ALL_SCALARS(npy_uint8, npy_uint64, KAT_UINT8, lookup_hash_uint, uint_to_hash, PyLong_FromUnsignedLongLong,);
1860
+ break;
1861
+ case NPY_FLOAT64:
1862
+ GET_ALL_SCALARS(npy_double, npy_double, KAT_FLOAT64, lookup_hash_double, double_to_hash, PyFloat_FromDouble,);
1863
+ break;
1864
+ case NPY_FLOAT32:
1865
+ GET_ALL_SCALARS(npy_float, npy_double, KAT_FLOAT32, lookup_hash_double, double_to_hash, PyFloat_FromDouble,);
1866
+ break;
1867
+ case NPY_FLOAT16:
1868
+ GET_ALL_SCALARS(npy_half, npy_double, KAT_FLOAT16, lookup_hash_double, double_to_hash, PyFloat_FromDouble, npy_half_to_double);
1869
+ break;
1870
+ case NPY_UNICODE:
1871
+ GET_ALL_FLEXIBLE(Py_UCS4, ucs4_get_end_p, lookup_hash_unicode, unicode_to_hash, PyUnicode_FromUCS4AndData);
1872
+ break;
1873
+ case NPY_STRING:
1874
+ GET_ALL_FLEXIBLE(char, char_get_end_p, lookup_hash_string, string_to_hash, PyBytes_FromStringAndSize);
1875
+ break;
1876
+ case NPY_DATETIME: {
1877
+ NPY_DATETIMEUNIT key_unit = AK_dt_unit_from_array(key_array);
1878
+ if (!kat_is_datetime_unit(self->keys_array_type, key_unit)) {
1879
+ PyErr_SetString(PyExc_KeyError, "datetime64 units do not match");
1880
+ Py_DECREF(array);
1881
+ return NULL;
1882
+ }
1883
+ GET_ALL_DT64(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash);
1884
+ break;
1885
+ }
1886
+ }
1887
+ }
1888
+ else {
1889
+ for (; i < key_size; i++) {
1890
+ k = PyArray_ToScalar(PyArray_GETPTR1(key_array, i), key_array);
1891
+ if (k == NULL) {
1892
+ Py_DECREF(array);
1893
+ return NULL;
1894
+ }
1895
+ keys_pos = lookup(self, k);
1896
+ if (keys_pos < 0) {
1897
+ Py_DECREF(array);
1898
+ if (PyErr_Occurred()) {
1899
+ Py_DECREF(k);
1900
+ return NULL;
1901
+ }
1902
+ PyErr_SetObject(PyExc_KeyError, k);
1903
+ Py_DECREF(k);
1904
+ return NULL;
1905
+ }
1906
+ Py_DECREF(k);
1907
+ b[i] = (npy_int64)keys_pos;
1908
+ }
1909
+ }
1910
+ }
1911
+
1912
+ PyArray_CLEARFLAGS((PyArrayObject *)array, NPY_ARRAY_WRITEABLE);
1913
+ return array;
1914
+
1915
+ }
1916
+
1917
+
1918
+ # undef GET_ALL_SCALARS
1919
+ # undef GET_ALL_FLEXIBLE
1920
+
1921
+
1922
+ static inline int
1923
+ append_ssize_t(
1924
+ PyObject* list,
1925
+ Py_ssize_t value)
1926
+ {
1927
+ PyObject* v = PyLong_FromSsize_t(value);
1928
+ if (v == NULL) {
1929
+ return -1;
1930
+ }
1931
+ int err = PyList_Append(list, v);
1932
+ Py_DECREF(v);
1933
+ return err;
1934
+ }
1935
+
1936
+ // Give an array of the same kind as KAT, lookup and load any keys_pos. Depends on self, key_size, key_array, table_pos, i, k, values
1937
+ # define GET_ANY_SCALARS(npy_type_src, npy_type_dst, kat, lookup_func, hash_func, post_deref) \
1938
+ { \
1939
+ npy_type_dst v; \
1940
+ for (; i < key_size; i++) { \
1941
+ v = post_deref(*(npy_type_src*)PyArray_GETPTR1(key_array, i)); \
1942
+ table_pos = lookup_func(self, v, hash_func(v), kat); \
1943
+ if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1944
+ if (PyErr_Occurred()) { \
1945
+ Py_DECREF(values); \
1946
+ return NULL; \
1947
+ } \
1948
+ continue; \
1949
+ } \
1950
+ keys_pos = self->table[table_pos].keys_pos; \
1951
+ if (append_ssize_t(values, keys_pos)) { \
1952
+ Py_DECREF(values); \
1953
+ return NULL; \
1954
+ } \
1955
+ } \
1956
+ } \
1957
+
1958
+ # define GET_ANY_FLEXIBLE(char_type, get_end_func, lookup_func, hash_func) \
1959
+ { \
1960
+ char_type* v; \
1961
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(key_array) / sizeof(char_type);\
1962
+ Py_ssize_t k_size; \
1963
+ for (; i < key_size; i++) { \
1964
+ v = (char_type*)PyArray_GETPTR1(key_array, i); \
1965
+ k_size = get_end_func(v, dt_size) - v; \
1966
+ table_pos = lookup_func(self, v, k_size, hash_func(v, k_size)); \
1967
+ if (table_pos < 0 || (self->table[table_pos].hash == -1)) { \
1968
+ if (PyErr_Occurred()) { \
1969
+ Py_DECREF(values); \
1970
+ return NULL; \
1971
+ } \
1972
+ continue; \
1973
+ } \
1974
+ keys_pos = self->table[table_pos].keys_pos; \
1975
+ if (append_ssize_t(values, keys_pos)) { \
1976
+ Py_DECREF(values); \
1977
+ return NULL; \
1978
+ } \
1979
+ } \
1980
+ } \
1981
+
1982
+ // Given a list or array of keys, return a list of the lookup-up integer values. If any unmatched keys are found, they are ignored. A list is always returned.
1983
+ static PyObject *
1984
+ fam_get_any(FAMObject *self, PyObject *key) {
1985
+ Py_ssize_t key_size = 0;
1986
+ Py_ssize_t keys_pos = -1;
1987
+ Py_ssize_t i = 0;
1988
+ PyObject* k = NULL;
1989
+ PyObject* values = NULL;
1990
+
1991
+ int key_is_list;
1992
+ if (PyList_CheckExact(key)) {
1993
+ key_is_list = 1;
1994
+ key_size = PyList_GET_SIZE(key);
1995
+ }
1996
+ else if (PyArray_Check(key)) {
1997
+ key_is_list = 0;
1998
+ key_size = PyArray_SIZE((PyArrayObject *)key);
1999
+ }
2000
+ else {
2001
+ PyErr_SetString(PyExc_TypeError, "Must provide a list or array.");
2002
+ return NULL;
2003
+ }
2004
+
2005
+ values = PyList_New(0);
2006
+ if (!values) {
2007
+ return NULL;
2008
+ }
2009
+
2010
+ if (key_is_list) {
2011
+ for (; i < key_size; i++) {
2012
+ k = PyList_GET_ITEM(key, i); // borrow
2013
+ keys_pos = lookup(self, k);
2014
+ if (keys_pos < 0) {
2015
+ if (PyErr_Occurred()) { // only exit if exception set
2016
+ Py_DECREF(values);
2017
+ return NULL;
2018
+ }
2019
+ continue;
2020
+ }
2021
+ if (append_ssize_t(values, keys_pos)) {
2022
+ Py_DECREF(values);
2023
+ return NULL;
2024
+ }
2025
+ }
2026
+ }
2027
+ else {
2028
+ PyArrayObject* key_array = (PyArrayObject *)key;
2029
+ // if key is an np array of the same kind as this FAMs keys, we can do optimized lookups; otherwise, we have to go through scalar to do full branching and coercion into lookup
2030
+ int key_array_t = PyArray_TYPE(key_array);
2031
+
2032
+ if (kat_is_kind(self->keys_array_type, PyArray_DESCR(key_array)->kind)) {
2033
+ Py_ssize_t table_pos;
2034
+ switch (key_array_t) {
2035
+ case NPY_INT64:
2036
+ GET_ANY_SCALARS(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash,);
2037
+ break;
2038
+ case NPY_INT32:
2039
+ GET_ANY_SCALARS(npy_int32, npy_int64, KAT_INT32, lookup_hash_int, int_to_hash,);
2040
+ break;
2041
+ case NPY_INT16:
2042
+ GET_ANY_SCALARS(npy_int16, npy_int64, KAT_INT16, lookup_hash_int, int_to_hash,);
2043
+ break;
2044
+ case NPY_INT8:
2045
+ GET_ANY_SCALARS(npy_int8, npy_int64, KAT_INT8, lookup_hash_int, int_to_hash,);
2046
+ break;
2047
+ case NPY_UINT64:
2048
+ GET_ANY_SCALARS(npy_uint64, npy_uint64, KAT_UINT64, lookup_hash_uint, uint_to_hash,);
2049
+ break;
2050
+ case NPY_UINT32:
2051
+ GET_ANY_SCALARS(npy_uint32, npy_uint64, KAT_UINT32, lookup_hash_uint, uint_to_hash,);
2052
+ break;
2053
+ case NPY_UINT16:
2054
+ GET_ANY_SCALARS(npy_uint16, npy_uint64, KAT_UINT16, lookup_hash_uint, uint_to_hash,);
2055
+ break;
2056
+ case NPY_UINT8:
2057
+ GET_ANY_SCALARS(npy_uint8, npy_uint64, KAT_UINT8, lookup_hash_uint, uint_to_hash,);
2058
+ break;
2059
+ case NPY_FLOAT64:
2060
+ GET_ANY_SCALARS(npy_double, npy_double, KAT_FLOAT64, lookup_hash_double, double_to_hash,);
2061
+ break;
2062
+ case NPY_FLOAT32:
2063
+ GET_ANY_SCALARS(npy_float, npy_double, KAT_FLOAT32, lookup_hash_double, double_to_hash,);
2064
+ break;
2065
+ case NPY_FLOAT16:
2066
+ GET_ANY_SCALARS(npy_half, npy_double, KAT_FLOAT16, lookup_hash_double, double_to_hash, npy_half_to_double);
2067
+ break;
2068
+ case NPY_UNICODE:
2069
+ GET_ANY_FLEXIBLE(Py_UCS4, ucs4_get_end_p, lookup_hash_unicode, unicode_to_hash);
2070
+ break;
2071
+ case NPY_STRING:
2072
+ GET_ANY_FLEXIBLE(char, char_get_end_p, lookup_hash_string, string_to_hash);
2073
+ break;
2074
+ case NPY_DATETIME: {
2075
+ NPY_DATETIMEUNIT key_unit = AK_dt_unit_from_array(key_array);
2076
+ if (!kat_is_datetime_unit(self->keys_array_type, key_unit)) {
2077
+ return values;
2078
+ }
2079
+ GET_ANY_SCALARS(npy_int64, npy_int64, KAT_INT64, lookup_hash_int, int_to_hash,);
2080
+ break;
2081
+ }
2082
+ }
2083
+ }
2084
+ else {
2085
+ for (; i < key_size; i++) {
2086
+ k = PyArray_ToScalar(PyArray_GETPTR1(key_array, i), key_array);
2087
+ if (k == NULL) {
2088
+ Py_DECREF(values);
2089
+ return NULL;
2090
+ }
2091
+ keys_pos = lookup(self, k);
2092
+ Py_DECREF(k);
2093
+ if (keys_pos < 0) {
2094
+ if (PyErr_Occurred()) { // only exit if exception set
2095
+ Py_DECREF(values);
2096
+ return NULL;
2097
+ }
2098
+ continue; // do not raise
2099
+ }
2100
+ if (append_ssize_t(values, keys_pos)) {
2101
+ Py_DECREF(values);
2102
+ return NULL;
2103
+ }
2104
+ }
2105
+ }
2106
+ }
2107
+ return values; // might be empty
2108
+ }
2109
+
2110
+
2111
+ # undef GET_ANY_SCALARS
2112
+ # undef GET_ANY_FLEXIBLE
2113
+
2114
+
2115
+ static PyObject *
2116
+ fam_subscript(FAMObject *self, PyObject *key)
2117
+ {
2118
+ return get(self, key, NULL);
2119
+ }
2120
+
2121
+
2122
+ static PyMappingMethods fam_as_mapping = {
2123
+ .mp_length = (lenfunc) fam_length,
2124
+ .mp_subscript = (binaryfunc) fam_subscript,
2125
+ };
2126
+
2127
+
2128
+ static PyObject *
2129
+ fam_or(PyObject *left, PyObject *right)
2130
+ {
2131
+ if (!PyObject_TypeCheck(left, &FAMType) ||
2132
+ !PyObject_TypeCheck(right, &FAMType)
2133
+ ) {
2134
+ Py_RETURN_NOTIMPLEMENTED;
2135
+ }
2136
+ FAMObject *updated = copy(Py_TYPE(left), (FAMObject *)left);
2137
+ if (!updated) {
2138
+ return NULL;
2139
+ }
2140
+ if (extend(updated, ((FAMObject *)right)->keys)) {
2141
+ Py_DECREF(updated);
2142
+ return NULL;
2143
+ }
2144
+ return (PyObject *)updated;
2145
+ }
2146
+
2147
+
2148
+ static PyNumberMethods fam_as_number = {
2149
+ .nb_or = (binaryfunc) fam_or,
2150
+ };
2151
+
2152
+
2153
+ static int
2154
+ fam_contains(FAMObject *self, PyObject *key)
2155
+ {
2156
+ if (lookup(self, key) < 0) {
2157
+ if (PyErr_Occurred()) {
2158
+ return -1;
2159
+ }
2160
+ return 0;
2161
+ }
2162
+ return 1;
2163
+ }
2164
+
2165
+
2166
+ static PySequenceMethods fam_as_sequence = {
2167
+ .sq_contains = (objobjproc) fam_contains,
2168
+ };
2169
+
2170
+
2171
+ static void
2172
+ fam_dealloc(FAMObject *self)
2173
+ {
2174
+ if (self->table) {
2175
+ PyMem_Free(self->table);
2176
+ }
2177
+ if (self->key_buffer) {
2178
+ PyMem_Free(self->key_buffer);
2179
+ }
2180
+ if (self->keys) {
2181
+ Py_DECREF(self->keys);
2182
+ }
2183
+ Py_TYPE(self)->tp_free((PyObject *)self);
2184
+ }
2185
+
2186
+
2187
+ // Return a hash integer for an entire FAM by combining all stored hashes
2188
+ static Py_hash_t
2189
+ fam_hash(FAMObject *self)
2190
+ {
2191
+ Py_hash_t hash = 0;
2192
+ for (Py_ssize_t i = 0; i < self->table_size; i++) {
2193
+ hash = hash * 3 + self->table[i].hash;
2194
+ }
2195
+ if (hash == -1) { // most not return -1
2196
+ return 0;
2197
+ }
2198
+ return hash;
2199
+ }
2200
+
2201
+
2202
+ static PyObject *
2203
+ fam_iter(FAMObject *self)
2204
+ {
2205
+ return fami_new(self, KEYS, false);
2206
+ }
2207
+
2208
+
2209
+ static PyObject *
2210
+ fam_getnewargs(FAMObject *self)
2211
+ {
2212
+ return PyTuple_Pack(1, self->keys);
2213
+ }
2214
+
2215
+
2216
+ static PyObject *
2217
+ fam_reversed(FAMObject *self)
2218
+ {
2219
+ return fami_new(self, KEYS, true);
2220
+ }
2221
+
2222
+
2223
+ static PyObject *
2224
+ fam_sizeof(FAMObject *self)
2225
+ {
2226
+ PyObject *listsizeof = PyObject_CallMethod(self->keys, "__sizeof__", NULL);
2227
+ if (!listsizeof) {
2228
+ return NULL;
2229
+ }
2230
+ Py_ssize_t listbytes = PyLong_AsSsize_t(listsizeof);
2231
+ Py_DECREF(listsizeof);
2232
+ if (listbytes == -1 && PyErr_Occurred()) {
2233
+ return NULL;
2234
+ }
2235
+ return PyLong_FromSsize_t(
2236
+ Py_TYPE(self)->tp_basicsize
2237
+ + listbytes
2238
+ + (self->table_size + SCAN - 1) * sizeof(TableElement)
2239
+ );
2240
+ }
2241
+
2242
+
2243
+ static PyObject *
2244
+ fam_get(FAMObject *self, PyObject *args)
2245
+ {
2246
+ PyObject *key, *missing = Py_None;
2247
+ if (!PyArg_UnpackTuple(args, Py_TYPE(self)->tp_name, 1, 2, &key, &missing))
2248
+ {
2249
+ return NULL;
2250
+ }
2251
+ return get(self, key, missing);
2252
+ }
2253
+
2254
+
2255
+ static PyObject *
2256
+ fam_items(FAMObject *self)
2257
+ {
2258
+ return famv_new(self, ITEMS);
2259
+ }
2260
+
2261
+
2262
+ static PyObject *
2263
+ fam_keys(FAMObject *self)
2264
+ {
2265
+ return famv_new(self, KEYS);
2266
+ }
2267
+
2268
+
2269
+ static PyObject *
2270
+ fam_values(FAMObject *self)
2271
+ {
2272
+ return famv_new(self, VALUES);
2273
+ }
2274
+
2275
+
2276
+ static PyObject *
2277
+ fam_new(PyTypeObject *cls, PyObject *args, PyObject *kwargs)
2278
+ {
2279
+ // NOTE: The original fam_new used to be able to provide a same reference back if a fam was in the args; this is tricky now that we have fam_init
2280
+ FAMObject *self = (FAMObject *)cls->tp_alloc(cls, 0);
2281
+ if (!self) {
2282
+ return NULL;
2283
+ }
2284
+ self->table = NULL;
2285
+ self->keys = NULL;
2286
+ self->key_buffer = NULL;
2287
+ self->keys_size = 0;
2288
+ return (PyObject*)self;
2289
+ }
2290
+
2291
+
2292
+ // This macro can be used with integer and floating point NumPy types, given an `npy_type` and a specialized `insert_func`. Uses context of `fam_init` to get `fam`, `contiguous`, `a`, `keys_size`, and `i`. An optional `post_deref` function can be supplied to transform extracted values before calling the appropriate insert function.
2293
+ # define INSERT_SCALARS(npy_type, insert_func, kat, post_deref) \
2294
+ { \
2295
+ if (contiguous) { \
2296
+ npy_type* b = (npy_type*)PyArray_DATA(a); \
2297
+ npy_type* b_end = b + keys_size; \
2298
+ while (b < b_end) { \
2299
+ if (insert_func(fam, post_deref(*b), i, -1, kat)) { \
2300
+ goto error; \
2301
+ } \
2302
+ b++; \
2303
+ i++; \
2304
+ } \
2305
+ } \
2306
+ else { \
2307
+ for (; i < keys_size; i++) { \
2308
+ if (insert_func(fam, \
2309
+ post_deref(*(npy_type*)PyArray_GETPTR1(a, i)),\
2310
+ i, \
2311
+ -1, \
2312
+ kat)) { \
2313
+ goto error; \
2314
+ } \
2315
+ } \
2316
+ } \
2317
+ } \
2318
+
2319
+ // This macro is for inserting flexible-sized types, Unicode (Py_UCS4) or strings (char). Uses context of `fam_init`.
2320
+ # define INSERT_FLEXIBLE(char_type, insert_func, get_end_func) \
2321
+ { \
2322
+ char_type* p = NULL; \
2323
+ if (contiguous) { \
2324
+ char_type *b = (char_type*)PyArray_DATA(a); \
2325
+ char_type *b_end = b + keys_size * dt_size; \
2326
+ while (b < b_end) { \
2327
+ p = get_end_func(b, dt_size); \
2328
+ if (insert_func(fam, b, p-b, i, -1)) { \
2329
+ goto error; \
2330
+ } \
2331
+ b += dt_size; \
2332
+ i++; \
2333
+ } \
2334
+ } \
2335
+ else { \
2336
+ for (; i < keys_size; i++) { \
2337
+ char_type* v = (char_type*)PyArray_GETPTR1(a, i); \
2338
+ p = get_end_func(v, dt_size); \
2339
+ if (insert_func(fam, v, p-v, i, -1)) { \
2340
+ goto error; \
2341
+ } \
2342
+ } \
2343
+ } \
2344
+ } \
2345
+
2346
+ // Initialize an allocated FAMObject. Returns 0 on success, -1 on error.
2347
+ int
2348
+ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
2349
+ {
2350
+ PyTypeObject* cls = Py_TYPE(self); // borrowed ref
2351
+ const char *name = cls->tp_name;
2352
+ FAMObject* fam = (FAMObject*)self;
2353
+
2354
+ if (kwargs) {
2355
+ PyErr_Format(PyExc_TypeError, "%s takes no keyword arguments", name);
2356
+ return -1;
2357
+ }
2358
+
2359
+ KeysArrayType keys_array_type = KAT_LIST; // default, will override if necessary
2360
+
2361
+ PyObject *keys = NULL;
2362
+ Py_ssize_t keys_size = 0;
2363
+
2364
+ if (!PyArg_UnpackTuple(args, name, 0, 1, &keys)) {
2365
+ return -1;
2366
+ }
2367
+
2368
+ if (!keys) {
2369
+ keys = PyList_New(0);
2370
+ if (!keys) {
2371
+ return -1;
2372
+ }
2373
+ }
2374
+ else if (PyObject_TypeCheck(keys, &FAMType)) {
2375
+ // Use `keys` as old, `self` as new, and fill from old to new. This returns the same error codes as this function.
2376
+ return copy_to_new(cls, (FAMObject*)keys, fam);
2377
+ }
2378
+ else if (PyArray_Check(keys)) {
2379
+ PyArrayObject *a = (PyArrayObject *)keys;
2380
+ if (PyArray_NDIM(a) != 1) {
2381
+ PyErr_SetString(PyExc_TypeError, "Arrays must be 1-dimensional");
2382
+ return -1;
2383
+ }
2384
+
2385
+ int array_t = PyArray_TYPE(a);
2386
+ keys_size = PyArray_SIZE(a);
2387
+
2388
+ if (cls != &AMType &&
2389
+ (PyTypeNum_ISINTEGER(array_t) // signed and unsigned
2390
+ || PyTypeNum_ISFLOAT(array_t)
2391
+ || PyTypeNum_ISFLEXIBLE(array_t)
2392
+ || array_t == NPY_DATETIME ))
2393
+ {
2394
+ if ((PyArray_FLAGS(a) & NPY_ARRAY_WRITEABLE)) {
2395
+ PyErr_Format(PyExc_TypeError, "Arrays must be immutable when given to a %s", name);
2396
+ return -1;
2397
+ }
2398
+ // NOTE: this might return 0 (list) given a dt64 array without a unit
2399
+ keys_array_type = at_to_kat(array_t, a);
2400
+ }
2401
+
2402
+ if (keys_array_type) { // we have a usable array
2403
+ Py_INCREF(keys);
2404
+ }
2405
+ else {
2406
+ // DEBUG_MSG_OBJ("got KAT", PyLong_FromLongLong(keys_array_type));
2407
+ // if an AutoMap or an array that we do not handle, create a list
2408
+ if (array_t == NPY_DATETIME || array_t == NPY_TIMEDELTA){
2409
+ keys = PySequence_List(keys); // force scalars
2410
+ }
2411
+ else {
2412
+ keys = PyArray_ToList(a); // converts to objs
2413
+ }
2414
+ if (!keys) {
2415
+ return -1;
2416
+ }
2417
+ }
2418
+ }
2419
+ else { // assume an arbitrary iterable
2420
+ keys = PySequence_List(keys);
2421
+ if (!keys) {
2422
+ return -1;
2423
+ }
2424
+ keys_size = PyList_GET_SIZE(keys);
2425
+ }
2426
+
2427
+ fam->keys = keys;
2428
+ fam->keys_array_type = keys_array_type;
2429
+ fam->keys_size = keys_size;
2430
+ fam->key_buffer = NULL;
2431
+
2432
+ // NOTE: on itialization, grow_table() does not use keys
2433
+ if (grow_table(fam, keys_size)) {
2434
+ return -1;
2435
+ }
2436
+ Py_ssize_t i = 0;
2437
+ if (keys_array_type) {
2438
+ PyArrayObject *a = (PyArrayObject *)fam->keys;
2439
+ int contiguous = PyArray_IS_C_CONTIGUOUS(a);
2440
+ switch (keys_array_type) {
2441
+ case KAT_INT64:
2442
+ INSERT_SCALARS(npy_int64, insert_int, keys_array_type,);
2443
+ break;
2444
+ case KAT_INT32:
2445
+ INSERT_SCALARS(npy_int32, insert_int, keys_array_type,);
2446
+ break;
2447
+ case KAT_INT16:
2448
+ INSERT_SCALARS(npy_int16, insert_int, keys_array_type,);
2449
+ break;
2450
+ case KAT_INT8:
2451
+ INSERT_SCALARS(npy_int8, insert_int, keys_array_type,);
2452
+ break;
2453
+ case KAT_UINT64:
2454
+ INSERT_SCALARS(npy_uint64, insert_uint, keys_array_type,);
2455
+ break;
2456
+ case KAT_UINT32:
2457
+ INSERT_SCALARS(npy_uint32, insert_uint, keys_array_type,);
2458
+ break;
2459
+ case KAT_UINT16:
2460
+ INSERT_SCALARS(npy_uint16, insert_uint, keys_array_type,);
2461
+ break;
2462
+ case KAT_UINT8:
2463
+ INSERT_SCALARS(npy_uint8, insert_uint, keys_array_type,);
2464
+ break;
2465
+ case KAT_FLOAT64:
2466
+ INSERT_SCALARS(npy_double, insert_double, keys_array_type,);
2467
+ break;
2468
+ case KAT_FLOAT32:
2469
+ INSERT_SCALARS(npy_float, insert_double, keys_array_type,);
2470
+ break;
2471
+ case KAT_FLOAT16:
2472
+ INSERT_SCALARS(npy_half, insert_double, keys_array_type, npy_half_to_double);
2473
+ break;
2474
+ case KAT_UNICODE: {
2475
+ // Over allocate buffer by 1 so there is room for null at end. This buffer is only used in lookup();
2476
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a) / UCS4_SIZE;
2477
+ fam->key_buffer = (Py_UCS4*)PyMem_Malloc((dt_size+1) * UCS4_SIZE);
2478
+ INSERT_FLEXIBLE(Py_UCS4, insert_unicode, ucs4_get_end_p);
2479
+ break;
2480
+ }
2481
+ case KAT_STRING: {
2482
+ Py_ssize_t dt_size = PyArray_ITEMSIZE(a);
2483
+ INSERT_FLEXIBLE(char, insert_string, char_get_end_p);
2484
+ break;
2485
+ }
2486
+ case KAT_DTY:
2487
+ case KAT_DTM:
2488
+ case KAT_DTW:
2489
+ case KAT_DTD:
2490
+ case KAT_DTh:
2491
+ case KAT_DTm:
2492
+ case KAT_DTs:
2493
+ case KAT_DTms:
2494
+ case KAT_DTus:
2495
+ case KAT_DTns:
2496
+ case KAT_DTps:
2497
+ case KAT_DTfs:
2498
+ case KAT_DTas:
2499
+ INSERT_SCALARS(npy_int64, insert_int, KAT_INT64,);
2500
+ break;
2501
+ default:
2502
+ return -1;
2503
+ }
2504
+ }
2505
+ else {
2506
+ for (; i < keys_size; i++) {
2507
+ if (insert_obj(fam, PyList_GET_ITEM(keys, i), i, -1)) {
2508
+ goto error;
2509
+ }
2510
+ }
2511
+ }
2512
+ return 0;
2513
+ error:
2514
+ // assume all dynamic memory assigned to struct attrs that will be cleaned
2515
+ return -1;
2516
+ }
2517
+
2518
+
2519
+ # undef INSERT_SCALARS
2520
+ # undef INSERT_FLEXIBLE
2521
+
2522
+
2523
+ static PyObject *
2524
+ fam_repr(FAMObject *self)
2525
+ {
2526
+ return PyUnicode_FromFormat("%s(%R)", Py_TYPE(self)->tp_name, self->keys);
2527
+ }
2528
+
2529
+
2530
+ static PyObject *
2531
+ fam_richcompare(FAMObject *self, PyObject *other, int op)
2532
+ {
2533
+ if (!PyObject_TypeCheck(other, &FAMType)) {
2534
+ Py_RETURN_NOTIMPLEMENTED;
2535
+ }
2536
+ return PyObject_RichCompare(self->keys, ((FAMObject *)other)->keys, op);
2537
+ }
2538
+
2539
+
2540
+ static PyObject*
2541
+ fam_getstate(FAMObject *self)
2542
+ {
2543
+ PyObject* state = PyTuple_Pack(1, self->keys);
2544
+ return state;
2545
+ }
2546
+
2547
+
2548
+ // State returned here is a tuple of keys, suitable for usage as an `args` argument.
2549
+ static PyObject*
2550
+ fam_setstate(FAMObject *self, PyObject *state)
2551
+ {
2552
+ if (!PyTuple_CheckExact(state) || !PyTuple_GET_SIZE(state)) {
2553
+ PyErr_SetString(PyExc_ValueError, "Unexpected pickled object.");
2554
+ return NULL;
2555
+ }
2556
+ PyObject *keys = PyTuple_GetItem(state, 0);
2557
+ if (PyArray_Check(keys)) {
2558
+ // if we an array, make it immutable
2559
+ PyArray_CLEARFLAGS((PyArrayObject*)keys, NPY_ARRAY_WRITEABLE);
2560
+ }
2561
+ fam_init((PyObject*)self, state, NULL);
2562
+ Py_RETURN_NONE;
2563
+ }
2564
+
2565
+
2566
+ static PyMethodDef fam_methods[] = {
2567
+ {"__getnewargs__", (PyCFunction) fam_getnewargs, METH_NOARGS, NULL},
2568
+ {"__reversed__", (PyCFunction) fam_reversed, METH_NOARGS, NULL},
2569
+ {"__sizeof__", (PyCFunction) fam_sizeof, METH_NOARGS, NULL},
2570
+ {"__getstate__", (PyCFunction) fam_getstate, METH_NOARGS, NULL},
2571
+ {"__setstate__", (PyCFunction) fam_setstate, METH_O, NULL},
2572
+ {"get", (PyCFunction) fam_get, METH_VARARGS, NULL},
2573
+ {"items", (PyCFunction) fam_items, METH_NOARGS, NULL},
2574
+ {"keys", (PyCFunction) fam_keys, METH_NOARGS, NULL},
2575
+ {"values", (PyCFunction) fam_values, METH_NOARGS, NULL},
2576
+ {"get_all", (PyCFunction) fam_get_all, METH_O, NULL},
2577
+ {"get_any", (PyCFunction) fam_get_any, METH_O, NULL},
2578
+ {NULL},
2579
+ };
2580
+
2581
+
2582
+ PyTypeObject FAMType = {
2583
+ PyVarObject_HEAD_INIT(NULL, 0)
2584
+ .tp_as_mapping = &fam_as_mapping,
2585
+ .tp_as_number = &fam_as_number,
2586
+ .tp_as_sequence = &fam_as_sequence,
2587
+ .tp_basicsize = sizeof(FAMObject),
2588
+ .tp_dealloc = (destructor) fam_dealloc,
2589
+ .tp_doc = "An immutable auto-incremented integer-valued mapping.",
2590
+ .tp_hash = (hashfunc) fam_hash,
2591
+ .tp_iter = (getiterfunc) fam_iter,
2592
+ .tp_methods = fam_methods,
2593
+ .tp_name = "arraykit.FrozenAutoMap",
2594
+ .tp_new = fam_new,
2595
+ .tp_init = fam_init,
2596
+ .tp_repr = (reprfunc) fam_repr,
2597
+ .tp_richcompare = (richcmpfunc) fam_richcompare,
2598
+ };
2599
+
2600
+
2601
+ //------------------------------------------------------------------------------
2602
+ // AutoMap subclass
2603
+
2604
+ static PyObject *
2605
+ am_inplace_or(FAMObject *self, PyObject *other)
2606
+ {
2607
+ if (PyObject_TypeCheck(other, &FAMType)) {
2608
+ other = ((FAMObject *)other)->keys;
2609
+ }
2610
+ if (extend(self, other)) {
2611
+ return NULL;
2612
+ }
2613
+ Py_INCREF(self);
2614
+ return (PyObject *)self;
2615
+ }
2616
+
2617
+
2618
+ static PyNumberMethods am_as_number = {
2619
+ .nb_inplace_or = (binaryfunc) am_inplace_or,
2620
+ };
2621
+
2622
+
2623
+ static PyObject *
2624
+ am_add(FAMObject *self, PyObject *other)
2625
+ {
2626
+ if (append(self, other)) {
2627
+ return NULL;
2628
+ }
2629
+ Py_RETURN_NONE;
2630
+ }
2631
+
2632
+
2633
+ static PyObject *
2634
+ am_update(FAMObject *self, PyObject *other)
2635
+ {
2636
+ if (PyObject_TypeCheck(other, &FAMType)) {
2637
+ other = ((FAMObject *)other)->keys;
2638
+ }
2639
+ if (extend(self, other)) {
2640
+ return NULL;
2641
+ }
2642
+ Py_RETURN_NONE;
2643
+ }
2644
+
2645
+
2646
+ static PyMethodDef am_methods[] = {
2647
+ {"add", (PyCFunction) am_add, METH_O, NULL},
2648
+ {"update", (PyCFunction) am_update, METH_O, NULL},
2649
+ {NULL},
2650
+ };
2651
+
2652
+ PyTypeObject AMType = {
2653
+ PyVarObject_HEAD_INIT(NULL, 0)
2654
+ .tp_as_number = &am_as_number,
2655
+ .tp_base = &FAMType,
2656
+ .tp_doc = "A grow-only autoincremented integer-valued mapping.",
2657
+ .tp_methods = am_methods,
2658
+ .tp_name = "arraykit.AutoMap",
2659
+ .tp_richcompare = (richcmpfunc) fam_richcompare,
2660
+ };
2661
+