arraykit 1.2.0__cp314-cp314-win_amd64.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/block_index.c ADDED
@@ -0,0 +1,1445 @@
1
+ # include "Python.h"
2
+ # include "structmember.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
+
10
+ # include "block_index.h"
11
+ # include "utilities.h"
12
+
13
+ PyObject * ErrorInitTypeBlocks;
14
+
15
+ // Returns NULL on error. Returns a new reference. Note that a reference is stolen from the PyObject argument.
16
+ static inline PyObject *
17
+ AK_build_pair_ssize_t_pyo(Py_ssize_t a, PyObject* py_b)
18
+ {
19
+ if (py_b == NULL) { // construction failed
20
+ return NULL;
21
+ }
22
+ PyObject* t = PyTuple_New(2);
23
+ if (t == NULL) {
24
+ return NULL;
25
+ }
26
+ PyObject* py_a = PyLong_FromSsize_t(a);
27
+ if (py_a == NULL) {
28
+ Py_DECREF(t);
29
+ return NULL;
30
+ }
31
+ // steals refs
32
+ PyTuple_SET_ITEM(t, 0, py_a);
33
+ PyTuple_SET_ITEM(t, 1, py_b);
34
+ return t;
35
+ }
36
+
37
+ // Given inclusive start, end indices, returns a new reference to a slice. Returns NULL on error. If `reduce` is True, single width slices return an integer.
38
+ static inline PyObject *
39
+ AK_build_slice_inclusive(Py_ssize_t start, Py_ssize_t end, bool reduce)
40
+ {
41
+ if (reduce && start == end) {
42
+ return PyLong_FromSsize_t(start); // new ref
43
+ }
44
+ // assert(start >= 0);
45
+ if (start <= end) {
46
+ return AK_build_slice(start, end + 1, 1);
47
+ }
48
+ // end of 0 goes to -1, gets converted to None
49
+ return AK_build_slice(start, end - 1, -1);
50
+ }
51
+
52
+ // NOTE: we use platform size types here, which are appropriate for the values, but might pose issues if trying to pass pickles between 32 and 64 bit machines.
53
+ typedef struct BlockIndexRecord {
54
+ Py_ssize_t block; // signed
55
+ Py_ssize_t column;
56
+ } BlockIndexRecord;
57
+
58
+ typedef struct BlockIndexObject {
59
+ PyObject_HEAD
60
+ Py_ssize_t block_count;
61
+ Py_ssize_t row_count;
62
+ Py_ssize_t bir_count;
63
+ Py_ssize_t bir_capacity;
64
+ BlockIndexRecord* bir;
65
+ PyArray_Descr* dtype;
66
+ bool shape_recache;
67
+ PyObject* shape;
68
+ } BlockIndexObject;
69
+
70
+ // Returns a new reference to tuple. Returns NULL on error. Python already wraps negative numbers up to negative length when used in the sequence slot
71
+ static inline PyObject *
72
+ AK_BI_item(BlockIndexObject* self, Py_ssize_t i) {
73
+ if (!((size_t)i < (size_t)self->bir_count)) {
74
+ PyErr_SetString(PyExc_IndexError, "index out of range");
75
+ return NULL;
76
+ }
77
+ BlockIndexRecord* biri = &self->bir[i];
78
+ return AK_build_pair_ssize_t(biri->block, biri->column); // may be NULL
79
+ }
80
+
81
+ //------------------------------------------------------------------------------
82
+ // BI Iterator
83
+
84
+ typedef struct BIIterObject {
85
+ PyObject_HEAD
86
+ BlockIndexObject *bi;
87
+ bool reversed;
88
+ Py_ssize_t pos; // current index state, mutated in-place
89
+ } BIIterObject;
90
+
91
+ static inline PyObject *
92
+ BIIter_new(BlockIndexObject *bi, bool reversed) {
93
+ BIIterObject *bii = PyObject_New(BIIterObject, &BIIterType);
94
+ if (!bii) {
95
+ return NULL;
96
+ }
97
+ Py_INCREF((PyObject*)bi);
98
+ bii->bi = bi;
99
+ bii->reversed = reversed;
100
+ bii->pos = 0;
101
+ return (PyObject *)bii;
102
+ }
103
+
104
+ void
105
+ BIIter_dealloc(BIIterObject *self) {
106
+ Py_DECREF((PyObject*)self->bi);
107
+ PyObject_Del((PyObject*)self);
108
+ }
109
+
110
+ PyObject *
111
+ BIIter_iter(BIIterObject *self) {
112
+ Py_INCREF(self);
113
+ return (PyObject*)self;
114
+ }
115
+
116
+ PyObject *
117
+ BIIter_iternext(BIIterObject *self) {
118
+ Py_ssize_t i;
119
+ if (self->reversed) {
120
+ i = self->bi->bir_count - ++self->pos;
121
+ if (i < 0) {
122
+ return NULL;
123
+ }
124
+ }
125
+ else {
126
+ i = self->pos++;
127
+ }
128
+ if (self->bi->bir_count <= i) {
129
+ return NULL;
130
+ }
131
+ return AK_BI_item(self->bi, i); // return new ref
132
+ }
133
+
134
+ PyObject *
135
+ BIIter_reversed(BIIterObject *self) {
136
+ return BIIter_new(self->bi, !self->reversed);
137
+ }
138
+
139
+ PyObject *
140
+ BIIter_length_hint(BIIterObject *self) {
141
+ // this works for reversed as we use self->pos to subtract from length
142
+ Py_ssize_t len = Py_MAX(0, self->bi->bir_count - self->pos);
143
+ return PyLong_FromSsize_t(len);
144
+ }
145
+
146
+ static PyMethodDef BIIter_methods[] = {
147
+ {"__length_hint__", (PyCFunction)BIIter_length_hint, METH_NOARGS, NULL},
148
+ {"__reversed__", (PyCFunction)BIIter_reversed, METH_NOARGS, NULL},
149
+ {NULL},
150
+ };
151
+
152
+ PyTypeObject BIIterType = {
153
+ PyVarObject_HEAD_INIT(NULL, 0)
154
+ .tp_basicsize = sizeof(BIIterObject),
155
+ .tp_dealloc = (destructor) BIIter_dealloc,
156
+ .tp_iter = (getiterfunc) BIIter_iter,
157
+ .tp_iternext = (iternextfunc) BIIter_iternext,
158
+ .tp_methods = BIIter_methods,
159
+ .tp_name = "arraykit.BlockIndexIterator",
160
+ };
161
+
162
+ //------------------------------------------------------------------------------
163
+ // BI Iterator sequence selection
164
+
165
+ typedef enum BIIterSelectorKind {
166
+ BIIS_SEQ, // BIIterSeqType
167
+ BIIS_SLICE,
168
+ BIIS_BOOLEAN, // BIIterBoolType
169
+ BIIS_UNKNOWN
170
+ } BIIterSelectorKind;
171
+
172
+ // Forward-def
173
+ static inline PyObject *
174
+ BIIterSelector_new(BlockIndexObject *bi,
175
+ PyObject* selector,
176
+ bool reversed,
177
+ BIIterSelectorKind kind,
178
+ bool ascending
179
+ );
180
+
181
+ typedef struct BIIterSeqObject {
182
+ PyObject_HEAD
183
+ BlockIndexObject *bi;
184
+ bool reversed;
185
+ PyObject* selector;
186
+ Py_ssize_t pos; // current pos in sequence, mutated in-place
187
+ Py_ssize_t len;
188
+ bool is_array;
189
+ } BIIterSeqObject;
190
+
191
+ void
192
+ BIIterSeq_dealloc(BIIterSeqObject *self) {
193
+ Py_DECREF((PyObject*)self->bi);
194
+ Py_DECREF(self->selector);
195
+ PyObject_Del((PyObject*)self);
196
+ }
197
+
198
+ PyObject *
199
+ BIIterSeq_iter(BIIterSeqObject *self) {
200
+ Py_INCREF(self);
201
+ return (PyObject*)self;
202
+ }
203
+
204
+ // Returns -1 on end of sequence; return -1 with exception set on
205
+ static inline Py_ssize_t
206
+ BIIterSeq_iternext_index(BIIterSeqObject *self)
207
+ {
208
+ Py_ssize_t i;
209
+ if (self->reversed) {
210
+ i = self->len - ++self->pos;
211
+ if (i < 0) {
212
+ return -1;
213
+ }
214
+ }
215
+ else {
216
+ i = self->pos++;
217
+ }
218
+ if (self->len <= i) {
219
+ return -1;
220
+ }
221
+ // use i to get index from selector
222
+ Py_ssize_t t = 0;
223
+ if (self->is_array) {
224
+ PyArrayObject *a = (PyArrayObject *)self->selector;
225
+ switch (PyArray_TYPE(a)) { // type of passed in array
226
+ case NPY_INT64:
227
+ t = (Py_ssize_t)*(npy_int64*)PyArray_GETPTR1(a, i);
228
+ break;
229
+ case NPY_INT32:
230
+ t = *(npy_int32*)PyArray_GETPTR1(a, i);
231
+ break;
232
+ case NPY_INT16:
233
+ t = *(npy_int16*)PyArray_GETPTR1(a, i);
234
+ break;
235
+ case NPY_INT8:
236
+ t = *(npy_int8*)PyArray_GETPTR1(a, i);
237
+ break;
238
+ case NPY_UINT64:
239
+ t = (Py_ssize_t)*(npy_uint64*)PyArray_GETPTR1(a, i);
240
+ break;
241
+ case NPY_UINT32:
242
+ t = *(npy_uint32*)PyArray_GETPTR1(a, i);
243
+ break;
244
+ case NPY_UINT16:
245
+ t = *(npy_uint16*)PyArray_GETPTR1(a, i);
246
+ break;
247
+ case NPY_UINT8:
248
+ t = *(npy_uint8*)PyArray_GETPTR1(a, i);
249
+ break;
250
+ }
251
+ }
252
+ else { // is a list
253
+ PyObject* o = PyList_GET_ITEM(self->selector, i); // borrow
254
+ if (PyNumber_Check(o)) { // handles scalars
255
+ t = PyNumber_AsSsize_t(o, NULL);
256
+ }
257
+ else {
258
+ PyErr_SetString(PyExc_TypeError, "element type not suitable for indexing");
259
+ return -1;
260
+ }
261
+ }
262
+ if (t < 0) {
263
+ t = self->bi->bir_count + t;
264
+ }
265
+ // we have to ensure valid range here to set an index error and distinguish from end of iteration
266
+ if (!((size_t)t < (size_t)self->bi->bir_count)) {
267
+ PyErr_SetString(PyExc_IndexError, "index out of range");
268
+ return -1;
269
+ }
270
+ return t;
271
+ }
272
+
273
+ PyObject *
274
+ BIIterSeq_iternext(BIIterSeqObject *self)
275
+ {
276
+ Py_ssize_t i = BIIterSeq_iternext_index(self);
277
+ if (i == -1) {
278
+ return NULL; // an error is set
279
+ }
280
+ return AK_BI_item(self->bi, i); // return new ref
281
+ }
282
+
283
+ PyObject *
284
+ BIIterSeq_reversed(BIIterSeqObject *self)
285
+ {
286
+ return BIIterSelector_new(self->bi, self->selector, !self->reversed, BIIS_SEQ, false);
287
+ }
288
+
289
+ PyObject *
290
+ BIIterSeq_length_hint(BIIterSeqObject *self)
291
+ {
292
+ // this works for reversed as we use self-> index to subtract from length
293
+ Py_ssize_t len = Py_MAX(0, self->len - self->pos);
294
+ return PyLong_FromSsize_t(len);
295
+ }
296
+
297
+ static PyMethodDef BIiterSeq_methods[] = {
298
+ {"__length_hint__", (PyCFunction)BIIterSeq_length_hint, METH_NOARGS, NULL},
299
+ {"__reversed__", (PyCFunction)BIIterSeq_reversed, METH_NOARGS, NULL},
300
+ {NULL},
301
+ };
302
+
303
+ PyTypeObject BIIterSeqType = {
304
+ PyVarObject_HEAD_INIT(NULL, 0)
305
+ .tp_basicsize = sizeof(BIIterSeqObject),
306
+ .tp_dealloc = (destructor) BIIterSeq_dealloc,
307
+ .tp_iter = (getiterfunc) BIIterSeq_iter,
308
+ .tp_iternext = (iternextfunc) BIIterSeq_iternext,
309
+ .tp_methods = BIiterSeq_methods,
310
+ .tp_name = "arraykit.BlockIndexIteratorSequence",
311
+ };
312
+
313
+ //------------------------------------------------------------------------------
314
+ // BI Iterator slice selection
315
+
316
+ typedef struct BIIterSliceObject {
317
+ PyObject_HEAD
318
+ BlockIndexObject *bi;
319
+ bool reversed;
320
+ PyObject* selector; // slice
321
+ Py_ssize_t count; // count of , mutated in-place
322
+ // these are the normalized values truncated to the span of the bir_count; len is the realized length after extraction; step is always set to 1 if missing; len is 0 if no realized values
323
+ Py_ssize_t pos;
324
+ Py_ssize_t step;
325
+ Py_ssize_t len;
326
+ } BIIterSliceObject;
327
+
328
+ void
329
+ BIIterSlice_dealloc(BIIterSliceObject *self) {
330
+ // AK_DEBUG_MSG_REFCNT("start dealloc", self);
331
+ Py_DECREF((PyObject*)self->bi);
332
+ Py_DECREF(self->selector);
333
+ PyObject_Del((PyObject*)self);
334
+ }
335
+
336
+ PyObject *
337
+ BIIterSlice_iter(BIIterSliceObject *self) {
338
+ Py_INCREF(self);
339
+ return (PyObject*)self;
340
+ }
341
+
342
+ // NOTE: this does not use `reversed`, as pos, step, and count are set in BIIterSelector_new
343
+ static inline Py_ssize_t
344
+ BIIterSlice_iternext_index(BIIterSliceObject *self)
345
+ {
346
+ if (self->len == 0 || self->count >= self->len) {
347
+ return -1;
348
+ }
349
+ Py_ssize_t i = self->pos;
350
+ self->pos += self->step;
351
+ self->count++; // by counting index we we do not need to compare to stop
352
+ // i will never be out of range
353
+ return i;
354
+ }
355
+
356
+ PyObject *
357
+ BIIterSlice_iternext(BIIterSliceObject *self) {
358
+ Py_ssize_t i = BIIterSlice_iternext_index(self);
359
+ if (i == -1) {
360
+ return NULL;
361
+ }
362
+ return AK_BI_item(self->bi, i); // return new ref
363
+ }
364
+
365
+ PyObject *
366
+ BIIterSlice_reversed(BIIterSliceObject *self)
367
+ {
368
+ return BIIterSelector_new(self->bi, self->selector, !self->reversed, BIIS_SLICE, false);
369
+ }
370
+
371
+ PyObject *
372
+ BIIterSlice_length_hint(BIIterSliceObject *self)
373
+ {
374
+ // this works for reversed as we use self-> index to subtract from length
375
+ Py_ssize_t len = Py_MAX(0, self->len - self->count);
376
+ return PyLong_FromSsize_t(len);
377
+ }
378
+
379
+ static PyMethodDef BIiterSlice_methods[] = {
380
+ {"__length_hint__", (PyCFunction)BIIterSlice_length_hint, METH_NOARGS, NULL},
381
+ {"__reversed__", (PyCFunction)BIIterSlice_reversed, METH_NOARGS, NULL},
382
+ {NULL},
383
+ };
384
+
385
+ PyTypeObject BIIterSliceType = {
386
+ PyVarObject_HEAD_INIT(NULL, 0)
387
+ .tp_basicsize = sizeof(BIIterSliceObject),
388
+ .tp_dealloc = (destructor) BIIterSlice_dealloc,
389
+ .tp_iter = (getiterfunc) BIIterSlice_iter,
390
+ .tp_iternext = (iternextfunc) BIIterSlice_iternext,
391
+ .tp_methods = BIiterSlice_methods,
392
+ .tp_name = "arraykit.BlockIndexIteratorSlice",
393
+ };
394
+
395
+ //------------------------------------------------------------------------------
396
+ // BI Iterator Boolean array selection
397
+
398
+ typedef struct BIIterBooleanObject {
399
+ PyObject_HEAD
400
+ BlockIndexObject *bi;
401
+ bool reversed;
402
+ PyObject* selector;
403
+ Py_ssize_t pos; // current index, mutated in-place
404
+ Py_ssize_t len;
405
+ } BIIterBooleanObject;
406
+
407
+ void
408
+ BIIterBoolean_dealloc(BIIterBooleanObject *self) {
409
+ Py_DECREF((PyObject*)self->bi);
410
+ Py_DECREF(self->selector);
411
+ PyObject_Del((PyObject*)self);
412
+ }
413
+
414
+ PyObject *
415
+ BIIterBoolean_iter(BIIterBooleanObject *self)
416
+ {
417
+ Py_INCREF(self);
418
+ return (PyObject*)self;
419
+ }
420
+
421
+ static inline Py_ssize_t
422
+ BIIterBoolean_iternext_index(BIIterBooleanObject *self)
423
+ {
424
+ npy_bool v = 0;
425
+ Py_ssize_t i = -1;
426
+ PyArrayObject* a = (PyArrayObject*) self->selector;
427
+
428
+ if (!self->reversed) {
429
+ while (self->pos < self->len) {
430
+ v = *(npy_bool*)PyArray_GETPTR1(a, self->pos);
431
+ if (v) {
432
+ i = self->pos;
433
+ self->pos++;
434
+ break;
435
+ }
436
+ self->pos++;
437
+ }
438
+ }
439
+ else { // reversed
440
+ while (self->pos >= 0) {
441
+ v = *(npy_bool*)PyArray_GETPTR1(a, self->pos);
442
+ if (v) {
443
+ i = self->pos;
444
+ self->pos--;
445
+ break;
446
+ }
447
+ self->pos--;
448
+ }
449
+ }
450
+ if (i != -1) {
451
+ return i;
452
+ }
453
+ return -1; // no True remain
454
+ }
455
+
456
+ PyObject *
457
+ BIIterBoolean_iternext(BIIterBooleanObject *self) {
458
+ Py_ssize_t i = BIIterBoolean_iternext_index(self);
459
+ if (i == -1) {
460
+ return NULL;
461
+ }
462
+ return AK_BI_item(self->bi, i); // return new ref
463
+ }
464
+
465
+ PyObject *
466
+ BIIterBoolean_reversed(BIIterBooleanObject *self)
467
+ {
468
+ return BIIterSelector_new(self->bi, self->selector, !self->reversed, BIIS_BOOLEAN, 0);
469
+ }
470
+
471
+ // NOTE: no length hint given as we would have to traverse whole array and count True... not sure it is worht it.
472
+ static PyMethodDef BIiterBoolean_methods[] = {
473
+ {"__reversed__", (PyCFunction)BIIterBoolean_reversed, METH_NOARGS, NULL},
474
+ {NULL},
475
+ };
476
+
477
+ PyTypeObject BIIterBoolType = {
478
+ PyVarObject_HEAD_INIT(NULL, 0)
479
+ .tp_basicsize = sizeof(BIIterBooleanObject),
480
+ .tp_dealloc = (destructor) BIIterBoolean_dealloc,
481
+ .tp_iter = (getiterfunc) BIIterBoolean_iter,
482
+ .tp_iternext = (iternextfunc) BIIterBoolean_iternext,
483
+ .tp_methods = BIiterBoolean_methods,
484
+ .tp_name = "arraykit.BlockIndexIteratorBoolean",
485
+ };
486
+
487
+ //------------------------------------------------------------------------------
488
+ // BI Iterator Contigous
489
+
490
+ typedef struct BIIterContiguousObject {
491
+ PyObject_HEAD
492
+ BlockIndexObject *bi;
493
+ PyObject* iter; // own reference to core iterator
494
+ bool reversed;
495
+ Py_ssize_t last_block;
496
+ Py_ssize_t last_column;
497
+ Py_ssize_t next_block;
498
+ Py_ssize_t next_column;
499
+ bool reduce; // optionally reduce slices to integers
500
+ } BIIterContiguousObject;
501
+
502
+ // Create a new contiguous slice iterator. Return NULL on error. Steals a reference to PyObject* iter.
503
+ static inline PyObject *
504
+ BIIterContiguous_new(BlockIndexObject *bi,
505
+ bool reversed,
506
+ PyObject* iter,
507
+ bool reduce)
508
+ {
509
+ BIIterContiguousObject *bii = PyObject_New(BIIterContiguousObject, &BIIterContiguousType);
510
+ if (!bii) {
511
+ return NULL;
512
+ }
513
+ Py_INCREF((PyObject*)bi);
514
+ bii->bi = bi;
515
+
516
+ bii->iter = iter; // steals ref
517
+ bii->reversed = reversed;
518
+
519
+ bii->last_block = -1;
520
+ bii->last_column = -1;
521
+ bii->next_block = -1;
522
+ bii->next_column = -1;
523
+ bii->reduce = reduce;
524
+
525
+ return (PyObject *)bii;
526
+ }
527
+
528
+ void
529
+ BIIterContiguous_dealloc(BIIterContiguousObject *self)
530
+ {
531
+ Py_DECREF((PyObject*)self->bi);
532
+ Py_DECREF(self->iter);
533
+ PyObject_Del((PyObject*)self);
534
+ }
535
+
536
+ // Simply incref this object and return it.
537
+ PyObject *
538
+ BIIterContiguous_iter(BIIterContiguousObject *self)
539
+ {
540
+ Py_INCREF(self);
541
+ return (PyObject*)self;
542
+ }
543
+
544
+ // Returns a new reference.
545
+ PyObject *
546
+ BIIterContiguous_reversed(BIIterContiguousObject *self)
547
+ {
548
+ bool reversed = !self->reversed;
549
+
550
+ PyObject* selector = NULL;
551
+ PyTypeObject* type = Py_TYPE(self->iter);
552
+ if (type == &BIIterSeqType) {
553
+ selector = ((BIIterSeqObject*)self->iter)->selector;
554
+ }
555
+ else if (type == &BIIterSliceType) {
556
+ selector = ((BIIterSliceObject*)self->iter)->selector;
557
+ }
558
+ else if (type == &BIIterBoolType) {
559
+ selector = ((BIIterBooleanObject*)self->iter)->selector;
560
+ }
561
+ if (selector == NULL) {
562
+ return NULL;
563
+ }
564
+
565
+ PyObject* iter = BIIterSelector_new(self->bi,
566
+ selector,
567
+ reversed,
568
+ BIIS_UNKNOWN, // let type be determined by selector
569
+ 0);
570
+ if (iter == NULL) {
571
+ return NULL;
572
+ }
573
+ PyObject* biiter = BIIterContiguous_new(self->bi,
574
+ reversed,
575
+ iter, // steals ref
576
+ self->reduce);
577
+ return biiter;
578
+ }
579
+
580
+ PyObject *
581
+ BIIterContiguous_iternext(BIIterContiguousObject *self)
582
+ {
583
+ Py_ssize_t i = -1;
584
+ PyObject* iter = self->iter;
585
+ PyTypeObject* type = Py_TYPE(iter);
586
+
587
+ Py_ssize_t slice_start = -1;
588
+ Py_ssize_t block;
589
+ Py_ssize_t column;
590
+
591
+ while (1) {
592
+ if (self->next_block == -2) {
593
+ break; // terminate
594
+ }
595
+ if (self->next_block != -1) {
596
+ // discontinuity found on last iteration, set new start
597
+ self->last_block = self->next_block;
598
+ self->last_column = slice_start = self->next_column;
599
+ self->next_block = self->next_column = -1; // clear next state
600
+ }
601
+ if (type == &BIIterSeqType) {
602
+ i = BIIterSeq_iternext_index((BIIterSeqObject*)iter);
603
+ }
604
+ else if (type == &BIIterSliceType) {
605
+ i = BIIterSlice_iternext_index((BIIterSliceObject*)iter);
606
+ }
607
+ else if (type == &BIIterBoolType) {
608
+ i = BIIterBoolean_iternext_index((BIIterBooleanObject*)iter);
609
+ }
610
+ if (i == -1) { // end of iteration or error
611
+ if (PyErr_Occurred()) {
612
+ break;
613
+ }
614
+ // no more pairs, return previous slice_start, flag for end on next call
615
+ self->next_block = -2;
616
+ if (self->last_block == -1) { // iter produced no values, terminate
617
+ break;
618
+ }
619
+ return AK_build_pair_ssize_t_pyo( // steals ref
620
+ self->last_block,
621
+ AK_build_slice_inclusive(slice_start,
622
+ self->last_column,
623
+ self->reduce));
624
+ }
625
+ // i is gauranteed to be within the range of self->bit_count at this point; the only source of arbitrary indices is in BIIterSeq_iternext_index, and that function validates the range
626
+ BlockIndexRecord* biri = &self->bi->bir[i];
627
+ block = biri->block;
628
+ column = biri->column;
629
+
630
+ // inititialization
631
+ if (self->last_block == -1) {
632
+ self->last_block = block;
633
+ self->last_column = column;
634
+ slice_start = column;
635
+ continue;
636
+ }
637
+ if (self->last_block == block && llabs(column - self->last_column) == 1) {
638
+ // contiguious region found, can be postive or negative
639
+ self->last_column = column;
640
+ continue;
641
+ }
642
+ self->next_block = block;
643
+ self->next_column = column;
644
+ return AK_build_pair_ssize_t_pyo( // steals ref
645
+ self->last_block,
646
+ AK_build_slice_inclusive(slice_start,
647
+ self->last_column,
648
+ self->reduce));
649
+ }
650
+ return NULL;
651
+ }
652
+
653
+ // not implementing __length_hint__
654
+ static PyMethodDef BIIterContiguous_methods[] = {
655
+ {"__reversed__", (PyCFunction)BIIterContiguous_reversed, METH_NOARGS, NULL},
656
+ {NULL},
657
+ };
658
+
659
+ PyTypeObject BIIterContiguousType = {
660
+ PyVarObject_HEAD_INIT(NULL, 0)
661
+ .tp_basicsize = sizeof(BIIterContiguousObject),
662
+ .tp_dealloc = (destructor) BIIterContiguous_dealloc,
663
+ .tp_iter = (getiterfunc) BIIterContiguous_iter,
664
+ .tp_iternext = (iternextfunc) BIIterContiguous_iternext,
665
+ .tp_methods = BIIterContiguous_methods,
666
+ .tp_name = "arraykit.BlockIndexContiguousIterator",
667
+ };
668
+
669
+ //------------------------------------------------------------------------------
670
+ // BI Iterator Block Slice
671
+
672
+ typedef struct BIIterBlockObject {
673
+ PyObject_HEAD
674
+ BlockIndexObject *bi;
675
+ bool reversed;
676
+ Py_ssize_t pos; // current index state, mutated in-place
677
+ PyObject* null_slice;
678
+ } BIIterBlockObject;
679
+
680
+ static inline PyObject *
681
+ BIIterBlock_new(BlockIndexObject *bi, bool reversed) {
682
+ BIIterBlockObject *bii = PyObject_New(BIIterBlockObject, &BIIterBlockType);
683
+ if (!bii) {
684
+ return NULL;
685
+ }
686
+ Py_INCREF((PyObject*)bi);
687
+ bii->bi = bi;
688
+ bii->reversed = reversed;
689
+ bii->pos = 0;
690
+
691
+ // create a new ref of the null slice
692
+ PyObject* ns = AK_build_slice(-1, -1, 1); // get all null; new ref
693
+ if (ns == NULL) {
694
+ return NULL;
695
+ }
696
+ bii->null_slice = ns;
697
+ return (PyObject *)bii;
698
+ }
699
+
700
+ void
701
+ BIIterBlock_dealloc(BIIterBlockObject *self) {
702
+ Py_DECREF((PyObject*)self->bi);
703
+ Py_DECREF(self->null_slice);
704
+ PyObject_Del((PyObject*)self);
705
+ }
706
+
707
+ PyObject *
708
+ BIIterBlock_iter(BIIterBlockObject *self) {
709
+ Py_INCREF(self);
710
+ return (PyObject*)self;
711
+ }
712
+
713
+ PyObject *
714
+ BIIterBlock_iternext(BIIterBlockObject *self) {
715
+ Py_ssize_t i;
716
+ if (self->reversed) {
717
+ i = self->bi->block_count - ++self->pos;
718
+ if (i < 0) {
719
+ return NULL;
720
+ }
721
+ }
722
+ else {
723
+ i = self->pos++;
724
+ }
725
+ if (self->bi->block_count <= i) {
726
+ return NULL;
727
+ }
728
+ // AK_build_pair_ssize_t_pyo steals the reference to the object; so incref here
729
+ Py_INCREF(self->null_slice);
730
+ PyObject* t = AK_build_pair_ssize_t_pyo(i, self->null_slice); // return new ref
731
+ if (t == NULL) {
732
+ // if tuple creation failed need to undo incref
733
+ Py_DECREF(self->null_slice);
734
+ }
735
+ return t;
736
+ }
737
+
738
+ PyObject *
739
+ BIIterBlock_reversed(BIIterBlockObject *self) {
740
+ return BIIterBlock_new(self->bi, !self->reversed);
741
+ }
742
+
743
+ PyObject *
744
+ BIIterBlock_length_hint(BIIterBlockObject *self) {
745
+ // this works for reversed as we use self->pos to subtract from length
746
+ Py_ssize_t len = Py_MAX(0, self->bi->block_count - self->pos);
747
+ return PyLong_FromSsize_t(len);
748
+ }
749
+
750
+ static PyMethodDef BIIterBlock_methods[] = {
751
+ {"__length_hint__", (PyCFunction)BIIterBlock_length_hint, METH_NOARGS, NULL},
752
+ {"__reversed__", (PyCFunction)BIIterBlock_reversed, METH_NOARGS, NULL},
753
+ {NULL},
754
+ };
755
+
756
+ PyTypeObject BIIterBlockType = {
757
+ PyVarObject_HEAD_INIT(NULL, 0)
758
+ .tp_basicsize = sizeof(BIIterBlockObject),
759
+ .tp_dealloc = (destructor) BIIterBlock_dealloc,
760
+ .tp_iter = (getiterfunc) BIIterBlock_iter,
761
+ .tp_iternext = (iternextfunc) BIIterBlock_iternext,
762
+ .tp_methods = BIIterBlock_methods,
763
+ .tp_name = "arraykit.BlockIndexBlockIterator",
764
+ };
765
+
766
+ //------------------------------------------------------------------------------
767
+
768
+ // NOTE: this constructor returns one of three different PyObject types. We do this to consolidate error reporting and type checks.
769
+ // The ascending argument is applied before consideration of a reverse iterator
770
+ PyObject *
771
+ BIIterSelector_new(BlockIndexObject *bi,
772
+ PyObject* selector,
773
+ bool reversed,
774
+ BIIterSelectorKind kind,
775
+ bool ascending) {
776
+
777
+ bool is_array = false;
778
+ bool incref_selector = true; // incref borrowed selector; but if a new ref is made, do not
779
+
780
+ Py_ssize_t len = -1;
781
+ Py_ssize_t pos = 0;
782
+ Py_ssize_t stop = 0;
783
+ Py_ssize_t step = 0;
784
+
785
+ if (PyArray_Check(selector)) {
786
+ if (kind == BIIS_SLICE) {
787
+ PyErr_SetString(PyExc_TypeError, "Arrays cannot be used as selectors for slice iterators");
788
+ return NULL;
789
+ }
790
+ is_array = true;
791
+ PyArrayObject *a = (PyArrayObject *)selector;
792
+ if (PyArray_NDIM(a) != 1) {
793
+ PyErr_SetString(PyExc_TypeError, "Arrays must be 1-dimensional");
794
+ return NULL;
795
+ }
796
+ len = PyArray_SIZE(a);
797
+
798
+ char k = PyArray_DESCR(a)->kind;
799
+ if (kind == BIIS_UNKNOWN) {
800
+ if (k == 'i' || k == 'u') {
801
+ kind = BIIS_SEQ;
802
+ }
803
+ else if (k == 'b') {
804
+ kind = BIIS_BOOLEAN;
805
+ }
806
+ else {
807
+ PyErr_SetString(PyExc_TypeError, "Arrays kind not supported");
808
+ return NULL;
809
+ }
810
+ }
811
+ else if (kind == BIIS_SEQ && k != 'i' && k != 'u') {
812
+ PyErr_SetString(PyExc_TypeError, "Arrays must be integer kind");
813
+ return NULL;
814
+ }
815
+ else if (kind == BIIS_BOOLEAN && k != 'b') {
816
+ PyErr_SetString(PyExc_TypeError, "Arrays must be Boolean kind");
817
+ return NULL;
818
+ }
819
+
820
+ if (kind == BIIS_BOOLEAN) {
821
+ if (len != bi->bir_count) {
822
+ PyErr_SetString(PyExc_TypeError, "Boolean arrays must match BlockIndex size");
823
+ return NULL;
824
+ }
825
+ }
826
+ else if (ascending) { // not Boolean
827
+ // NOTE: we can overwrite selector here as we have a borrowed refernce; sorting gives us a new reference, so we do not need to incref below
828
+ selector = PyArray_NewCopy(a, NPY_CORDER);
829
+ // sort in-place; can use a non-stable sort
830
+ if (PyArray_Sort((PyArrayObject*)selector, 0, NPY_QUICKSORT)) {
831
+ return NULL; // returns -1 on error
832
+ }; // new ref
833
+ incref_selector = false;
834
+ }
835
+ }
836
+ else if (PySlice_Check(selector)) {
837
+ if (kind == BIIS_UNKNOWN) {
838
+ kind = BIIS_SLICE;
839
+ }
840
+ else if (kind != BIIS_SLICE) {
841
+ PyErr_SetString(PyExc_TypeError, "Slices cannot be used as selectors for this type of iterator");
842
+ return NULL;
843
+ }
844
+
845
+ if (ascending) {
846
+ // NOTE: we are abandoning the borrowed reference
847
+ selector = AK_slice_to_ascending_slice(selector, bi->bir_count); // new ref
848
+ incref_selector = false;
849
+ }
850
+ if (PySlice_Unpack(selector, &pos, &stop, &step)) {
851
+ return NULL;
852
+ }
853
+ len = PySlice_AdjustIndices(bi->bir_count, &pos, &stop, step);
854
+
855
+ if (reversed) {
856
+ pos += (step * (len - 1));
857
+ step *= -1;
858
+ }
859
+ }
860
+ else if (PyList_CheckExact(selector)) {
861
+ if (kind == BIIS_UNKNOWN) {
862
+ kind = BIIS_SEQ;
863
+ }
864
+ else if (kind != BIIS_SEQ) {
865
+ PyErr_SetString(PyExc_TypeError, "Lists cannot be used as for non-sequence iterators");
866
+ return NULL;
867
+ }
868
+ len = PyList_GET_SIZE(selector);
869
+
870
+ if (ascending) {
871
+ // abandoning borrowed ref
872
+ selector = PyObject_CallMethod(selector, "copy", NULL); // new ref
873
+ if (selector == NULL) {
874
+ return NULL;
875
+ }
876
+ PyObject* post = PyObject_CallMethod(selector, "sort", NULL); // new ref
877
+ if (post == NULL) {
878
+ return NULL;
879
+ }
880
+ Py_DECREF(post); // just a None
881
+ incref_selector = false;
882
+ }
883
+ }
884
+ else {
885
+ PyErr_SetString(PyExc_TypeError, "Input type not supported");
886
+ return NULL;
887
+ }
888
+
889
+ PyObject *bii = NULL;
890
+ switch (kind) {
891
+ case BIIS_SEQ: {
892
+ BIIterSeqObject* it = PyObject_New(BIIterSeqObject, &BIIterSeqType);
893
+ if (it == NULL) {goto error;}
894
+ it->bi = bi;
895
+ it->selector = selector;
896
+ it->reversed = reversed;
897
+ it->len = len;
898
+ it->pos = 0;
899
+ it->is_array = is_array;
900
+ bii = (PyObject*)it;
901
+ break;
902
+ }
903
+ case BIIS_SLICE: {
904
+ BIIterSliceObject* it = PyObject_New(BIIterSliceObject, &BIIterSliceType);
905
+ if (it == NULL) {goto error;}
906
+ it->bi = bi;
907
+ it->selector = selector;
908
+ it->reversed = reversed;
909
+ it->len = len;
910
+ it->pos = pos;
911
+ it->step = step;
912
+ it->count = 0;
913
+ bii = (PyObject*)it;
914
+ break;
915
+ }
916
+ case BIIS_BOOLEAN: {
917
+ BIIterBooleanObject* it = PyObject_New(BIIterBooleanObject, &BIIterBoolType);
918
+ if (it == NULL) {goto error;}
919
+ it->bi = bi;
920
+ it->selector = selector;
921
+ it->reversed = reversed;
922
+ it->len = len;
923
+ it->pos = reversed ? len - 1 : 0;
924
+ bii = (PyObject*)it;
925
+ break;
926
+ }
927
+ case BIIS_UNKNOWN:
928
+ goto error; // should not get here!
929
+ }
930
+ Py_INCREF((PyObject*)bi);
931
+
932
+ if (incref_selector) {
933
+ Py_INCREF(selector);
934
+ }
935
+ return bii;
936
+ error: // nothing shold be increfed when we get here
937
+ return NULL;
938
+ }
939
+
940
+ //------------------------------------------------------------------------------
941
+ // block index new, init, memory
942
+
943
+ // Returns 0 on succes, -1 on error.
944
+ static inline int
945
+ AK_BI_BIR_new(BlockIndexObject* bi) {
946
+ BlockIndexRecord* bir = (BlockIndexRecord*)PyMem_Malloc(
947
+ sizeof(BlockIndexRecord) * bi->bir_capacity);
948
+ if (bir == NULL) {
949
+ PyErr_SetNone(PyExc_MemoryError);
950
+ return -1;
951
+ }
952
+ bi->bir = bir;
953
+ return 0;
954
+ }
955
+
956
+ // Returns 0 on success, -1 on error
957
+ static inline int
958
+ AK_BI_BIR_resize(BlockIndexObject* bi, Py_ssize_t increment) {
959
+ Py_ssize_t target = bi->bir_count + increment;
960
+ Py_ssize_t capacity = bi->bir_capacity;
961
+ if (AK_UNLIKELY(target >= capacity)) {
962
+ while (capacity < target) {
963
+ capacity <<= 1; // get 2x the capacity
964
+ }
965
+ bi->bir = PyMem_Realloc(bi->bir,
966
+ sizeof(BlockIndexRecord) * capacity);
967
+ if (bi->bir == NULL) {
968
+ PyErr_SetNone(PyExc_MemoryError);
969
+ return -1;
970
+ }
971
+ bi->bir_capacity = capacity;
972
+ }
973
+ return 0;
974
+ }
975
+
976
+ PyDoc_STRVAR(
977
+ BlockIndex_doc,
978
+ "\n"
979
+ "A grow only, reference lookup of realized columns to block, block columns."
980
+ );
981
+
982
+ PyObject *
983
+ BlockIndex_new(PyTypeObject *cls, PyObject *args, PyObject *kwargs) {
984
+ BlockIndexObject *self = (BlockIndexObject *)cls->tp_alloc(cls, 0);
985
+ if (!self) {
986
+ return NULL;
987
+ }
988
+ return (PyObject *)self;
989
+ }
990
+
991
+ // Returns 0 on success, -1 on error.
992
+ int
993
+ BlockIndex_init(PyObject *self, PyObject *args, PyObject *kwargs) {
994
+ BlockIndexObject* bi = (BlockIndexObject*)self;
995
+
996
+ Py_ssize_t block_count = 0;
997
+ Py_ssize_t row_count = -1; // mark as unset
998
+ Py_ssize_t bir_count = 0;
999
+ Py_ssize_t bir_capacity = 8;
1000
+ PyObject* bir_bytes = NULL;
1001
+ PyObject* dtype = NULL;
1002
+
1003
+ if (!PyArg_ParseTuple(args,
1004
+ "|nnnnO!O:__init__",
1005
+ &block_count,
1006
+ &row_count,
1007
+ &bir_count,
1008
+ &bir_capacity,
1009
+ &PyBytes_Type, &bir_bytes,
1010
+ &dtype)) {
1011
+ return -1;
1012
+ }
1013
+ if (bir_count > bir_capacity) {
1014
+ PyErr_SetString(PyExc_ValueError, "record count exceeds capacity");
1015
+ return -1;
1016
+ }
1017
+ // handle all Py_ssize_t
1018
+ bi->block_count = block_count;
1019
+ bi->row_count = row_count;
1020
+ bi->bir_count = bir_count;
1021
+ bi->bir_capacity = bir_capacity;
1022
+
1023
+ bi->shape_recache = true; // always init to true
1024
+ bi->shape = NULL;
1025
+
1026
+ // Load the bi->bir struct array, if defined
1027
+ bi->bir = NULL;
1028
+ // always set bi to capacity defined at this point
1029
+ if (AK_BI_BIR_new(bi)) {
1030
+ return -1;
1031
+ }
1032
+ if (bir_bytes != NULL) {
1033
+ // already know bir is a bytes object
1034
+ char* data = PyBytes_AS_STRING(bir_bytes);
1035
+ memcpy(bi->bir, data, bi->bir_count * sizeof(BlockIndexRecord));
1036
+ // bir_bytes is a borrowed ref
1037
+ }
1038
+
1039
+ bi->dtype = NULL;
1040
+ if (dtype != NULL && dtype != Py_None) {
1041
+ if (PyArray_DescrCheck(dtype)) {
1042
+ Py_INCREF(dtype);
1043
+ bi->dtype = (PyArray_Descr*)dtype;
1044
+ }
1045
+ else {
1046
+ PyErr_SetString(PyExc_TypeError, "dtype argument must be a dtype");
1047
+ return -1;
1048
+ }
1049
+ }
1050
+ return 0;
1051
+ }
1052
+
1053
+ void
1054
+ BlockIndex_dealloc(BlockIndexObject *self) {
1055
+ if (self->bir != NULL) {
1056
+ PyMem_Free(self->bir);
1057
+ }
1058
+ // both dtype and shape might not be set
1059
+ Py_XDECREF((PyObject*)self->dtype);
1060
+ Py_XDECREF(self->shape);
1061
+ Py_TYPE(self)->tp_free((PyObject *)self);
1062
+ }
1063
+
1064
+ //------------------------------------------------------------------------------
1065
+
1066
+ // Returns NULL on error, True if the block should be reatained, False if the block has zero columns and should not be retained. This checks and raises on non-array inputs, dimensions other than 1 or 2, and mis-aligned columns.
1067
+ PyObject *
1068
+ BlockIndex_register(BlockIndexObject *self, PyObject *value) {
1069
+ if (!PyArray_Check(value)) {
1070
+ PyErr_Format(ErrorInitTypeBlocks, "Found non-array block: %R", value);
1071
+ return NULL;
1072
+ }
1073
+ PyArrayObject *a = (PyArrayObject *)value;
1074
+ int ndim = PyArray_NDIM(a);
1075
+
1076
+ if (ndim < 1 || ndim > 2) {
1077
+ PyErr_Format(ErrorInitTypeBlocks, "Array block has invalid dimensions: %i", ndim);
1078
+ return NULL;
1079
+ }
1080
+ Py_ssize_t increment = ndim == 1 ? 1 : PyArray_DIM(a, 1);
1081
+
1082
+ // assign alignment on first observation; otherwise force alignemnt. We do this regardless of if the array has no columns.
1083
+ Py_ssize_t alignment = PyArray_DIM(a, 0);
1084
+ if (self->row_count == -1) {
1085
+ self->row_count = alignment;
1086
+ self->shape_recache = true; // setting rows, must recache shape
1087
+ }
1088
+ else if (self->row_count != alignment) {
1089
+ PyErr_Format(ErrorInitTypeBlocks,
1090
+ "Array block has unaligned row count: found %i, expected %i",
1091
+ alignment,
1092
+ self->row_count);
1093
+ return NULL;
1094
+ }
1095
+ // if we are not adding columns, we are not adding types, so we are not changing the dtype or shape
1096
+ if (increment == 0) {
1097
+ Py_RETURN_FALSE;
1098
+ }
1099
+
1100
+ PyArray_Descr* dt = PyArray_DESCR(a); // borrowed ref
1101
+ self->shape_recache = true; // adjusting columns, must recache shape
1102
+
1103
+ if (self->dtype == NULL) { // if not already set
1104
+ Py_INCREF((PyObject*)dt);
1105
+ self->dtype = dt;
1106
+ }
1107
+ else if (!PyDataType_ISOBJECT(self->dtype)) { // if object cannot resolve further
1108
+ PyArray_Descr* dtr = AK_resolve_dtype(self->dtype, dt); // new ref
1109
+ if (dtr == NULL) {
1110
+ return NULL;
1111
+ }
1112
+ Py_DECREF((PyObject*)self->dtype);
1113
+ self->dtype = dtr;
1114
+ }
1115
+
1116
+ // create space for increment new records
1117
+ if (AK_BI_BIR_resize(self, increment)) {
1118
+ return NULL;
1119
+ };
1120
+
1121
+ // pull out references
1122
+ BlockIndexRecord* bir = self->bir;
1123
+ Py_ssize_t bc = self->block_count;
1124
+ Py_ssize_t birc = self->bir_count;
1125
+ for (Py_ssize_t i = 0; i < increment; i++) {
1126
+ bir[birc] = (BlockIndexRecord){bc, i};
1127
+ birc++;
1128
+ }
1129
+ self->bir_count = birc;
1130
+ self->block_count++;
1131
+ Py_RETURN_TRUE;
1132
+ }
1133
+
1134
+ //------------------------------------------------------------------------------
1135
+ // exporters
1136
+
1137
+ PyObject *
1138
+ BlockIndex_to_list(BlockIndexObject *self, PyObject *Py_UNUSED(unused)) {
1139
+ PyObject* list = PyList_New(self->bir_count);
1140
+ if (list == NULL) {
1141
+ return NULL;
1142
+ }
1143
+ BlockIndexRecord* bir = self->bir;
1144
+
1145
+ for (Py_ssize_t i = 0; i < self->bir_count; i++) {
1146
+ PyObject* item = AK_build_pair_ssize_t(bir[i].block, bir[i].column);
1147
+ if (item == NULL) {
1148
+ Py_DECREF(list);
1149
+ return NULL;
1150
+ }
1151
+ // set_item steals reference
1152
+ PyList_SET_ITEM(list, i, item);
1153
+ }
1154
+ return list;
1155
+ }
1156
+
1157
+ // Returns NULL on error
1158
+ static inline PyObject *
1159
+ AK_BI_to_bytes(BlockIndexObject *self) {
1160
+ Py_ssize_t size = self->bir_count * sizeof(BlockIndexRecord);
1161
+ // bytes might be null on error
1162
+ PyObject* bytes = PyBytes_FromStringAndSize((const char*)self->bir, size);
1163
+ return bytes;
1164
+ }
1165
+
1166
+ // Returns NULL on error
1167
+ PyObject *
1168
+ BlockIndex_to_bytes(BlockIndexObject *self, PyObject *Py_UNUSED(unused)) {
1169
+ return AK_BI_to_bytes(self);
1170
+ }
1171
+
1172
+ //------------------------------------------------------------------------------
1173
+ // pickle support
1174
+
1175
+ // Returns NULL on error, PyObject* otherwise.
1176
+ PyObject *
1177
+ BlockIndex_getstate(BlockIndexObject *self) {
1178
+ PyObject* bi = AK_BI_to_bytes(self);
1179
+ if (bi == NULL) {
1180
+ return NULL;
1181
+ }
1182
+ PyObject* dt = self->dtype == NULL ? Py_None : (PyObject*) self->dtype;
1183
+ // state might be NULL on failure; assume exception set
1184
+ PyObject* state = Py_BuildValue("nnnnNO", // use N to steal ref of bytes
1185
+ self->block_count,
1186
+ self->row_count,
1187
+ self->bir_count,
1188
+ self->bir_capacity,
1189
+ bi, // stolen new ref
1190
+ dt); // increfs passed object
1191
+ return state;
1192
+ }
1193
+
1194
+ // State returned here is a tuple of keys, suitable for usage as an `args` argument.
1195
+ PyObject *
1196
+ BlockIndex_setstate(BlockIndexObject *self, PyObject *state)
1197
+ {
1198
+ if (!PyTuple_CheckExact(state) || !PyTuple_GET_SIZE(state)) {
1199
+ PyErr_SetString(PyExc_ValueError, "Unexpected pickled object.");
1200
+ return NULL;
1201
+ }
1202
+ BlockIndex_init((PyObject*)self, state, NULL);
1203
+ Py_RETURN_NONE;
1204
+ }
1205
+
1206
+ //------------------------------------------------------------------------------
1207
+ // getters
1208
+
1209
+ // In a shape tuple, rows will never be negative.
1210
+ PyObject *
1211
+ BlockIndex_shape_getter(BlockIndexObject *self, void* Py_UNUSED(closure))
1212
+ {
1213
+ if (self->shape == NULL || self->shape_recache) {
1214
+ Py_XDECREF(self->shape); // get rid of old if it exists
1215
+ self->shape = AK_build_pair_ssize_t(
1216
+ self->row_count < 0 ? 0 : self->row_count,
1217
+ self->bir_count);
1218
+ }
1219
+ // shape is not null and shape_recache is false
1220
+ Py_INCREF(self->shape); // for caller
1221
+ self->shape_recache = false;
1222
+ return self->shape;
1223
+ }
1224
+
1225
+ // Unset rows will be -1.
1226
+ PyObject *
1227
+ BlockIndex_rows_getter(BlockIndexObject *self, void* Py_UNUSED(closure)){
1228
+ return PyLong_FromSsize_t(self->row_count);
1229
+ }
1230
+
1231
+ PyObject *
1232
+ BlockIndex_columns_getter(BlockIndexObject *self, void* Py_UNUSED(closure)){
1233
+ return PyLong_FromSsize_t(self->bir_count);
1234
+ }
1235
+
1236
+ // Return the resolved dtype for all registered blocks. If no block have been registered, this will return a float dtype.
1237
+ PyObject *
1238
+ BlockIndex_dtype_getter(BlockIndexObject *self, void* Py_UNUSED(closure)){
1239
+ if (self->dtype != NULL) {
1240
+ Py_INCREF(self->dtype);
1241
+ return (PyObject*)self->dtype;
1242
+ }
1243
+ // NOTE: could use NPY_DEFAULT_TYPE here; SF defines this explicitly as float64
1244
+ return (PyObject*)PyArray_DescrFromType(NPY_FLOAT64);
1245
+ }
1246
+
1247
+ static struct PyGetSetDef BlockIndex_getset[] = {
1248
+ {"shape", (getter)BlockIndex_shape_getter, NULL, NULL, NULL},
1249
+ {"rows", (getter)BlockIndex_rows_getter, NULL, NULL, NULL},
1250
+ {"columns", (getter)BlockIndex_columns_getter, NULL, NULL, NULL},
1251
+ {"dtype", (getter)BlockIndex_dtype_getter, NULL, NULL, NULL},
1252
+ {NULL},
1253
+ };
1254
+
1255
+ //------------------------------------------------------------------------------
1256
+ // general methods
1257
+
1258
+ PyObject *
1259
+ BlockIndex_repr(BlockIndexObject *self) {
1260
+ PyObject* dt = self->dtype == NULL ? Py_None : (PyObject*) self->dtype;
1261
+ return PyUnicode_FromFormat("<%s(blocks: %i, rows: %i, columns: %i, dtype: %R)>",
1262
+ Py_TYPE(self)->tp_name,
1263
+ self->block_count,
1264
+ self->row_count,
1265
+ self->bir_count,
1266
+ dt);
1267
+ }
1268
+
1269
+ PyObject *
1270
+ BlockIndex_copy(BlockIndexObject *self, PyObject *Py_UNUSED(unused))
1271
+ {
1272
+ PyTypeObject* cls = Py_TYPE(self); // borrowed ref
1273
+ BlockIndexObject *bi = (BlockIndexObject *)cls->tp_alloc(cls, 0);
1274
+ if (bi == NULL) {
1275
+ return NULL;
1276
+ }
1277
+ bi->block_count = self->block_count;
1278
+ bi->row_count = self->row_count;
1279
+ bi->bir_count = self->bir_count;
1280
+ bi->bir_capacity = self->bir_capacity;
1281
+
1282
+ bi->shape_recache = true; // could copy, but do not want to copy a pending cache state
1283
+ bi->shape = NULL;
1284
+
1285
+ bi->bir = NULL;
1286
+ AK_BI_BIR_new(bi); // do initial alloc to self->bir_capacity
1287
+ memcpy(bi->bir,
1288
+ self->bir,
1289
+ self->bir_count * sizeof(BlockIndexRecord));
1290
+
1291
+ bi->dtype = NULL;
1292
+ if (self->dtype != NULL) {
1293
+ bi->dtype = self->dtype;
1294
+ Py_INCREF((PyObject*)bi->dtype);
1295
+ }
1296
+ return (PyObject *)bi;
1297
+ }
1298
+
1299
+ Py_ssize_t
1300
+ BlockIndex_length(BlockIndexObject *self){
1301
+ return self->bir_count;
1302
+ }
1303
+
1304
+ PyObject *
1305
+ BlockIndex_sizeof(BlockIndexObject *self) {
1306
+ return PyLong_FromSsize_t(
1307
+ Py_TYPE(self)->tp_basicsize
1308
+ + (self->bir_capacity) * sizeof(BlockIndexRecord)
1309
+ );
1310
+ }
1311
+
1312
+ // Given an index, return just the block index.
1313
+ PyObject *
1314
+ BlockIndex_get_block(BlockIndexObject *self, PyObject *key){
1315
+ if (PyNumber_Check(key)) {
1316
+ Py_ssize_t i = PyNumber_AsSsize_t(key, NULL);
1317
+ if (!((size_t)i < (size_t)self->bir_count)) {
1318
+ PyErr_SetString(PyExc_IndexError, "index out of range");
1319
+ return NULL;
1320
+ }
1321
+ return PyLong_FromSsize_t(self->bir[i].block); // maybe NULL, exception will be set
1322
+ }
1323
+ PyErr_SetString(PyExc_TypeError, "An integer is required.");
1324
+ return NULL;
1325
+ }
1326
+
1327
+ // Given an index, return just the column index.
1328
+ PyObject *
1329
+ BlockIndex_get_column(BlockIndexObject *self, PyObject *key){
1330
+ if (PyNumber_Check(key)) {
1331
+ Py_ssize_t i = PyNumber_AsSsize_t(key, NULL);
1332
+ if (!((size_t)i < (size_t)self->bir_count)) {
1333
+ PyErr_SetString(PyExc_IndexError, "index out of range");
1334
+ return NULL;
1335
+ }
1336
+ return PyLong_FromSsize_t(self->bir[i].column); // maybe NULL, exception will be set
1337
+ }
1338
+ PyErr_SetString(PyExc_TypeError, "An integer is required.");
1339
+ return NULL;
1340
+ }
1341
+
1342
+ //------------------------------------------------------------------------------
1343
+ // iterators
1344
+
1345
+ PyObject *
1346
+ BlockIndex_iter(BlockIndexObject* self) {
1347
+ return BIIter_new(self, false);
1348
+ }
1349
+
1350
+ PyObject *
1351
+ BlockIndex_reversed(BlockIndexObject* self) {
1352
+ return BIIter_new(self, true);
1353
+ }
1354
+
1355
+ // Given key, return an iterator of a selection.
1356
+ PyObject *
1357
+ BlockIndex_iter_select(BlockIndexObject *self, PyObject *selector){
1358
+ return BIIterSelector_new(self, selector, false, BIIS_UNKNOWN, false);
1359
+ }
1360
+
1361
+ static char *iter_contiguous_kargs_names[] = {
1362
+ "selector",
1363
+ "ascending",
1364
+ "reduce",
1365
+ NULL
1366
+ };
1367
+
1368
+ // Given key, return an iterator of a selection.
1369
+ PyObject *
1370
+ BlockIndex_iter_contiguous(BlockIndexObject *self, PyObject *args, PyObject *kwargs)
1371
+ {
1372
+ PyObject* selector;
1373
+ int ascending = 0; // must be int for parsing to "p"
1374
+ int reduce = 0;
1375
+
1376
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs,
1377
+ "O|$pp:iter_contiguous",
1378
+ iter_contiguous_kargs_names,
1379
+ &selector,
1380
+ &ascending,
1381
+ &reduce
1382
+ )) {
1383
+ return NULL;
1384
+ }
1385
+ PyObject* iter = BIIterSelector_new(self, selector, false, BIIS_UNKNOWN, ascending);
1386
+ if (iter == NULL) {
1387
+ return NULL; // exception set
1388
+ }
1389
+ PyObject* biiter = BIIterContiguous_new(self, false, iter, reduce); // might be NULL, steals iter ref
1390
+ return biiter;
1391
+ }
1392
+
1393
+ // Given key, return an iterator of a selection.
1394
+ PyObject *
1395
+ BlockIndex_iter_block(BlockIndexObject *self){
1396
+ return BIIterBlock_new(self, false);
1397
+ }
1398
+
1399
+ //------------------------------------------------------------------------------
1400
+ // slot / method def
1401
+
1402
+ static PySequenceMethods BlockIndex_as_sequece = {
1403
+ .sq_length = (lenfunc)BlockIndex_length,
1404
+ .sq_item = (ssizeargfunc)AK_BI_item,
1405
+ };
1406
+
1407
+ static PyMethodDef BlockIndex_methods[] = {
1408
+ // {"__getitem__", (PyCFunction)BlockIndex_subscript, METH_O, NULL},
1409
+ {"register", (PyCFunction)BlockIndex_register, METH_O, NULL},
1410
+ {"__getstate__", (PyCFunction) BlockIndex_getstate, METH_NOARGS, NULL},
1411
+ {"__setstate__", (PyCFunction) BlockIndex_setstate, METH_O, NULL},
1412
+ {"__sizeof__", (PyCFunction) BlockIndex_sizeof, METH_NOARGS, NULL},
1413
+ {"__reversed__", (PyCFunction) BlockIndex_reversed, METH_NOARGS, NULL},
1414
+ {"to_list", (PyCFunction)BlockIndex_to_list, METH_NOARGS, NULL},
1415
+ {"to_bytes", (PyCFunction)BlockIndex_to_bytes, METH_NOARGS, NULL},
1416
+ {"copy", (PyCFunction)BlockIndex_copy, METH_NOARGS, NULL},
1417
+ {"get_block", (PyCFunction) BlockIndex_get_block, METH_O, NULL},
1418
+ {"get_column", (PyCFunction) BlockIndex_get_column, METH_O, NULL},
1419
+ {"iter_select", (PyCFunction) BlockIndex_iter_select, METH_O, NULL},
1420
+ {"iter_contiguous",
1421
+ (PyCFunction) BlockIndex_iter_contiguous,
1422
+ METH_VARARGS | METH_KEYWORDS,
1423
+ NULL},
1424
+ {"iter_block", (PyCFunction) BlockIndex_iter_block, METH_NOARGS, NULL},
1425
+ // {"__getnewargs__", (PyCFunction)BlockIndex_getnewargs, METH_NOARGS, NULL},
1426
+ {NULL},
1427
+ };
1428
+
1429
+ PyTypeObject BlockIndexType = {
1430
+ PyVarObject_HEAD_INIT(NULL, 0)
1431
+ // .tp_as_mapping = &BlockIndex_as_mapping,
1432
+ .tp_as_sequence = &BlockIndex_as_sequece,
1433
+ .tp_basicsize = sizeof(BlockIndexObject),
1434
+ .tp_dealloc = (destructor)BlockIndex_dealloc,
1435
+ .tp_doc = BlockIndex_doc,
1436
+ .tp_flags = Py_TPFLAGS_DEFAULT,
1437
+ .tp_getset = BlockIndex_getset,
1438
+ .tp_iter = (getiterfunc)BlockIndex_iter,
1439
+ .tp_methods = BlockIndex_methods,
1440
+ .tp_name = "arraykit.BlockIndex",
1441
+ .tp_new = BlockIndex_new,
1442
+ .tp_init = BlockIndex_init,
1443
+ .tp_repr = (reprfunc) BlockIndex_repr,
1444
+ // .tp_traverse = (traverseproc)BlockIndex_traverse,
1445
+ };