├── 2023
    ├── Makefile
    ├── native
    │   ├── __init__.py
    │   ├── chunked_stream.h
    │   ├── chunked_stream.pxd
    │   ├── cpython.pxd
    │   ├── mi_heap_destroy_stl_allocator.h
    │   ├── mi_heap_destroy_stl_allocator.pxd
    │   ├── mi_heap_destroy_stl_allocator.pyx
    │   ├── numpy.pxd
    │   ├── optional.pxd
    │   ├── sorted_intersection.h
    │   ├── string_view.pxd
    │   ├── sum_repeated_with_step.h
    │   └── utf8.pxd
    └── pyx
    │   ├── asyncpg_recordobj.h
    │   ├── dag_accelerated.h
    │   ├── dag_accelerated.pyx
    │   ├── interval_intersections.h
    │   ├── interval_intersections.pyx
    │   ├── io.pyx
    │   ├── object_arrays.pyx
    │   ├── sentry_native.pyx
    │   ├── sorted_ops.pyx
    │   ├── sql_builders.h
    │   ├── sql_builders.pyx
    │   ├── types_accelerated.pyx
    │   ├── unordered_unique.pyx
    │   ├── utils_accelerated.h
    │   ├── utils_accelerated.pyx
    │   ├── web_model_io.h
    │   └── web_model_io.pyx
├── LICENSE
├── README.md
├── async_utils.py
├── asyncpg_recordobj.h
├── to_object_arrays.pyx
└── typing_utils.py


/2023/Makefile:
--------------------------------------------------------------------------------
 1 | .PHONY: build-native
 2 | build-native:
 3 | 	cmake -S athenian/api/sentry_native -B athenian/api/sentry_native/build -D SENTRY_BACKEND=crashpad -D SENTRY_BUILD_EXAMPLES=OFF -D SENTRY_BUILD_TESTS=OFF -D CMAKE_BUILD_TYPE=RelWithDebInfo
 4 | 	cmake --build athenian/api/sentry_native/build --parallel
 5 | 	cmake -S athenian/api/mimalloc -B athenian/api/mimalloc/build -D mi_cflags=-flto -D MI_BUILD_STATIC=OFF -D MI_BUILD_OBJECT=OFF -D MI_BUILD_TESTS=OFF -D MI_INSTALL_TOPLEVEL=ON -D MI_USE_CXX=OFF -D CMAKE_BUILD_TYPE=RelWithDebInfo
 6 | 	cmake --build athenian/api/mimalloc/build --parallel
 7 | 
 8 | .PHONY: install-native
 9 | install-native: build-native
10 | 	sudo cmake --install athenian/api/sentry_native/build
11 | 	sudo cmake --install athenian/api/mimalloc/build
12 | 
13 | PHONY: install-native-user
14 | install-native-user: build-native
15 | 	cmake --install athenian/api/sentry_native/build
16 | 	cmake --install athenian/api/mimalloc/build
17 | 
18 | PHONY: clean-native
19 | clean-native:
20 | 	rm -rf athenian/api/sentry_native/build
21 | 	rm -rf athenian/api/mimalloc/build
22 | 


--------------------------------------------------------------------------------
/2023/native/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/athenianco/athenian-api-open/810d2b894ddf1a2ea28d22e78d7eb34f52a0f32a/2023/native/__init__.py


--------------------------------------------------------------------------------
/2023/native/chunked_stream.h:
--------------------------------------------------------------------------------
 1 | #include <list>
 2 | #include "mi_heap_destroy_stl_allocator.h"
 3 | 
 4 | template <int chunk_size = (1 << 16)>
 5 | class chunked_stream {
 6 |  public:
 7 |   template <typename T>
 8 |   explicit chunked_stream(mi_heap_destroy_stl_allocator<T> &alloc): chunks_(alloc), pos_(0) {
 9 |     chunks_.emplace_back(alloc).reserve(chunk_size);
10 |   }
11 | 
12 |   void write(const void *buffer, size_t size) {
13 |     const char *input = reinterpret_cast<const char *>(buffer);
14 |     int avail = chunk_size - pos_;
15 | 	while (size > static_cast<size_t>(avail)) {
16 | 	  memcpy(chunks_.back().data() + pos_, input, avail);
17 | 	  size -= avail;
18 | 	  input += avail;
19 | 	  pos_ = 0;
20 | 	  chunks_.emplace_back(chunks_.get_allocator()).reserve(chunk_size);
21 | 	  avail = chunk_size;
22 | 	}
23 | 	memcpy(chunks_.back().data() + pos_, input, size);
24 | 	pos_ += size;
25 |   }
26 | 
27 |   size_t dump(char *output, size_t output_size) noexcept {
28 |     size_t total_size = size();
29 |     if (output_size > total_size) {
30 |       output_size = total_size;
31 |     }
32 |     size_t left = output_size;
33 |     auto it = chunks_.begin();
34 |     while (left > chunk_size) {
35 |       memcpy(output, it->data(), chunk_size);
36 |       left -= chunk_size;
37 |       output += chunk_size;
38 |       it++;
39 |     }
40 |     if (left > 0) {
41 |       memcpy(output, it->data(), left);
42 |     }
43 |     return output_size;
44 |   }
45 | 
46 |   size_t size() const noexcept { return (chunks_.size() - 1) * chunk_size + pos_; }
47 | 
48 |  private:
49 |   std::list<mi_vector<char>, mi_heap_destroy_stl_allocator<mi_vector<char>>> chunks_;
50 |   int pos_;
51 | };
52 | 


--------------------------------------------------------------------------------
/2023/native/chunked_stream.pxd:
--------------------------------------------------------------------------------
 1 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport mi_heap_destroy_stl_allocator
 2 | 
 3 | 
 4 | cdef extern from "chunked_stream.h" nogil:
 5 |     cdef cppclass chunked_stream[I=*]:
 6 |         chunked_stream chunked_stream[X](mi_heap_destroy_stl_allocator[X] &) except +
 7 |         void write(const void *buffer, size_t size)
 8 |         size_t dump(char *output, size_t output_size)
 9 |         size_t size()
10 | 


--------------------------------------------------------------------------------
/2023/native/cpython.pxd:
--------------------------------------------------------------------------------
  1 | from cpython cimport PyObject
  2 | 
  3 | ctypedef PyObject *PyObjectPtr
  4 | 
  5 | cdef extern from "structmember.h":
  6 |     ctypedef struct PyMemberDef:
  7 |         const char *name
  8 |         int type
  9 |         Py_ssize_t offset
 10 |         int flags
 11 |         const char *doc
 12 | 
 13 | cdef extern from "Python.h":
 14 |     ctypedef PyObject *(*allocfunc)(PyTypeObject *cls, Py_ssize_t nitems)
 15 | 
 16 |     ctypedef struct PyTypeObject:
 17 |         allocfunc tp_alloc
 18 |         PyMemberDef *tp_members
 19 | 
 20 |     bint PyObject_TypeCheck(PyObject *, PyTypeObject *) nogil
 21 |     PyTypeObject *Py_TYPE(const PyObject *) nogil
 22 | 
 23 |     bint PyLong_CheckExact(PyObject *) nogil
 24 |     long PyLong_AsLong(PyObject *) nogil
 25 | 
 26 |     double PyFloat_AS_DOUBLE(PyObject *) nogil
 27 |     bint PyFloat_CheckExact(PyObject *) nogil
 28 |     bint PyFloat_Check(PyObject *) nogil
 29 | 
 30 |     PyObject *PyList_New(Py_ssize_t len)
 31 |     bint PyList_CheckExact(PyObject *) nogil
 32 |     Py_ssize_t PyList_GET_SIZE(PyObject *) nogil
 33 |     PyObject *PyList_GET_ITEM(PyObject *, Py_ssize_t) nogil
 34 |     void PyList_SET_ITEM(PyObject *list, Py_ssize_t i, PyObject *o) nogil
 35 | 
 36 |     PyObject *PyTuple_GET_ITEM(PyObject *, Py_ssize_t) nogil
 37 | 
 38 |     bint PyDict_CheckExact(PyObject *) nogil
 39 |     int PyDict_Next(PyObject *p, Py_ssize_t *ppos, PyObject **pkey, PyObject **pvalue) nogil
 40 |     Py_ssize_t PyDict_Size(PyObject *p) nogil
 41 | 
 42 |     bint PyUnicode_Check(PyObject *) nogil
 43 |     Py_ssize_t PyUnicode_GET_LENGTH(PyObject *) nogil
 44 |     unsigned int PyUnicode_KIND(PyObject *) nogil
 45 |     void *PyUnicode_DATA(PyObject *) nogil
 46 |     Py_ssize_t PyUnicode_FindChar(
 47 |         PyObject *str, Py_UCS4 ch, Py_ssize_t start, Py_ssize_t end, int direction
 48 |     ) nogil
 49 | 
 50 |     bint PyBytes_Check(PyObject *) nogil
 51 |     char *PyBytes_AS_STRING(PyObject *) nogil
 52 |     Py_ssize_t PyBytes_GET_SIZE(PyObject *) nogil
 53 | 
 54 |     bint PyByteArray_CheckExact(PyObject *) nogil
 55 |     char *PyByteArray_AS_STRING(PyObject *) nogil
 56 | 
 57 |     unsigned int PyUnicode_1BYTE_KIND
 58 |     unsigned int PyUnicode_2BYTE_KIND
 59 |     unsigned int PyUnicode_4BYTE_KIND
 60 | 
 61 |     PyObject *Py_None
 62 |     PyObject *Py_True
 63 |     PyObject *Py_False
 64 | 
 65 |     PyTypeObject PyLong_Type
 66 |     PyTypeObject PyFloat_Type
 67 |     PyTypeObject PyUnicode_Type
 68 |     PyTypeObject PyBool_Type
 69 |     PyTypeObject PyList_Type
 70 |     PyTypeObject PyDict_Type
 71 |     PyTypeObject PyBaseObject_Type
 72 | 
 73 |     void Py_INCREF(PyObject *)
 74 |     void Py_DECREF(PyObject *)
 75 | 
 76 |     object PyUnicode_FromStringAndSize(const char *, Py_ssize_t)
 77 |     str PyUnicode_FromKindAndData(unsigned int kind, void *buffer, Py_ssize_t size)
 78 |     str PyUnicode_FromString(const char *)
 79 |     object PyUnicode_New(Py_ssize_t, Py_UCS4)
 80 |     PyObject *PyBytes_FromStringAndSize(char *v, Py_ssize_t len)
 81 |     PyObject *PyLong_FromLong(long v)
 82 |     PyObject *PyObject_GetItem(PyObject *o, PyObject *key)
 83 |     PyObject *PyObject_GetAttr(PyObject *o, PyObject *attr_name)
 84 | 
 85 | cdef extern from "datetime.h" nogil:
 86 |     bint PyDateTime_Check(PyObject *)
 87 |     bint PyDelta_Check(PyObject *)
 88 | 
 89 |     int PyDateTime_GET_YEAR(PyObject *)
 90 |     int PyDateTime_GET_MONTH(PyObject *)
 91 |     int PyDateTime_GET_DAY(PyObject *)
 92 | 
 93 |     int PyDateTime_DATE_GET_HOUR(PyObject *)
 94 |     int PyDateTime_DATE_GET_MINUTE(PyObject *)
 95 |     int PyDateTime_DATE_GET_SECOND(PyObject *)
 96 | 
 97 |     int PyDateTime_DELTA_GET_DAYS(PyObject *)
 98 |     int PyDateTime_DELTA_GET_SECONDS(PyObject *)
 99 | 
100 |     PyObject *PyDateTime_DATE_GET_TZINFO(PyObject *)
101 | 
102 |     ctypedef struct PyDateTime_CAPI:
103 |         PyObject *TimeZone_UTC
104 | 


--------------------------------------------------------------------------------
/2023/native/mi_heap_destroy_stl_allocator.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | #include <unordered_map>
 5 | #include <unordered_set>
 6 | #include <vector>
 7 | 
 8 | #include <mimalloc.h>
 9 | 
10 | template<
11 |     class T,
12 |     class U,
13 |     class HASH = std::hash<T>,
14 |     class PRED = std::equal_to<T>
15 | >
16 | using mi_unordered_map = std::unordered_map<T, U, HASH, PRED, mi_heap_destroy_stl_allocator<std::pair<const T, U>>>;
17 | 
18 | template<
19 |     class T,
20 |     class HASH = std::hash<T>,
21 |     class PRED = std::equal_to<T>
22 | >
23 | using mi_unordered_set = std::unordered_set<T, HASH, PRED, mi_heap_destroy_stl_allocator<T>>;
24 | 
25 | template<class T>
26 | using mi_vector = std::vector<T, mi_heap_destroy_stl_allocator<T>>;
27 | 
28 | using mi_string = std::basic_string<char, std::char_traits<char>, mi_heap_destroy_stl_allocator<char>>;
29 | 
30 | namespace std {
31 |     template<> struct hash<mi_string> {
32 |         size_t operator()(const mi_string &s) const {
33 |             return std::hash<std::string_view>()(s);
34 |         }
35 |     };
36 | }
37 | 
38 | struct empty_deleter {
39 |   template <typename T>
40 |   void operator()(T *) const noexcept {}
41 | };
42 | 


--------------------------------------------------------------------------------
/2023/native/mi_heap_destroy_stl_allocator.pxd:
--------------------------------------------------------------------------------
 1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
 2 | # cython: warn.maybe_uninitialized=True
 3 | 
 4 | from cpython.pycapsule cimport PyCapsule_GetPointer
 5 | from libcpp cimport bool
 6 | from libcpp.string cimport string
 7 | from libcpp.unordered_map cimport pair, unordered_map
 8 | from libcpp.unordered_set cimport unordered_set
 9 | from libcpp.vector cimport vector
10 | 
11 | 
12 | cdef extern from "mi_heap_destroy_stl_allocator.h" nogil:
13 |     cdef cppclass mi_heap_destroy_stl_allocator[T]:
14 |         mi_heap_destroy_stl_allocator() except +
15 |         mi_heap_destroy_stl_allocator(const mi_heap_destroy_stl_allocator &)
16 |         T* allocate(size_t count) except +
17 |         void deallocate(T*)
18 | 
19 |     cdef cppclass empty_deleter:
20 |         empty_deleter()
21 | 
22 |     cdef cppclass mi_unordered_map[T, U, HASH=*, PRED=*](unordered_map[T, U, HASH, PRED]):
23 |         mi_unordered_map mi_unordered_map[X](mi_heap_destroy_stl_allocator[X]&) except +
24 |         pair[mi_unordered_map.iterator, bool] try_emplace(...) except +
25 |         mi_heap_destroy_stl_allocator[T] get_allocator()
26 | 
27 |     cdef cppclass mi_unordered_set[T, HASH=*, PRED=*](unordered_set[T, HASH, PRED]):
28 |         mi_unordered_set mi_unordered_set[X](mi_heap_destroy_stl_allocator[X]&) except +
29 |         pair[mi_unordered_set.iterator, bool] emplace(...) except +
30 |         mi_heap_destroy_stl_allocator[T] get_allocator()
31 | 
32 |         mi_unordered_set.iterator erase(mi_unordered_set.iterator)
33 |         mi_unordered_set.iterator erase(mi_unordered_set.iterator, mi_unordered_set.iterator)
34 |         size_t erase(T&)
35 | 
36 |     cdef cppclass mi_vector[T](vector[T]):
37 |         mi_vector mi_vector[X](mi_heap_destroy_stl_allocator[X]&) except +
38 |         T& emplace_back(...) except +
39 |         mi_heap_destroy_stl_allocator[T] get_allocator()
40 | 
41 |     cdef cppclass mi_string(string):
42 |         mi_string mi_string[X](const char *, size_t, mi_heap_destroy_stl_allocator[X]&) except +
43 |         mi_heap_destroy_stl_allocator[char] get_allocator()
44 | 
45 | 
46 | cdef inline mi_heap_destroy_stl_allocator[char] *mi_heap_allocator_from_capsule(obj) except? NULL:
47 |     return <mi_heap_destroy_stl_allocator[char] *> PyCapsule_GetPointer(obj, b"mi_heap_destroy_stl_allocator")
48 | 
49 | 
50 | cdef inline void _delete_mi_heap_allocator_in_capsule(obj):
51 |     cdef mi_heap_destroy_stl_allocator[char] *alloc = mi_heap_allocator_from_capsule(obj)
52 |     del alloc
53 | 


--------------------------------------------------------------------------------
/2023/native/mi_heap_destroy_stl_allocator.pyx:
--------------------------------------------------------------------------------
 1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
 2 | # cython: warn.maybe_uninitialized=True
 3 | # distutils: language = c++
 4 | # distutils: extra_compile_args = -std=c++17
 5 | # distutils: libraries = mimalloc
 6 | # distutils: runtime_library_dirs = /usr/local/lib
 7 | 
 8 | from cpython.pycapsule cimport PyCapsule_New
 9 | 
10 | 
11 | def make_mi_heap_allocator_capsule() -> object:
12 |     cdef mi_heap_destroy_stl_allocator[char] *alloc = new mi_heap_destroy_stl_allocator[char]()
13 |     return PyCapsule_New(alloc, b"mi_heap_destroy_stl_allocator", _delete_mi_heap_allocator_in_capsule)
14 | 


--------------------------------------------------------------------------------
/2023/native/numpy.pxd:
--------------------------------------------------------------------------------
 1 | from cpython cimport PyObject
 2 | from numpy cimport dtype as npdtype, npy_int64, npy_intp
 3 | 
 4 | from athenian.api.native.cpython cimport PyTypeObject
 5 | 
 6 | 
 7 | cdef extern from "numpy/arrayobject.h":
 8 |     PyTypeObject PyArray_Type
 9 |     PyTypeObject PyDatetimeArrType_Type
10 |     PyTypeObject PyDoubleArrType_Type
11 |     PyTypeObject PyIntegerArrType_Type
12 |     PyTypeObject PyFloatArrType_Type
13 |     PyTypeObject PyTimedeltaArrType_Type
14 | 
15 |     enum: NPY_DATETIME_NAT
16 | 
17 |     ctypedef struct PyArray_Descr:
18 |         char kind
19 |         char type
20 |         char byteorder
21 |         char flags
22 |         int type_num
23 |         int itemsize "elsize"
24 |         int alignment
25 | 
26 |     PyObject *PyArray_NewFromDescr(
27 |         PyTypeObject *subtype,
28 |         PyArray_Descr *descr,
29 |         int nd,
30 |         const npy_intp *dims,
31 |         const npy_intp *strides,
32 |         void *data,
33 |         int flags,
34 |         PyObject *obj,
35 |     )
36 |     npdtype PyArray_DescrNew(npdtype)
37 | 
38 |     void *PyArray_DATA(PyObject *) nogil
39 |     char *PyArray_BYTES(PyObject *) nogil
40 |     npy_intp PyArray_DIM(PyObject *, size_t) nogil
41 |     npy_intp PyArray_STRIDE(PyObject *, size_t) nogil
42 |     int PyArray_NDIM(PyObject *) nogil
43 |     npy_intp PyArray_ITEMSIZE(PyObject *) nogil
44 |     bint PyArray_CheckExact(PyObject *) nogil
45 |     PyArray_Descr *PyArray_DESCR(PyObject *) nogil
46 |     int PyArray_TYPE(PyObject *) nogil
47 |     bint PyArray_IS_C_CONTIGUOUS(PyObject *) nogil
48 |     bint PyArray_IS_F_CONTIGUOUS(PyObject *) nogil
49 |     void PyArray_ScalarAsCtype(PyObject *scalar, void *ctypeptr) nogil
50 | 
51 |     ctypedef enum NPY_DATETIMEUNIT:
52 |         NPY_FR_ERROR = -1
53 |         NPY_FR_M = 1
54 |         NPY_FR_W = 2
55 |         NPY_FR_D = 4
56 |         NPY_FR_h = 5
57 |         NPY_FR_m = 6
58 |         NPY_FR_s = 7
59 |         NPY_FR_ms = 8
60 |         NPY_FR_us = 9
61 |         NPY_FR_ns = 10
62 |         NPY_FR_ps = 11
63 |         NPY_FR_fs = 12
64 |         NPY_FR_as = 13
65 |         NPY_FR_GENERIC = 14
66 | 
67 |     ctypedef struct PyArray_DatetimeMetaData:
68 |         NPY_DATETIMEUNIT base
69 |         int num
70 | 
71 |     ctypedef struct PyDatetimeScalarObject:
72 |         npy_int64 obval
73 |         PyArray_DatetimeMetaData obmeta
74 | 


--------------------------------------------------------------------------------
/2023/native/optional.pxd:
--------------------------------------------------------------------------------
 1 | # backported from Cython 3.0 alpha
 2 | 
 3 | from libcpp cimport bool
 4 | 
 5 | 
 6 | cdef extern from "<optional>" namespace "std" nogil:
 7 |     cdef cppclass nullopt_t:
 8 |         nullopt_t()
 9 | 
10 |     cdef nullopt_t nullopt
11 | 
12 |     cdef cppclass optional[T]:
13 |         ctypedef T value_type
14 |         optional()
15 |         optional(nullopt_t)
16 |         optional(optional&) except +
17 |         optional(T&) except +
18 |         bool has_value()
19 |         T& value()
20 |         T& value_or[U](U& default_value)
21 |         void swap(optional&)
22 |         void reset()
23 |         T& emplace(...)
24 |         T& operator*()
25 |         #T* operator->() # Not Supported
26 |         optional& operator=(optional&)
27 |         optional& operator=[U](U&)
28 |         bool operator bool()
29 |         bool operator!()
30 |         bool operator==[U](optional&, U&)
31 |         bool operator!=[U](optional&, U&)
32 |         bool operator<[U](optional&, U&)
33 |         bool operator>[U](optional&, U&)
34 |         bool operator<=[U](optional&, U&)
35 |         bool operator>=[U](optional&, U&)
36 | 
37 |     optional[T] make_optional[T](...) except +


--------------------------------------------------------------------------------
/2023/native/string_view.pxd:
--------------------------------------------------------------------------------
1 | cdef extern from "<string_view>" namespace "std" nogil:
2 |     cppclass string_view:
3 |         string_view() except +
4 |         string_view(const char *, size_t) except +
5 |         const char *data()
6 |         size_t size()
7 | 


--------------------------------------------------------------------------------
/2023/native/sum_repeated_with_step.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <stdint.h>
 4 | #include <immintrin.h>
 5 | #define restrict __restrict__
 6 | 
 7 | void sum_repeated_with_step_avx2(
 8 |     const int64_t * restrict src,
 9 |     int64_t src_len,
10 |     int64_t repeats,
11 |     int64_t step,
12 |     int64_t * restrict dst
13 | ) {
14 |     __m256i offset = _mm256_set1_epi64x(0);
15 |     const __m256i step_vec = _mm256_set1_epi64x(step);
16 | 	for (int64_t i = 0; i < repeats; i++) {
17 | 		int64_t j;
18 | 		for (j = 0; j < src_len - 3; j += 4) {
19 | 			_mm256_storeu_si256(
20 | 				reinterpret_cast<__m256i *>(dst),
21 | 				_mm256_add_epi64(
22 | 					_mm256_loadu_si256(reinterpret_cast<const __m256i *>(src + j)),
23 | 					offset
24 | 				)
25 | 			);
26 | 			dst += 4;
27 | 		}
28 | 		for (; j < src_len; j++) {
29 | 			*dst++ = src[j] + step * i;
30 | 		}
31 | 		offset = _mm256_add_epi64(offset, step_vec);
32 | 	}
33 | }


--------------------------------------------------------------------------------
/2023/native/utf8.pxd:
--------------------------------------------------------------------------------
 1 | from libc.stdint cimport uint32_t
 2 | 
 3 | 
 4 | # Adapted from CPython, licensed under PSF2 (BSD-like)
 5 | cdef inline int ucs4_to_utf8_json(uint32_t ucs4, char *utf8) nogil:
 6 |     if ucs4 == 0:
 7 |         return 0
 8 |     if ucs4 == b"\\" or ucs4 == b'"':
 9 |         utf8[0] = b"\\"
10 |         utf8[1] = ucs4
11 |         return 2
12 |     if ucs4 < 0x20:
13 |         # Escape control chars
14 |         utf8[0] = b"\\"
15 |         utf8[1] = b"u"
16 |         utf8[2] = b"0"
17 |         utf8[3] = b"0"
18 |         utf8[4] = b"0" if ucs4 < 0x10 else b"1"
19 |         ucs4 &= 0x0F
20 |         if ucs4 > 0x09:
21 |             utf8[5] = (ucs4 - 0x0A) + ord(b"A")
22 |         else:
23 |             utf8[5] = ucs4 + ord(b"0")
24 |         return 6
25 |     if ucs4 < 0x80:
26 |         # Encode ASCII
27 |         utf8[0] = ucs4
28 |         return 1
29 |     if ucs4 < 0x0800:
30 |         # Encode Latin-1
31 |         utf8[0] = 0xc0 | (ucs4 >> 6)
32 |         utf8[1] = 0x80 | (ucs4 & 0x3f)
33 |         return 2
34 |     if 0xD800 <= ucs4 <= 0xDFFF:
35 |         return 0
36 |     if ucs4 < 0x10000:
37 |         utf8[0] = 0xe0 | (ucs4 >> 12)
38 |         utf8[1] = 0x80 | ((ucs4 >> 6) & 0x3f)
39 |         utf8[2] = 0x80 | (ucs4 & 0x3f)
40 |         return 3
41 |     # Encode UCS4 Unicode ordinals
42 |     utf8[0] = 0xf0 | (ucs4 >> 18)
43 |     utf8[1] = 0x80 | ((ucs4 >> 12) & 0x3f)
44 |     utf8[2] = 0x80 | ((ucs4 >> 6) & 0x3f)
45 |     utf8[3] = 0x80 | (ucs4 & 0x3f)
46 |     return 4
47 | 


--------------------------------------------------------------------------------
/2023/pyx/asyncpg_recordobj.h:
--------------------------------------------------------------------------------
 1 | // no need to #include anything, this file is used internally by to_object_arrays.pyx
 2 | 
 3 | typedef struct {
 4 |     PyObject_VAR_HEAD
 5 | 
 6 |     // asyncpg specifics begin here
 7 |     // if they add another field, we will break spectacularly
 8 |     Py_hash_t self_hash;
 9 |     PyObject *desc;  // we don't care of the actual type
10 |     PyObject *ob_item[1];  // embedded in the tail, the count matches len()
11 | } ApgRecordObject;
12 | 
13 | #define ApgRecord_GET_ITEM(op, i) (((ApgRecordObject *)(op))->ob_item[i])
14 | #define ApgRecord_SET_ITEM(op, i, v) (((ApgRecordObject *)(op))->ob_item[i] = v)
15 | #define ApgRecord_GET_DESC(op) (((ApgRecordObject *)(op))->desc)
16 | 


--------------------------------------------------------------------------------
/2023/pyx/dag_accelerated.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | #include <immintrin.h>
 3 | 
 4 | size_t sorted_set_difference_avx2(
 5 | 	const uint32_t *__restrict__ set1,
 6 | 	const size_t length1,
 7 | 	const uint32_t *__restrict__ set2,
 8 |     const size_t length2,
 9 |     uint32_t *__restrict__ out) {
10 |   __m256i left, right;
11 |   const __m256i ones = _mm256_set1_epi8(0xff);
12 |   size_t passed = 0;
13 |   const uint32_t *border_left = set1 + length1 - 8;
14 |   const uint32_t *border_right = set2 + length2 - 8;
15 |   while (set1 <= border_left && set2 <= border_right) {
16 |     left = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(set1));
17 |     right = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(set2));
18 |     __m256i c = _mm256_cmpeq_epi32(left, right);
19 |     if (_mm256_testc_si256(c, ones)) {
20 |     	set1 += 8;
21 |     	set2 += 8;
22 |     	continue;
23 |     }
24 |     int offset = __builtin_ctz(~static_cast<unsigned>(_mm256_movemask_epi8(c))) >> 2;
25 |     set1 += offset;
26 |     set2 += offset;
27 |     if (*set1 < *set2) {
28 |       out[passed++] = *set1++;
29 |     } else {
30 |       set2++;
31 |     }
32 |   }
33 |   border_left += 8;
34 |   border_right += 8;
35 |   while (set1 < border_left && set2 < border_right) {
36 |     uint32_t ileft = *set1;
37 |     uint32_t iright = *set2;
38 |     if (ileft == iright) {
39 |       set1++;
40 |       set2++;
41 |     } else if (ileft < iright) {
42 |       out[passed++] = ileft;
43 |       set1++;
44 |     } else {
45 |       set2++;
46 |     }
47 |   }
48 |   if (set2 == border_right) {
49 |     while (set1 < border_left) {
50 |       out[passed++] = *set1++;
51 |     }
52 |   }
53 |   return passed;
54 | }


--------------------------------------------------------------------------------
/2023/pyx/interval_intersections.h:
--------------------------------------------------------------------------------
1 | #include <algorithm>
2 | 
3 | template <typename T, typename I>
4 | void argsort_bodies(const mi_vector<T> &bodies, mi_vector<I> &indexes) noexcept {
5 |   std::sort(indexes.begin(), indexes.end(), [&bodies](I left, I right) -> bool {
6 | 	  return bodies[left] < bodies[right];
7 |   });
8 | }


--------------------------------------------------------------------------------
/2023/pyx/interval_intersections.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: extra_compile_args = -std=c++17
  5 | # distutils: libraries = mimalloc
  6 | # distutils: runtime_library_dirs = /usr/local/lib
  7 | 
  8 | cimport cython
  9 | from cython.operator cimport dereference as deref, postincrement
 10 | from libc.stdint cimport int64_t, uint64_t
 11 | from libcpp cimport bool
 12 | from libcpp.set cimport set
 13 | from libcpp.utility cimport move
 14 | 
 15 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport (
 16 |     mi_heap_allocator_from_capsule,
 17 |     mi_heap_destroy_stl_allocator,
 18 |     mi_string,
 19 |     mi_unordered_map,
 20 |     mi_vector,
 21 | )
 22 | from athenian.api.native.optional cimport optional
 23 | 
 24 | import numpy as np
 25 | 
 26 | 
 27 | # __builtin_clzl is a compiler built-in that counts the number of leading zeros
 28 | cdef extern int __builtin_clzl(unsigned long)
 29 | 
 30 | cdef inline int _leading_zero_bits(unsigned long n):
 31 |    if n == 0:
 32 |        # __builtin_clzl is undefined when called with 0
 33 |        return sizeof(n) * 8
 34 |    return __builtin_clzl(n)
 35 | 
 36 | 
 37 | def calculate_interval_intersections(starts: np.ndarray,
 38 |                                      finishes: np.ndarray,
 39 |                                      borders: np.ndarray,
 40 |                                      ) -> np.ndarray:
 41 |     cdef unsigned long max_intervals, groups_count, time_offset
 42 |     assert len(starts) == len(finishes)
 43 |     assert starts.dtype == np.uint64
 44 |     assert finishes.dtype == np.uint64
 45 |     if len(starts) == 0:
 46 |         return np.array([], dtype=float)
 47 |     time_offset = starts.min()
 48 |     # require less bits for the timestamps
 49 |     starts -= time_offset
 50 |     finishes -= time_offset
 51 |     # there can be intervals of zero length, make them 1-second
 52 |     finishes[starts >= finishes] += 1
 53 |     group_lengths = np.diff(borders, prepend=0)
 54 |     max_intervals = group_lengths.max()
 55 |     series = np.arange(max_intervals, dtype=np.uint64)
 56 |     intervals = np.concatenate([starts, finishes])
 57 |     size = len(starts)
 58 |     time_offset = 64 - _leading_zero_bits(max_intervals - 1)
 59 |     intervals <<= time_offset
 60 |     groups_count = len(borders)
 61 |     group_offset = _leading_zero_bits(groups_count - 1)
 62 |     group_indexes = np.repeat(np.arange(groups_count, dtype=np.uint64), group_lengths)
 63 | 
 64 |     intervals[:size] |= group_indexes << group_offset
 65 |     intervals[size:] |= group_indexes << group_offset
 66 |     # https://codereview.stackexchange.com/questions/83018/vectorized-numpy-version-of-arange-with-multiple-start-stop
 67 |     indexes = (
 68 |         np.repeat(group_lengths - group_lengths.cumsum(), group_lengths) +
 69 |         np.arange(group_lengths.sum())
 70 |     ).view(np.uint64)
 71 |     intervals[:size] |= indexes
 72 |     intervals[size:] |= indexes
 73 |     # bits 0..time_offset - interval indexes, each in range 0..group length
 74 |     # bits time_offset..(64 - group_offset) - timestamps
 75 |     # bits (64 - group_offset)..63 - group index
 76 |     # stable sort because starts must come before finishes if the timestamps are equal
 77 |     intervals = np.sort(intervals, kind="stable")
 78 |     # remove the group indexes
 79 |     intervals &= (1 << group_offset) - 1
 80 |     raw = np.zeros(size, dtype=np.uint64)
 81 |     cdef:
 82 |         const uint64_t[:] intervals_view = intervals
 83 |         const int64_t[:] borders_view = borders * 2
 84 |         uint64_t[:] raw_view = raw
 85 |     with nogil:
 86 |         _calculate_interval_intersections(intervals_view, borders_view, time_offset, raw_view)
 87 |     result = raw.astype(float) / (finishes - starts)
 88 |     return result
 89 | 
 90 | 
 91 | @cython.boundscheck(False)
 92 | @cython.wraparound(False)
 93 | cdef void _calculate_interval_intersections(const uint64_t[:] intervals,
 94 |                                             const int64_t[:] borders,
 95 |                                             char time_offset,
 96 |                                             uint64_t[:] intersections) nogil:
 97 |     cdef:
 98 |         int64_t i, j, border_index, group_start, group_finish, ii_open, intersections_offset
 99 |         uint64_t item, index_mask, timestamp, previous_timestamp, delta
100 |         # set faster than unordered_set because we iterate over all elements on each step
101 |         set[int64_t] open_intervals
102 |         set[int64_t].iterator ii
103 |     index_mask = (1 << time_offset) - 1
104 |     previous_timestamp = 0  # not really needed but removes the warning
105 |     for border_index in range(len(borders)):
106 |         group_start = borders[border_index - 1] if border_index > 0 else 0
107 |         intersections_offset = group_start >> 1
108 |         group_finish = borders[border_index]
109 |         for i in range(group_start, group_finish):
110 |             item = intervals[i]
111 |             timestamp = item >> time_offset
112 |             delta = (timestamp - previous_timestamp) * open_intervals.size()
113 |             for ii_open in open_intervals:
114 |                 intersections[intersections_offset + ii_open] += delta
115 |             interval_index = item & index_mask
116 |             ii = open_intervals.find(interval_index)
117 |             if ii == open_intervals.end():
118 |                 open_intervals.insert(interval_index)
119 |             else:
120 |                 open_intervals.erase(ii)
121 |             previous_timestamp = timestamp
122 |         open_intervals.clear()
123 | 
124 | 


--------------------------------------------------------------------------------
/2023/pyx/io.pyx:
--------------------------------------------------------------------------------
 1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
 2 | # cython: warn.maybe_uninitialized=True
 3 | # distutils: language = c++
 4 | # distutils: extra_compile_args = -std=c++17
 5 | 
 6 | import pickle
 7 | from typing import Any
 8 | 
 9 | from cpython cimport Py_INCREF, PyBytes_FromStringAndSize, PyObject, PyTuple_New, PyTuple_SET_ITEM
10 | from libc.stdint cimport uint32_t
11 | from libc.string cimport memcpy
12 | 
13 | from athenian.api.native.cpython cimport PyBytes_AS_STRING
14 | 
15 | from medvedi import DataFrame
16 | 
17 | 
18 | cdef extern from "<stdlib.h>" nogil:
19 |     char *gcvt(double number, int ndigit, char *buf)
20 | 
21 | 
22 | cdef extern from "<string.h>" nogil:
23 |     size_t strnlen(const char *, size_t)
24 | 
25 | 
26 | def serialize_args(tuple args, alloc_capsule=None) -> bytes:
27 |     cdef:
28 |         bytes result, buffer
29 |         Py_ssize_t size = 4
30 |         list buffers = []
31 |         char *output
32 |         bint is_df
33 | 
34 |     for arg in args:
35 |         if isinstance(arg, DataFrame):
36 |             is_df = True
37 |             buffer = arg.serialize_unsafe()
38 |         else:
39 |             is_df = False
40 |             buffer = pickle.dumps(arg)
41 |         size += len(buffer) + 5
42 |         buffers.append((is_df, buffer))
43 |     result = PyBytes_FromStringAndSize(NULL, size)
44 |     output = PyBytes_AS_STRING(<PyObject *> result)
45 |     (<uint32_t *> output)[0] = len(buffers)
46 |     output += 4
47 |     for is_df, buffer in buffers:
48 |         output[0] = is_df
49 |         output += 1
50 |         size = len(buffer)
51 |         (<uint32_t *> output)[0] = size
52 |         output += 4
53 |         memcpy(output, PyBytes_AS_STRING(<PyObject *> buffer), size)
54 |         output += size
55 |     return result
56 | 
57 | 
58 | def deserialize_args(bytes buffer) -> tuple[Any]:
59 |     cdef:
60 |         uint32_t size, i
61 |         tuple result
62 |         long offset = 4
63 |         object item
64 |         char is_df
65 | 
66 |     input = PyBytes_AS_STRING(<PyObject *> buffer)
67 |     size = (<uint32_t *> input)[0]
68 |     input += 4
69 |     result = PyTuple_New(size)
70 |     for i in range(size):
71 |         is_df = input[0]
72 |         input += 1
73 |         size = (<uint32_t *> input)[0]
74 |         input += 4
75 |         offset += 5
76 |         if is_df:
77 |             item = DataFrame.deserialize_unsafe(buffer[offset: offset + size])
78 |         else:
79 |             item = pickle.loads(buffer[offset: offset + size])
80 |         offset += size
81 |         input += size
82 |         Py_INCREF(item)
83 |         PyTuple_SET_ITEM(result, i, item)
84 |     return result
85 | 


--------------------------------------------------------------------------------
/2023/pyx/object_arrays.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: libraries = mimalloc
  5 | # distutils: runtime_library_dirs = /usr/local/lib
  6 | # distutils: extra_compile_args = -std=c++17 -mavx2 -ftree-vectorize
  7 | 
  8 | cimport cython
  9 | from cpython cimport Py_INCREF, PyObject
 10 | from cpython.bytearray cimport PyByteArray_AS_STRING, PyByteArray_Check
 11 | from cpython.bytes cimport PyBytes_AS_STRING, PyBytes_Check
 12 | from cpython.memoryview cimport PyMemoryView_Check, PyMemoryView_GET_BUFFER
 13 | from cpython.unicode cimport PyUnicode_Check
 14 | from cython.operator cimport dereference
 15 | from libc.stdint cimport int32_t, int64_t
 16 | from libc.string cimport memcpy, memset
 17 | from numpy cimport (
 18 |     NPY_ARRAY_C_CONTIGUOUS,
 19 |     NPY_OBJECT,
 20 |     PyArray_CheckExact,
 21 |     PyArray_DATA,
 22 |     PyArray_Descr,
 23 |     PyArray_DescrFromType,
 24 |     PyArray_DIM,
 25 |     PyArray_GETCONTIGUOUS,
 26 |     PyArray_ISOBJECT,
 27 |     PyArray_ISSTRING,
 28 |     PyArray_NDIM,
 29 |     PyArray_SetBaseObject,
 30 |     dtype as npdtype,
 31 |     import_array,
 32 |     ndarray,
 33 |     npy_bool,
 34 |     npy_intp,
 35 | )
 36 | 
 37 | from athenian.api.native.chunked_stream cimport chunked_stream
 38 | from athenian.api.native.cpython cimport (
 39 |     Py_None,
 40 |     Py_True,
 41 |     PyBytes_GET_SIZE,
 42 |     PyList_CheckExact,
 43 |     PyList_GET_ITEM,
 44 |     PyList_GET_SIZE,
 45 |     PyTuple_GET_ITEM,
 46 |     PyTypeObject,
 47 |     PyUnicode_DATA,
 48 |     PyUnicode_GET_LENGTH,
 49 |     PyUnicode_KIND,
 50 | )
 51 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport mi_heap_destroy_stl_allocator
 52 | from athenian.api.native.numpy cimport (
 53 |     PyArray_DESCR,
 54 |     PyArray_DescrNew,
 55 |     PyArray_NewFromDescr,
 56 |     PyArray_Type,
 57 | )
 58 | from athenian.api.native.optional cimport optional
 59 | 
 60 | import asyncpg
 61 | import numpy as np
 62 | 
 63 | import_array()
 64 | 
 65 | 
 66 | cdef extern from "asyncpg_recordobj.h":
 67 |     PyObject *ApgRecord_GET_ITEM(PyObject *, int)
 68 | 
 69 | 
 70 | @cython.boundscheck(False)
 71 | def to_object_arrays(list rows not None, int columns) -> np.ndarray:
 72 |     """
 73 |     Convert a list of tuples or asyncpg.Record-s into an object array. Any subclass of
 74 |     tuple in `rows` will be casted to tuple.
 75 | 
 76 |     Parameters
 77 |     ----------
 78 |     rows: 2-d array (N, K)
 79 |         list of tuples to be converted into an array. Each tuple must be of equal length,
 80 |         otherwise, the results are undefined.
 81 |     columns: number of columns in each row.
 82 | 
 83 |     Returns
 84 |     -------
 85 |     np.ndarray[object, ndim=2]
 86 |     """
 87 |     cdef:
 88 |         Py_ssize_t i, j, size
 89 |         ndarray[object, ndim=2] result
 90 |         PyObject *record
 91 | 
 92 |     size = len(rows)
 93 | 
 94 |     result = np.empty((columns, size), dtype=object)
 95 |     if size == 0:
 96 |         return result
 97 | 
 98 |     if isinstance(rows[0], asyncpg.Record):
 99 |         for i in range(size):
100 |             record = PyList_GET_ITEM(<PyObject *>rows, i)
101 |             for j in range(columns):
102 |                 result[j, i] = <object>ApgRecord_GET_ITEM(record, j)
103 |     elif isinstance(rows[0], tuple):
104 |         for i in range(size):
105 |             record = PyList_GET_ITEM(<PyObject *>rows, i)
106 |             for j in range(columns):
107 |                 result[j, i] = <object>PyTuple_GET_ITEM(record, j)
108 |     else:
109 |         # convert to tuple
110 |         for i in range(size):
111 |             row = tuple(rows[i])
112 |             for j in range(columns):
113 |                 result[j, i] = row[j]
114 | 
115 |     return result
116 | 
117 | 
118 | def as_bool(ndarray arr not None) -> np.ndarray:
119 |     if arr.dtype == bool:
120 |         return arr
121 |     assert arr.dtype == object
122 |     assert arr.ndim == 1
123 |     new_arr = np.empty(len(arr), dtype=bool)
124 |     cdef:
125 |         const char *arr_obj = <const char *> PyArray_DATA(arr)
126 |         long size = len(arr), stride = arr.strides[0]
127 |         npy_bool *out_bools = <npy_bool *> PyArray_DATA(new_arr)
128 |     with nogil:
129 |         _as_bool_vec(arr_obj, stride, size, out_bools)
130 |     return new_arr
131 | 
132 | 
133 | @cython.boundscheck(False)
134 | @cython.wraparound(False)
135 | cdef void _as_bool_vec(const char *obj_arr,
136 |                        const long stride,
137 |                        const long size,
138 |                        npy_bool *out_arr) nogil:
139 |     cdef long i
140 |     for i in range(size):
141 |         # Py_None and Py_False become 0
142 |         out_arr[i] = Py_True == (<const PyObject **> (obj_arr + i * stride))[0]
143 | 
144 | 
145 | 
146 | def is_null(ndarray arr not None) -> np.ndarray:
147 |     if arr.dtype != object:
148 |         return np.zeros(len(arr), dtype=bool)
149 |     assert arr.ndim == 1
150 |     new_arr = np.zeros(len(arr), dtype=bool)
151 |     cdef:
152 |         const char *arr_obj = <const char *> PyArray_DATA(arr)
153 |         long size = len(arr), stride = arr.strides[0]
154 |         npy_bool *out_bools = <npy_bool *> PyArray_DATA(new_arr)
155 |     with nogil:
156 |         _is_null_vec(arr_obj, stride, size, out_bools)
157 |     return new_arr
158 | 
159 | 
160 | @cython.boundscheck(False)
161 | @cython.wraparound(False)
162 | cdef void _is_null_vec(const char *obj_arr,
163 |                        const long stride,
164 |                        const long size,
165 |                        npy_bool *out_arr) nogil:
166 |     cdef long i
167 |     for i in range(size):
168 |         out_arr[i] = Py_None == (<const PyObject **> (obj_arr + i * stride))[0]
169 | 
170 | 
171 | def is_not_null(ndarray arr not None) -> np.ndarray:
172 |     if arr.dtype != object:
173 |         return np.ones(len(arr), dtype=bool)
174 |     assert arr.ndim == 1
175 |     new_arr = np.zeros(len(arr), dtype=bool)
176 |     cdef:
177 |         const char *arr_obj = <const char *> PyArray_DATA(arr)
178 |         long size = len(arr), stride = arr.strides[0]
179 |         npy_bool *out_bools = <npy_bool *> PyArray_DATA(new_arr)
180 |     with nogil:
181 |         _is_not_null(arr_obj, stride, size, out_bools)
182 |     return new_arr
183 | 
184 | 
185 | @cython.boundscheck(False)
186 | @cython.wraparound(False)
187 | cdef void _is_not_null(const char *obj_arr,
188 |                        const long stride,
189 |                        const long size,
190 |                        npy_bool *out_arr) nogil:
191 |     cdef long i
192 |     for i in range(size):
193 |         out_arr[i] = Py_None != (<const PyObject **> (obj_arr + i * stride))[0]
194 | 
195 | 
196 | def nested_lengths(arr not None, output=None) -> np.ndarray:
197 |     cdef:
198 |         long size
199 |         bint is_array = PyArray_CheckExact(arr)
200 |         ndarray result
201 | 
202 |     if is_array:
203 |         assert PyArray_ISOBJECT(arr) or PyArray_ISSTRING(arr)
204 |         assert PyArray_NDIM(arr) == 1
205 |         size = PyArray_DIM(arr, 0)
206 |     else:
207 |         assert PyList_CheckExact(<PyObject *> arr)
208 |         size = PyList_GET_SIZE(<PyObject *> arr)
209 | 
210 |     if output is None:
211 |         result = np.zeros(size, dtype=int)
212 |     else:
213 |         assert PyArray_CheckExact(output)
214 |         assert output.dtype == int
215 |         result = output
216 |     if size == 0:
217 |         return result
218 |     if is_array:
219 |         return _nested_lengths_arr(arr, size, result)
220 |     return _nested_lengths_list(<PyObject *> arr, size, result)
221 | 
222 | 
223 | cdef ndarray _nested_lengths_arr(ndarray arr, long size, ndarray result):
224 |     cdef:
225 |         PyObject **elements = <PyObject **>PyArray_DATA(arr)
226 |         PyObject *element
227 |         long i
228 |         long *result_data
229 | 
230 |     if PyArray_ISSTRING(arr):
231 |         return np.char.str_len(arr)
232 | 
233 |     result_data = <long *>PyArray_DATA(result)
234 |     element = elements[0]
235 |     if PyArray_CheckExact(<object> element):
236 |         for i in range(size):
237 |             result_data[i] = PyArray_DIM(<ndarray> elements[i], 0)
238 |     elif PyList_CheckExact(element):
239 |         for i in range(size):
240 |             result_data[i] = PyList_GET_SIZE(elements[i])
241 |     elif PyUnicode_Check(<object> element):
242 |         for i in range(size):
243 |             result_data[i] = PyUnicode_GET_LENGTH(elements[i])
244 |     elif PyBytes_Check(<object> element):
245 |         for i in range(size):
246 |             result_data[i] = PyBytes_GET_SIZE(elements[i])
247 |     else:
248 |         raise AssertionError(f"Unsupported nested type: {type(<object> element).__name__}")
249 |     return result
250 | 
251 | 
252 | cdef ndarray _nested_lengths_list(PyObject *arr, long size, ndarray result):
253 |     cdef:
254 |         PyObject *element
255 |         long i
256 |         long *result_data
257 | 
258 |     result_data = <long *>PyArray_DATA(result)
259 |     element = PyList_GET_ITEM(arr, 0)
260 |     if PyArray_CheckExact(<object> element):
261 |         for i in range(size):
262 |             result_data[i] = PyArray_DIM(<ndarray> PyList_GET_ITEM(arr, i), 0)
263 |     elif PyList_CheckExact(element):
264 |         for i in range(size):
265 |             result_data[i] = PyList_GET_SIZE(PyList_GET_ITEM(arr, i))
266 |     elif PyUnicode_Check(<object> element):
267 |         for i in range(size):
268 |             result_data[i] = PyUnicode_GET_LENGTH(PyList_GET_ITEM(arr, i))
269 |     elif PyBytes_Check(<object> element):
270 |         for i in range(size):
271 |             result_data[i] = PyBytes_GET_SIZE(PyList_GET_ITEM(arr, i))
272 |     else:
273 |         raise AssertionError(f"Unsupported nested type: {type(<object> element).__name__}")
274 |     return result
275 | 
276 | 
277 | def array_from_buffer(buffer not None, npdtype dtype, npy_intp count, npy_intp offset=0) -> ndarray:
278 |     cdef:
279 |         void *data
280 |         ndarray arr
281 |     if PyBytes_Check(buffer):
282 |         data = PyBytes_AS_STRING(buffer) + offset
283 |     elif PyByteArray_Check(buffer):
284 |         data = PyByteArray_AS_STRING(buffer) + offset
285 |     elif PyMemoryView_Check(buffer):
286 |         data = (<char *>PyMemoryView_GET_BUFFER(buffer).buf) + offset
287 |     else:
288 |         raise ValueError(f"Unsupported buffer type: {type(buffer).__name__}")
289 |     Py_INCREF(dtype)
290 |     Py_INCREF(buffer)
291 |     arr = <ndarray> PyArray_NewFromDescr(
292 |         &PyArray_Type,
293 |         <PyArray_Descr *> dtype,
294 |         1,
295 |         &count,
296 |         NULL,
297 |         data,
298 |         NPY_ARRAY_C_CONTIGUOUS,
299 |         NULL,
300 |     )
301 |     PyArray_SetBaseObject(arr, buffer)
302 |     return arr
303 | 
304 | 
305 | def array_of_objects(int length, fill_value) -> ndarray:
306 |     cdef:
307 |         ndarray arr
308 |         npdtype objdtype = PyArray_DescrNew(PyArray_DescrFromType(NPY_OBJECT))
309 |         npy_intp nplength = length, i
310 |         PyObject **data
311 |         PyObject *obj = <PyObject *> fill_value
312 | 
313 |     arr = <ndarray> PyArray_NewFromDescr(
314 |         &PyArray_Type,
315 |         <PyArray_Descr *> objdtype,
316 |         1,
317 |         &nplength,
318 |         NULL,
319 |         NULL,
320 |         NPY_ARRAY_C_CONTIGUOUS,
321 |         NULL,
322 |     )
323 |     Py_INCREF(objdtype)
324 |     data = <PyObject **> PyArray_DATA(arr)
325 |     for i in range(nplength):
326 |         data[i] = obj
327 |     obj.ob_refcnt += nplength
328 |     return arr
329 | 
330 | 
331 | def vectorize_numpy_struct_scalar_field(cls, structs, npdtype dtype, long offset) -> np.ndarray:
332 |     cdef:
333 |         ndarray arr
334 |         npy_intp length = len(structs), i = 0
335 |         npy_intp itemsize = dtype.itemsize
336 |         char *arr_data
337 |         Py_ssize_t struct_data_offset = (<PyTypeObject *> cls).tp_members[1].offset
338 |         PyObject *struct_data_obj
339 |         char *struct_data
340 |     arr = <ndarray> PyArray_NewFromDescr(
341 |         &PyArray_Type,
342 |         <PyArray_Descr *> dtype,
343 |         1,
344 |         &length,
345 |         NULL,
346 |         NULL,
347 |         NPY_ARRAY_C_CONTIGUOUS,
348 |         NULL,
349 |     )
350 |     Py_INCREF(dtype)
351 |     arr_data = <char *> PyArray_DATA(arr)
352 | 
353 |     for struct in structs:
354 |         struct_data_obj = dereference(
355 |             <PyObject **> ((<char *> <PyObject *> struct) + struct_data_offset)
356 |         )
357 |         if PyBytes_Check(<object> struct_data_obj):
358 |             struct_data = PyBytes_AS_STRING(<object> struct_data_obj)
359 |         elif PyMemoryView_Check(<object> struct_data_obj):
360 |             struct_data = <char *>PyMemoryView_GET_BUFFER(<object> struct_data_obj).buf
361 |         elif PyByteArray_Check(<object> struct_data_obj):
362 |             struct_data = PyByteArray_AS_STRING(<object> struct_data_obj)
363 |         else:
364 |             raise AssertionError(f"unsupported buffer type: {type(struct.data)}")
365 |         memcpy(arr_data + i * itemsize, struct_data + offset, itemsize)
366 |         i += 1
367 |     return arr
368 | 
369 | 
370 | def vectorize_numpy_struct_array_field(
371 |     cls,
372 |     structs,
373 |     npdtype dtype,
374 |     long offset,
375 | ) -> tuple[np.ndarray, np.ndarray]:
376 |     cdef:
377 |         ndarray arr, arr_offsets
378 |         npy_intp length = len(structs) + 1, i = 0
379 |         npy_intp itemsize = dtype.itemsize
380 |         char *arr_data
381 |         Py_ssize_t struct_data_offset = (<PyTypeObject *> cls).tp_members[1].offset
382 |         PyObject *struct_data_obj
383 |         char *struct_data
384 |         int32_t field_offset, field_count
385 |         int32_t *field
386 |         npdtype int_dtype = npdtype(int)
387 |         int64_t pos = 0, delta
388 |         optional[chunked_stream] dump
389 |         optional[mi_heap_destroy_stl_allocator[char]] alloc
390 | 
391 |     alloc.emplace()
392 |     dump.emplace(dereference(alloc))
393 |     arr_offsets = <ndarray> PyArray_NewFromDescr(
394 |         &PyArray_Type,
395 |         <PyArray_Descr *> int_dtype,
396 |         1,
397 |         &length,
398 |         NULL,
399 |         NULL,
400 |         NPY_ARRAY_C_CONTIGUOUS,
401 |         NULL,
402 |     )
403 |     Py_INCREF(int_dtype)
404 |     offsets_data = <int64_t *> PyArray_DATA(arr_offsets)
405 | 
406 |     length = 0
407 |     for struct in structs:
408 |         struct_data_obj = dereference(
409 |             <PyObject **> ((<char *> <PyObject *> struct) + struct_data_offset)
410 |         )
411 |         if PyBytes_Check(<object> struct_data_obj):
412 |             struct_data = PyBytes_AS_STRING(<object> struct_data_obj)
413 |         elif PyMemoryView_Check(<object> struct_data_obj):
414 |             struct_data = <char *>PyMemoryView_GET_BUFFER(<object> struct_data_obj).buf
415 |         elif PyByteArray_Check(<object> struct_data_obj):
416 |             struct_data = PyByteArray_AS_STRING(<object> struct_data_obj)
417 |         else:
418 |             raise AssertionError(f"unsupported buffer type: {type(struct.data)}")
419 |         field = <int32_t *>(struct_data + offset)
420 |         field_offset = field[0]
421 |         field_count = field[1]
422 |         offsets_data[i] = length
423 |         length += field_count
424 |         delta = itemsize * field_count
425 |         dereference(dump).write(struct_data + field_offset, delta)
426 |         pos += delta
427 |         i += 1
428 |     offsets_data[i] = length
429 | 
430 |     arr = <ndarray> PyArray_NewFromDescr(
431 |         &PyArray_Type,
432 |         <PyArray_Descr *> dtype,
433 |         1,
434 |         &length,
435 |         NULL,
436 |         NULL,
437 |         NPY_ARRAY_C_CONTIGUOUS,
438 |         NULL,
439 |     )
440 |     Py_INCREF(dtype)
441 |     dereference(dump).dump(<char *>PyArray_DATA(arr), pos)
442 | 
443 |     return arr, arr_offsets
444 | 
445 | 
446 | def objects_to_pyunicode_bytes(ndarray arr not None, char_limit=None) -> ndarray:
447 |     assert PyArray_NDIM(arr) == 1
448 |     assert PyArray_DESCR(<PyObject *> arr).kind == b"O"
449 | 
450 |     cdef npy_intp length = PyArray_DIM(arr, 0)
451 | 
452 |     if length == 0:
453 |         return np.array([], dtype="S")
454 | 
455 |     arr = PyArray_GETCONTIGUOUS(arr)
456 | 
457 |     cdef:
458 |         npy_intp i, max_itemsize = 0
459 |         PyObject **data = <PyObject **> PyArray_DATA(arr)
460 |         PyObject *obj
461 |         Py_ssize_t i_itemsize, char_limit_native
462 |         npdtype dtype
463 |         ndarray converted
464 |         char *converted_data
465 |         char *head
466 | 
467 |     if char_limit is None:
468 |         for i in range(length):
469 |             obj = data[i]
470 |             if obj == Py_None:
471 |                 i_itemsize = 4
472 |             else:
473 |                 assert PyUnicode_Check(<object> obj), f"arr[{i}]: {arr[i]}"
474 |                 i_itemsize = PyUnicode_GET_LENGTH(obj) * PyUnicode_KIND(obj)
475 |             if i_itemsize > max_itemsize:
476 |                 max_itemsize = i_itemsize
477 |     else:
478 |         assert char_limit > 0
479 |         char_limit_native = char_limit
480 |         for i in range(length):
481 |             obj = data[i]
482 |             if obj == Py_None:
483 |                 i_itemsize = 4
484 |             else:
485 |                 assert PyUnicode_Check(<object> obj), f"arr[{i}]: {arr[i]}"
486 |                 i_itemsize = PyUnicode_GET_LENGTH(obj) * PyUnicode_KIND(obj)
487 |             if i_itemsize >= char_limit_native:
488 |                 max_itemsize = char_limit_native
489 |                 break
490 |             if i_itemsize > max_itemsize:
491 |                 max_itemsize = i_itemsize
492 | 
493 |     dtype = npdtype("S" + str(max_itemsize))
494 |     converted = <ndarray> PyArray_NewFromDescr(
495 |         &PyArray_Type,
496 |         <PyArray_Descr *> dtype,
497 |         1,
498 |         &length,
499 |         NULL,
500 |         NULL,
501 |         NPY_ARRAY_C_CONTIGUOUS,
502 |         NULL,
503 |     )
504 |     Py_INCREF(dtype)
505 |     converted_data = <char *> PyArray_DATA(converted)
506 | 
507 |     for i in range(length):
508 |         obj = data[i]
509 |         head = converted_data + i * max_itemsize
510 | 
511 |         if obj == Py_None:
512 |             i_itemsize = 4
513 |             if i_itemsize > max_itemsize:
514 |                 i_itemsize = max_itemsize
515 |             memcpy(head, b"None", i_itemsize)
516 |         else:
517 |             i_itemsize = PyUnicode_GET_LENGTH(obj) * PyUnicode_KIND(obj)
518 |             if i_itemsize > max_itemsize:
519 |                 i_itemsize = max_itemsize
520 |             memcpy(head, PyUnicode_DATA(obj), i_itemsize)
521 |         memset(head + i_itemsize, 0, max_itemsize - i_itemsize)
522 | 
523 |     return converted
524 | 


--------------------------------------------------------------------------------
/2023/pyx/sentry_native.pyx:
--------------------------------------------------------------------------------
 1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
 2 | # cython: warn.maybe_uninitialized=True
 3 | # distutils: language = c++
 4 | # distutils: extra_compile_args = -std=c++17
 5 | # distutils: libraries = sentry
 6 | # distutils: runtime_library_dirs = /usr/local/lib
 7 | 
 8 | 
 9 | cdef extern from "sentry.h" nogil:
10 |     struct sentry_options_s
11 |     ctypedef sentry_options_s sentry_options_t
12 |     sentry_options_t *sentry_options_new()
13 |     void sentry_options_set_dsn(sentry_options_t *opts, const char *dsn)
14 |     void sentry_options_set_release(sentry_options_t *opts, const char *release)
15 |     void sentry_options_set_environment(sentry_options_t *opts, const char *environment)
16 |     void sentry_options_set_debug(sentry_options_t *opts, int debug)
17 |     void sentry_options_set_max_breadcrumbs(sentry_options_t *opts, size_t max_breadcrumbs)
18 |     void sentry_options_set_handler_path(sentry_options_t *opts, const char *path)
19 |     void sentry_options_set_symbolize_stacktraces(sentry_options_t *opts, int val)
20 |     int sentry_init(sentry_options_t *options)
21 |     int sentry_close()
22 | 
23 | 
24 | def init(str dsn not None, str release not None, str env not None) -> None:
25 |     cdef:
26 |         sentry_options_t *options = sentry_options_new()
27 |     sentry_options_set_dsn(options, dsn.encode())
28 |     sentry_options_set_release(options, release.encode())
29 |     sentry_options_set_environment(options, env.encode())
30 |     sentry_options_set_debug(options, env != "production")
31 |     sentry_options_set_handler_path(options, b"/usr/local/bin/crashpad_handler")
32 |     sentry_options_set_symbolize_stacktraces(options, 1)
33 |     sentry_options_set_max_breadcrumbs(options, 20)
34 |     sentry_init(options)
35 | 
36 | 
37 | def fini():
38 |     sentry_close()
39 | 


--------------------------------------------------------------------------------
/2023/pyx/sorted_ops.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: extra_compile_args = -mavx2 -ftree-vectorize -std=c++17
  5 | 
  6 | from cpython cimport Py_INCREF, PyObject
  7 | from libc.stdint cimport int64_t, uint32_t
  8 | from numpy cimport (
  9 |     NPY_ARRAY_C_CONTIGUOUS,
 10 |     NPY_INT64,
 11 |     NPY_UINT,
 12 |     PyArray_DATA,
 13 |     PyArray_DescrFromType,
 14 |     PyArray_DIM,
 15 |     PyArray_IS_C_CONTIGUOUS,
 16 |     PyArray_NDIM,
 17 |     dtype,
 18 |     import_array,
 19 |     ndarray,
 20 |     npy_intp,
 21 | )
 22 | 
 23 | from athenian.api.native.cpython cimport PyUnicode_DATA
 24 | from athenian.api.native.numpy cimport (
 25 |     PyArray_Descr,
 26 |     PyArray_DESCR,
 27 |     PyArray_DescrNew,
 28 |     PyArray_NewFromDescr,
 29 |     PyArray_Type,
 30 | )
 31 | 
 32 | 
 33 | cdef extern from "native/sorted_intersection.h" nogil:
 34 |     size_t intersect(
 35 |         const char *algorithm,
 36 |         const uint32_t *set1,
 37 |         const size_t length1,
 38 |         const uint32_t *set2,
 39 |         const size_t length2,
 40 |         uint32_t *out
 41 |     )
 42 | 
 43 | import_array()
 44 | 
 45 | """
 46 | Benchmark results on production metrics:
 47 | 
 48 | v1
 49 | 4.79 ms ± 180 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 50 | v3
 51 | 6.94 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 52 | simd
 53 | 4.92 ms ± 293 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 54 | galloping
 55 | 2.78 ms ± 30.5 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 56 | mut_part
 57 | 3.35 ms ± 294 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 58 | scalar
 59 | 3.21 ms ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 60 | simdgalloping
 61 | 7.08 ms ± 93.6 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 62 | simd_avx2
 63 | 4.6 ms ± 336 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 64 | v1_avx2
 65 | 4.22 ms ± 143 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 66 | v3_avx2
 67 | 10.5 ms ± 152 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 68 | simdgalloping_avx2
 69 | 10.7 ms ± 193 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 70 | highlyscalable_intersect_SIMD
 71 | 4.78 ms ± 340 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 72 | lemire_highlyscalable_intersect_SIMD
 73 | 3.97 ms ± 387 µs per loop (mean ± std. dev. of 7 runs, 200 loops each)
 74 | """
 75 | 
 76 | def sorted_intersect1d(
 77 |     ndarray arr1 not None,
 78 |     ndarray arr2 not None,
 79 |     str algo="galloping",
 80 | ) -> ndarray:
 81 |     assert PyArray_NDIM(arr1) == 1
 82 |     assert PyArray_IS_C_CONTIGUOUS(arr1)
 83 |     assert PyArray_DESCR(<PyObject *> arr1).kind == b"u"
 84 |     assert PyArray_NDIM(arr2) == 1
 85 |     assert PyArray_IS_C_CONTIGUOUS(arr2)
 86 |     assert PyArray_DESCR(<PyObject *> arr2).kind == b"u"
 87 | 
 88 |     cdef:
 89 |         uint32_t *arr1_data = <uint32_t *> PyArray_DATA(arr1)
 90 |         uint32_t *arr2_data = <uint32_t *> PyArray_DATA(arr2)
 91 |         npy_intp len1 = PyArray_DIM(arr1, 0)
 92 |         npy_intp len2 = PyArray_DIM(arr2, 0)
 93 |         ndarray output
 94 |         dtype u32dtype = PyArray_DescrNew(PyArray_DescrFromType(NPY_UINT))
 95 | 
 96 |     output = <ndarray> PyArray_NewFromDescr(
 97 |         &PyArray_Type,
 98 |         <PyArray_Descr *> u32dtype,
 99 |         1,
100 |         &len1 if len1 > len2 else &len2,
101 |         NULL,
102 |         NULL,
103 |         NPY_ARRAY_C_CONTIGUOUS,
104 |         NULL,
105 |     )
106 |     Py_INCREF(u32dtype)
107 |     output.shape[0] = intersect(
108 |         <char *>PyUnicode_DATA(<PyObject *> algo),
109 |         arr1_data,
110 |         len1,
111 |         arr2_data,
112 |         len2,
113 |         <uint32_t *> PyArray_DATA(output),
114 |     )
115 |     return output
116 | 
117 | 
118 | cdef extern from "native/sum_repeated_with_step.h" nogil:
119 |     void sum_repeated_with_step_avx2(
120 |         const int64_t *src,
121 |         int64_t src_len,
122 |         int64_t repeats,
123 |         int64_t step,
124 |         int64_t *dst
125 |     )
126 | 
127 | 
128 | def sum_repeated_with_step(ndarray arr not None, long repeats, long step) -> ndarray:
129 |     """
130 |     Calculate fused sum of wide range with repeated vector.
131 | 
132 |         np.repeat(arr[None, :], repeats, axis=0).ravel() + np.repeat(
133 |             np.arange(repeats, dtype=int) * step, len(arr),
134 |         )
135 |     """
136 |     assert PyArray_NDIM(arr) == 1
137 |     assert PyArray_IS_C_CONTIGUOUS(arr)
138 | 
139 |     cdef:
140 |         dtype arr_dtype = arr.dtype
141 |         int64_t src_len = PyArray_DIM(arr, 0)
142 |         npy_intp dst_len = src_len * repeats
143 |         int64_t *src_data = <int64_t *> PyArray_DATA(arr)
144 | 
145 |     assert arr_dtype.kind == b"i"
146 |     output = <ndarray> PyArray_NewFromDescr(
147 |         &PyArray_Type,
148 |         <PyArray_Descr *> arr_dtype,
149 |         1,
150 |         &dst_len,
151 |         NULL,
152 |         NULL,
153 |         NPY_ARRAY_C_CONTIGUOUS,
154 |         NULL,
155 |     )
156 |     Py_INCREF(arr_dtype)
157 |     sum_repeated_with_step_avx2(src_data, src_len, repeats, step, <int64_t *> PyArray_DATA(output))
158 |     return output
159 | 
160 | 
161 | def sorted_union1d(ndarray arr1 not None, ndarray arr2 not None) -> ndarray:
162 |     assert PyArray_NDIM(arr1) == 1
163 |     assert PyArray_NDIM(arr2) == 1
164 | 
165 |     if PyArray_DIM(arr1, 0) == 0:
166 |         return arr2
167 |     elif PyArray_DIM(arr2, 0) == 0:
168 |         return arr1
169 | 
170 |     cdef:
171 |         char kind1 = PyArray_DESCR(<PyObject *> arr1).kind
172 |         char kind2 = PyArray_DESCR(<PyObject *> arr2).kind
173 | 
174 |     assert PyArray_IS_C_CONTIGUOUS(arr1)
175 |     assert PyArray_IS_C_CONTIGUOUS(arr2)
176 |     assert kind1 == kind2, "dtypes must match"
177 | 
178 |     if kind1 == b"i":
179 |         return _sorted_union1d_i64(arr1, arr2)
180 |     elif kind1 == b"u":
181 |         return _sorted_union1d_u32(arr1, arr2)
182 | 
183 |     raise AssertionError(f"dtype is not supported: {arr1.dtype}")
184 | 
185 | 
186 | cdef ndarray _sorted_union1d_i64(ndarray arr1, ndarray arr2):
187 |     cdef:
188 |         npy_intp len1 = PyArray_DIM(arr1, 0)
189 |         npy_intp len2 = PyArray_DIM(arr2, 0)
190 |         int64_t *arr1_data = <int64_t *> PyArray_DATA(arr1)
191 |         int64_t *arr2_data = <int64_t *> PyArray_DATA(arr2)
192 |         npy_intp out_len = len1 + len2, pos1 = 0, pos2 = 0
193 |         int64_t head1, head2
194 |         ndarray output
195 |         dtype i64dtype = PyArray_DescrNew(PyArray_DescrFromType(NPY_INT64))
196 |         int64_t *output_data
197 | 
198 |     output = <ndarray> PyArray_NewFromDescr(
199 |         &PyArray_Type,
200 |         <PyArray_Descr *> i64dtype,
201 |         1,
202 |         &out_len,
203 |         NULL,
204 |         NULL,
205 |         NPY_ARRAY_C_CONTIGUOUS,
206 |         NULL,
207 |     )
208 |     Py_INCREF(i64dtype)
209 |     output_data = <int64_t *> PyArray_DATA(output)
210 | 
211 |     with nogil:
212 |         head1 = arr1_data[0]
213 |         head2 = arr2_data[0]
214 |         out_len = 0
215 |         while pos1 < len1 and pos2 < len2:
216 |             if head1 < head2:
217 |                 output_data[out_len] = head1
218 |                 pos1 += 1
219 |                 head1 = arr1_data[pos1]
220 |             elif head1 > head2:
221 |                 output_data[out_len] = head2
222 |                 pos2 += 1
223 |                 head2 = arr2_data[pos2]
224 |             else:
225 |                 output_data[out_len] = head1
226 |                 pos1 += 1
227 |                 head1 = arr1_data[pos1]
228 |                 pos2 += 1
229 |                 head2 = arr2_data[pos2]
230 |             out_len += 1
231 |         if pos1 == len1:
232 |             while pos2 < len2:
233 |                 output_data[out_len] = arr2_data[pos2]
234 |                 pos2 += 1
235 |                 out_len += 1
236 |         else:
237 |             while pos1 < len1:
238 |                 output_data[out_len] = arr1_data[pos1]
239 |                 pos1 += 1
240 |                 out_len += 1
241 | 
242 |     return output[:out_len]
243 | 
244 | 
245 | cdef ndarray _sorted_union1d_u32(ndarray arr1, ndarray arr2):
246 |     cdef:
247 |         npy_intp len1 = PyArray_DIM(arr1, 0)
248 |         npy_intp len2 = PyArray_DIM(arr2, 0)
249 |         uint32_t *arr1_data = <uint32_t *> PyArray_DATA(arr1)
250 |         uint32_t *arr2_data = <uint32_t *> PyArray_DATA(arr2)
251 |         npy_intp out_len = len1 + len2, pos1 = 0, pos2 = 0
252 |         uint32_t head1, head2
253 |         ndarray output
254 |         dtype u32dtype = PyArray_DescrNew(PyArray_DescrFromType(NPY_UINT))
255 |         uint32_t *output_data
256 | 
257 |     output = <ndarray> PyArray_NewFromDescr(
258 |         &PyArray_Type,
259 |         <PyArray_Descr *> u32dtype,
260 |         1,
261 |         &out_len,
262 |         NULL,
263 |         NULL,
264 |         NPY_ARRAY_C_CONTIGUOUS,
265 |         NULL,
266 |     )
267 |     Py_INCREF(u32dtype)
268 |     output_data = <uint32_t *> PyArray_DATA(output)
269 | 
270 |     with nogil:
271 |         head1 = arr1_data[0]
272 |         head2 = arr2_data[0]
273 |         out_len = 0
274 |         while pos1 < len1 and pos2 < len2:
275 |             if head1 < head2:
276 |                 output_data[out_len] = head1
277 |                 pos1 += 1
278 |                 head1 = arr1_data[pos1]
279 |             elif head1 > head2:
280 |                 output_data[out_len] = head2
281 |                 pos2 += 1
282 |                 head2 = arr2_data[pos2]
283 |             else:
284 |                 output_data[out_len] = head1
285 |                 pos1 += 1
286 |                 head1 = arr1_data[pos1]
287 |                 pos2 += 1
288 |                 head2 = arr2_data[pos2]
289 |             out_len += 1
290 |         if pos1 == len1:
291 |             while pos2 < len2:
292 |                 output_data[out_len] = arr2_data[pos2]
293 |                 pos2 += 1
294 |                 out_len += 1
295 |         else:
296 |             while pos1 < len1:
297 |                 output_data[out_len] = arr1_data[pos1]
298 |                 pos1 += 1
299 |                 out_len += 1
300 | 
301 |     return output[:out_len]
302 | 


--------------------------------------------------------------------------------
/2023/pyx/sql_builders.h:
--------------------------------------------------------------------------------
 1 | #include <immintrin.h>
 2 | 
 3 | static inline int scan_unicode_kind(const char *data, long length) {
 4 | 	__m256i word = _mm256_set1_epi8(0);
 5 | 	long i;
 6 | 	for (i = 0; i < length - 31; i += 32) {
 7 | 		word = _mm256_or_si256(word, _mm256_loadu_si256((const __m256i *)(data + i)));
 8 | 	}
 9 | 	word = _mm256_cmpeq_epi8(word, _mm256_set1_epi8(0));
10 | 	uint32_t mask = _mm256_movemask_epi8(word);
11 | 	mask = ~mask;
12 | 	for (; i < length; i += 4) {
13 | 	    mask |= data[i] != 0;
14 | 	    mask |= (data[i + 1] != 0) << 1;
15 | 	    mask |= (data[i + 2] != 0) << 2;
16 | 	    mask |= (data[i + 3] != 0) << 3;
17 | 	}
18 | 	if (mask & ((1 << 3) | (1 << 7) | (1 << 11) | (1 << 15) | (1 << 19) | (1 << 23) | (1 << 27) | (1 << 31))) {
19 | 	    return 4;
20 | 	}
21 | 	if (mask & ((1 << 2) | (1 << 6) | (1 << 10) | (1 << 14) | (1 << 18) | (1 << 22) | (1 << 26) | (1 << 30))) {
22 | 	    return 4;
23 | 	}
24 | 	if (mask & ((1 << 1) | (1 << 5) | (1 << 9) | (1 << 13) | (1 << 17) | (1 << 21) | (1 << 25) | (1 << 29))) {
25 | 	    return 2;
26 | 	}
27 | 	return 1;
28 | }


--------------------------------------------------------------------------------
/2023/pyx/sql_builders.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: extra_compile_args = -mavx2 -ftree-vectorize
  5 | 
  6 | from cpython cimport PyList_CheckExact, PyObject
  7 | 
  8 | import cython
  9 | 
 10 | from libc.stddef cimport wchar_t
 11 | from libc.stdint cimport int64_t
 12 | from libc.stdlib cimport lldiv, lldiv_t
 13 | from libc.string cimport memchr, memcpy
 14 | from numpy cimport (
 15 |     PyArray_CheckExact,
 16 |     PyArray_DATA,
 17 |     PyArray_DESCR,
 18 |     PyArray_DIM,
 19 |     PyArray_STRIDE,
 20 |     dtype as np_dtype,
 21 |     import_array,
 22 |     ndarray,
 23 | )
 24 | 
 25 | from athenian.api.native.cpython cimport (
 26 |     Py_None,
 27 |     PyBytes_AS_STRING,
 28 |     PyBytes_Check,
 29 |     PyBytes_GET_SIZE,
 30 |     PyList_GET_ITEM,
 31 |     PyLong_AsLong,
 32 |     PyLong_CheckExact,
 33 |     PyObject_TypeCheck,
 34 |     PyUnicode_Check,
 35 |     PyUnicode_DATA,
 36 |     PyUnicode_FindChar,
 37 |     PyUnicode_GET_LENGTH,
 38 |     PyUnicode_KIND,
 39 |     PyUnicode_New,
 40 | )
 41 | from athenian.api.native.numpy cimport PyArray_ScalarAsCtype, PyIntegerArrType_Type
 42 | 
 43 | import_array()
 44 | 
 45 | 
 46 | cdef extern from "wchar.h" nogil:
 47 |     wchar_t *wmemchr(const wchar_t *s, wchar_t c, size_t n)
 48 | 
 49 | 
 50 | cdef extern from "sql_builders.h" nogil:
 51 |     int scan_unicode_kind(const char *data, long length)
 52 | 
 53 | 
 54 | def in_any_values_inline(values) -> str:
 55 |     if PyArray_CheckExact(values):
 56 |         return in_any_values_inline_array(values)
 57 |     if PyList_CheckExact(values):
 58 |         return in_any_values_inline_list(values)
 59 |     raise ValueError(f"Only numpy arrays and lists are supported, got {type(values)}")
 60 | 
 61 | 
 62 | cdef str in_any_values_inline_array(ndarray values):
 63 |     if values.ndim != 1:
 64 |         raise ValueError(f"We support only 1-dimensional numpy arrays, got {values.ndim}")
 65 |     if values.dtype.kind not in ("S", "U", "i", "u"):
 66 |         raise ValueError(f"unsupported dtype {values.dtype}")
 67 |     if len(values) == 0:
 68 |         raise ValueError("= ANY(VALUES) is invalid syntax")
 69 |     cdef:
 70 |         np_dtype dtype = <np_dtype> PyArray_DESCR(values)
 71 |         int is_s = dtype.kind == b"S"
 72 |         int is_str = is_s or dtype.kind == b"U"
 73 |         int stride = PyArray_STRIDE(values, 0)
 74 |         int itemsize = dtype.itemsize
 75 |         int length = PyArray_DIM(values, 0)
 76 |         int effective_itemsize = \
 77 |             (itemsize if dtype.kind == b"S" else itemsize >> 2) + 2 \
 78 |             if is_str \
 79 |             else len(str(values.max()))
 80 |         Py_ssize_t size = (7 + (effective_itemsize + 3) * length - 1)
 81 |         result = PyUnicode_New(size, 255)
 82 |         char *buf = <char *> PyUnicode_DATA(<PyObject *> result)
 83 |         char *data = <char *> PyArray_DATA(values)
 84 | 
 85 |     if is_str and itemsize == stride:
 86 |         if is_s:
 87 |             if memchr(data, b"'", length * itemsize) != NULL:
 88 |                 raise NotImplementedError("One of the strings requires escaping the single quote")
 89 |         else:
 90 |             if wmemchr(<const wchar_t *> data, b"'", length * itemsize >> 2) != NULL:
 91 |                 raise NotImplementedError("One of the strings requires escaping the single quote")
 92 |     with nogil:
 93 |         if is_str:
 94 |             if is_s:
 95 |                 if not _in_any_values_array_s(data, stride, itemsize, length, buf):
 96 |                     raise NotImplementedError("One of the strings requires escaping the single quote")
 97 |             else:
 98 |                 if not _in_any_values_array_u(data, stride, itemsize, length, buf):
 99 |                     raise NotImplementedError("One of the strings requires escaping the single quote")
100 |         elif itemsize == 8:
101 |             _in_any_values_array_int64(data, stride, length, effective_itemsize, buf)
102 |         else:
103 |             raise ValueError(f"unsupported dtype {dtype}")
104 |     return result
105 | 
106 | 
107 | cdef bint _in_any_values_array_s(
108 |     const char *data,
109 |     int stride,
110 |     int itemsize,
111 |     int length,
112 |     char *output,
113 | ) nogil:
114 |     cdef:
115 |         int i, pos = 7
116 |         char *quoteptr
117 |         char *nullptr
118 |     memcpy(output, b"VALUES ", 7)
119 | 
120 |     for i in range(length):
121 |         output[pos] = b"("
122 |         pos += 1
123 |         output[pos] = b"'"
124 |         pos += 1
125 |         memcpy(output + pos, data + stride * i, itemsize)
126 |         if stride != itemsize and memchr(output + pos, b"'", itemsize) != NULL:
127 |             return False
128 |         pos += itemsize
129 |         output[pos] = b"'"
130 |         pos += 1
131 |         output[pos] = b")"
132 |         pos += 1
133 |         if i < length - 1:
134 |             output[pos] = b","
135 |             pos += 1
136 | 
137 |     nullptr = <char *> memchr(output, 0, pos)
138 |     while nullptr:
139 |         quoteptr = <char *> memchr(nullptr + 1, b"'", itemsize)
140 |         nullptr[0] = b"'"
141 |         for i in range(1, (quoteptr - nullptr) + 1):
142 |             nullptr[i] = b" "
143 |         nullptr = <char *> memchr(quoteptr, 0, pos - (quoteptr - output))
144 | 
145 |     return True
146 | 
147 | 
148 | cdef bint _in_any_values_array_u(
149 |     const char *data,
150 |     int stride,
151 |     int itemsize,
152 |     int length,
153 |     char *output,
154 | ) nogil:
155 |     cdef:
156 |         int i, j, fill, pos = 7
157 |         char c
158 |     memcpy(output, b"VALUES ", 7)
159 | 
160 |     for i in range(length):
161 |         output[pos] = b"("
162 |         pos += 1
163 |         output[pos] = b"'"
164 |         pos += 1
165 |         fill = False
166 |         for j in range(0, itemsize, 4):
167 |             c = data[stride * i + j]
168 |             if stride != itemsize and c == b"'":
169 |                 return False
170 |             if fill:
171 |                 c = b" "
172 |             elif c == 0:
173 |                 c = b"'"
174 |                 fill = True
175 |             output[pos] = c
176 |             pos += 1
177 |         output[pos] = b" " if fill else b"'"
178 |         pos += 1
179 |         output[pos] = b")"
180 |         pos += 1
181 |         if i < length - 1:
182 |             output[pos] = b","
183 |             pos += 1
184 | 
185 |     return True
186 | 
187 | 
188 | cdef void _in_any_values_array_int64(
189 |     const char *data,
190 |     int stride,
191 |     int length,
192 |     int alignment,
193 |     char *output,
194 | ) nogil:
195 |     cdef:
196 |         int i, pos = 7, valstart
197 |         lldiv_t qr
198 |     memcpy(output, b"VALUES ", 7)
199 | 
200 |     for i in range(length):
201 |         output[pos] = b"("
202 |         pos += 1
203 |         valstart = pos
204 |         pos += alignment - 1
205 |         qr.quot = (<const int64_t *>(data + i * stride))[0]
206 |         while True:
207 |             qr = lldiv(qr.quot, 10)
208 |             output[pos] = (<char>b'0') + (<char>qr.rem)
209 |             pos -= 1
210 |             if qr.quot == 0:
211 |                 break
212 |         while pos >= valstart:
213 |             output[pos] = b" "
214 |             pos -= 1
215 |         pos = valstart + alignment
216 |         output[pos] = b")"
217 |         pos += 1
218 |         if i < length - 1:
219 |             output[pos] = b","
220 |             pos += 1
221 | 
222 | 
223 | cdef str in_any_values_inline_list(list values):
224 |     cdef:
225 |         Py_ssize_t i, length = len(values)
226 |         PyObject *item
227 |         bint with_null = length == 0
228 | 
229 |     for i in range(length):
230 |         item = PyList_GET_ITEM(<PyObject *> values, i)
231 |         if item == Py_None:
232 |             with_null = True
233 |             continue
234 |         if PyUnicode_Check(item):
235 |             return _in_any_values_inline_list_u(values)
236 |         elif PyBytes_Check(item):
237 |             return _in_any_values_inline_list_s(values)
238 |         elif PyLong_CheckExact(item) or PyObject_TypeCheck(item, &PyIntegerArrType_Type):
239 |             return _in_any_values_inline_list_int64(values)
240 |         else:
241 |             raise ValueError(f"Unsupported type of list item #{i}: {type(values[i])}")
242 | 
243 |     assert with_null
244 |     return "VALUES (null)"
245 | 
246 | 
247 | cdef str _in_any_values_inline_list_u(list values):
248 |     cdef:
249 |         Py_ssize_t i, j, length = len(values), size1 = 7 - 1, size2 = 0, size4 = 0, str_len, result_len
250 |         int kind, effective_kind
251 |         Py_UCS4 max_char
252 |         str result
253 |         char *output
254 |         char *border
255 |         char *str_data
256 |         PyObject *item
257 | 
258 |     for i in range(length):
259 |         item = PyList_GET_ITEM(<PyObject *> values, i)
260 |         if item == Py_None:
261 |             continue
262 |         if not PyUnicode_Check(item):
263 |             raise ValueError(f"Mixed types in list: expected str, got {type(values[i])}")
264 |         str_len = PyUnicode_GET_LENGTH(item)
265 |         if PyUnicode_FindChar(item, ord(b"'"), 0, str_len, 1) >= 0:
266 |             raise NotImplementedError("One of the strings requires escaping the single quote")
267 |         kind = PyUnicode_KIND(item)
268 |         if kind == 1:
269 |             size1 += str_len
270 |         elif kind == 2:
271 |             size2 += str_len
272 |         else:
273 |             size4 += str_len
274 |         size1 += 5
275 |     if size4 > 0:
276 |         effective_kind = 4
277 |     elif size2 > 0:
278 |         effective_kind = 2
279 |     else:
280 |         effective_kind = 1
281 |     result_len = size1 + size2 + size4
282 |     max_char = (1 << (8 * effective_kind) - 1) if effective_kind != 4 else 1114111
283 |     result = PyUnicode_New(result_len, max_char)
284 |     output = <char *> PyUnicode_DATA(<PyObject *> result)
285 |     with nogil:
286 |         border = output + (result_len - 2) * effective_kind
287 |         if effective_kind == 1:
288 |             memcpy(output, b"VALUES ", 7)
289 |             output += 7
290 |         elif effective_kind == 2:
291 |             memcpy(output, b"V\x00A\x00L\x00U\x00E\x00S\x00 \x00", 14)
292 |             output += 14
293 |         else:
294 |             memcpy(
295 |                 output,
296 |                 b"V\x00\x00\x00A\x00\x00\x00L\x00\x00\x00U\x00\x00\x00E\x00\x00\x00S\x00\x00\x00 \x00\x00\x00",
297 |                 28,
298 |             )
299 |             output += 28
300 | 
301 |         for i in range(length):
302 |             item = PyList_GET_ITEM(<PyObject *> values, i)
303 |             if item == Py_None:
304 |                 continue
305 |             kind = PyUnicode_KIND(item)
306 |             str_len = PyUnicode_GET_LENGTH(item)
307 |             str_data = <char *> PyUnicode_DATA(item)
308 |             if effective_kind == 1:
309 |                 memcpy(output, b"('", 2)
310 |                 output += 2
311 |                 memcpy(output, str_data, str_len)
312 |                 output += str_len
313 |                 str_len = 3 if output != border else 2
314 |                 memcpy(output, b"'),", str_len)
315 |                 output += str_len
316 |             elif effective_kind == 2:
317 |                 memcpy(output, b"(\x00'\x00", 4)
318 |                 output += 4
319 |                 if kind == 1:
320 |                     for j in range(str_len):
321 |                         output[j * 2] = str_data[j]
322 |                         output[j * 2 + 1] = 0
323 |                 else:
324 |                     memcpy(output, str_data, str_len * 2)
325 |                 output += str_len * 2
326 |                 str_len = 6 if output != border else 4
327 |                 memcpy(output, b"'\x00)\x00,\x00", str_len)
328 |                 output += str_len
329 |             else:
330 |                 memcpy(output, b"(\x00\x00\x00'\x00\x00\x00", 8)
331 |                 output += 8
332 |                 if kind == 1:
333 |                     for j in range(str_len):
334 |                         output[j * 4] = str_data[j]
335 |                         output[j * 4 + 1] = 0
336 |                         output[j * 4 + 2] = 0
337 |                         output[j * 4 + 3] = 0
338 |                 elif kind == 2:
339 |                     for j in range(str_len):
340 |                         output[j * 4] = str_data[j * 2]
341 |                         output[j * 4 + 1] = str_data[j * 2 + 1]
342 |                         output[j * 4 + 2] = 0
343 |                         output[j * 4 + 3] = 0
344 |                 else:
345 |                     memcpy(output, str_data, str_len * 4)
346 |                 output += str_len * 4
347 |                 str_len = 12 if output != border else 8
348 |                 memcpy(output, b"'\x00\x00\x00)\x00\x00\x00,\x00\x00\x00", str_len)
349 |                 output += str_len
350 | 
351 |     return result
352 | 
353 | 
354 | cdef str _in_any_values_inline_list_s(list values):
355 |     cdef:
356 |         Py_ssize_t i, length = len(values), str_len, result_len = 0
357 |         str result
358 |         char *output
359 |         char *border
360 |         PyObject *item
361 | 
362 |     for i in range(length):
363 |         item = PyList_GET_ITEM(<PyObject *> values, i)
364 |         if item == Py_None:
365 |             continue
366 |         if not PyBytes_Check(item):
367 |             raise ValueError(f"Mixed types in list: expected bytes, got {type(values[i])}")
368 |         str_len = PyBytes_GET_SIZE(item)
369 |         if memchr(PyBytes_AS_STRING(item), b"'", str_len) != NULL:
370 |             raise NotImplementedError("One of the strings requires escaping the single quote")
371 |         result_len += str_len + 5
372 | 
373 |     result_len += 7 - 1
374 |     result = PyUnicode_New(result_len, 255)
375 |     output = <char *> PyUnicode_DATA(<PyObject *> result)
376 |     with nogil:
377 |         border = output + result_len - 2
378 |         memcpy(output, b"VALUES ", 7)
379 |         output += 7
380 | 
381 |         for i in range(length):
382 |             item = PyList_GET_ITEM(<PyObject *> values, i)
383 |             if item == Py_None:
384 |                 continue
385 |             str_len = PyBytes_GET_SIZE(item)
386 |             str_data = PyBytes_AS_STRING(item)
387 |             memcpy(output, b"('", 2)
388 |             output += 2
389 |             memcpy(output, str_data, str_len)
390 |             output += str_len
391 |             str_len = 3 if output != border else 2
392 |             memcpy(output, b"'),", str_len)
393 |             output += str_len
394 | 
395 |     return result
396 | 
397 | 
398 | @cython.cdivision(True)
399 | cdef str _in_any_values_inline_list_int64(list values):
400 |     cdef:
401 |         Py_ssize_t i, length = len(values), pos = 7, nulls = 0, result_len, valstart
402 |         long val = 0, max_val = 0
403 |         int digits = 0
404 |         str result
405 |         char *output
406 |         lldiv_t qr
407 |         PyObject *item
408 | 
409 |     for i in range(length):
410 |         item = PyList_GET_ITEM(<PyObject *> values, i)
411 |         if item == Py_None:
412 |             nulls += 1
413 |             continue
414 |         if PyLong_CheckExact(item):
415 |             val = PyLong_AsLong(item)
416 |         else:
417 |             PyArray_ScalarAsCtype(item, &val)
418 |         if val > max_val:
419 |             max_val = val
420 | 
421 |     while max_val:
422 |         max_val //= 10
423 |         digits += 1
424 | 
425 |     result_len = 6 + 1 + (3 + digits) * (length - nulls) - 1
426 |     result = PyUnicode_New(result_len, 255)
427 |     output = <char *> PyUnicode_DATA(<PyObject *> result)
428 |     with nogil:
429 |         memcpy(output, b"VALUES ", 7)
430 |         for i in range(length):
431 |             item = PyList_GET_ITEM(<PyObject *> values, i)
432 |             if item == Py_None:
433 |                 continue
434 |             if PyLong_CheckExact(item):
435 |                 val = PyLong_AsLong(item)
436 |             else:
437 |                 PyArray_ScalarAsCtype(item, &val)
438 |             output[pos] = b"("
439 |             pos += 1
440 |             valstart = pos
441 |             pos += digits - 1
442 |             qr.quot = val
443 |             while True:
444 |                 qr = lldiv(qr.quot, 10)
445 |                 output[pos] = (<char> b'0') + (<char> qr.rem)
446 |                 pos -= 1
447 |                 if qr.quot == 0:
448 |                     break
449 |             while pos >= valstart:
450 |                 output[pos] = b" "
451 |                 pos -= 1
452 |             pos = valstart + digits
453 |             output[pos] = b")"
454 |             pos += 1
455 |             if pos < result_len:
456 |                 output[pos] = b","
457 |                 pos += 1
458 |     return result
459 | 
460 | 
461 | def in_inline(values) -> str:
462 |     if PyArray_CheckExact(values):
463 |         return in_inline_array(values)
464 |     if PyList_CheckExact(values):
465 |         return in_inline_list(values)
466 |     raise ValueError(f"Only numpy arrays and lists are supported, got {type(values)}")
467 | 
468 | 
469 | cdef str in_inline_array(ndarray values):
470 |     if values.ndim != 1:
471 |         raise ValueError(f"We support only 1-dimensional numpy arrays, got {values.ndim}")
472 |     if values.dtype.kind not in ("S", "U", "i", "u"):
473 |         raise ValueError(f"unsupported dtype {values.dtype}")
474 |     if len(values) == 0:
475 |         return "null"
476 | 
477 |     cdef:
478 |         np_dtype dtype = <np_dtype> PyArray_DESCR(values)
479 |         int is_s = dtype.kind == b"S"
480 |         int is_str = is_s or dtype.kind == b"U"
481 |         int stride = PyArray_STRIDE(values, 0)
482 |         int itemsize = dtype.itemsize
483 |         int length = PyArray_DIM(values, 0)
484 |         int effective_itemsize = (
485 |             (itemsize if is_s else (itemsize >> 2)) + 2
486 |         ) if is_str else len(str(values.max()))
487 |         Py_ssize_t size = (effective_itemsize + 1) * length - 1
488 |         int kind = 1
489 |         Py_UCS4 max_char
490 |         char *data = <char *> PyArray_DATA(values)
491 | 
492 |     if is_str and not is_s:
493 |         kind = scan_unicode_kind(data, length * itemsize)
494 |         max_char = (1 << (8 * kind) - 1) if kind != 4 else 1114111
495 |     else:
496 |         max_char = 255
497 | 
498 |     cdef:
499 |         result = PyUnicode_New(size, max_char)
500 |         char *buf = <char *> PyUnicode_DATA(<PyObject *> result)
501 | 
502 |     if is_str and itemsize == stride:
503 |         if is_s:
504 |             if memchr(data, b"'", length * itemsize) != NULL:
505 |                 raise NotImplementedError("One of the strings requires escaping the single quote")
506 |         else:
507 |             if wmemchr(<const wchar_t *> data, b"'", length * itemsize >> 2) != NULL:
508 |                 raise NotImplementedError("One of the strings requires escaping the single quote")
509 | 
510 |     with nogil:
511 |         if is_str:
512 |             if is_s:
513 |                 if not _in_s(data, stride, itemsize, length, buf):
514 |                     raise NotImplementedError("One of the strings requires escaping the single quote")
515 |             else:
516 |                 if not _in_u(data, stride, itemsize, length, kind, buf):
517 |                     raise NotImplementedError("One of the strings requires escaping the single quote")
518 |         elif itemsize == 8:
519 |             _in_int64(data, stride, length, effective_itemsize, buf)
520 |         else:
521 |             raise ValueError(f"unsupported dtype {dtype}")
522 |     return result
523 | 
524 | 
525 | cdef bint _in_s(const char *data,
526 |                 int stride,
527 |                 int itemsize,
528 |                 int length,
529 |                 char *output) nogil:
530 |     cdef:
531 |         int i
532 |         char *output_start = output
533 |         char *quoteptr
534 |         char *nullptr
535 | 
536 |     for i in range(length):
537 |         output[0] = b"'"
538 |         output += 1
539 |         memcpy(output, data + stride * i, itemsize)
540 |         if stride != itemsize and memchr(output, b"'", itemsize) != NULL:
541 |             return False
542 |         output += itemsize
543 |         output[0] = b"'"
544 |         output += 1
545 |         if i < length - 1:
546 |             output[0] = b","
547 |             output += 1
548 | 
549 |     nullptr = <char *> memchr(output_start, 0, output - output_start)
550 |     while nullptr:
551 |         quoteptr = <char *> memchr(nullptr + 1, b"'", itemsize)
552 |         nullptr[0] = b"'"
553 |         for i in range(1, (quoteptr - nullptr) + 1):
554 |             nullptr[i] = b" "
555 |         nullptr = <char *> memchr(quoteptr, 0, output - quoteptr)
556 | 
557 |     return True
558 | 
559 | 
560 | cdef bint _in_u(const char *data,
561 |                 int stride,
562 |                 int itemsize,
563 |                 int length,
564 |                 int kind,
565 |                 char *output) nogil:
566 |     cdef:
567 |         int i, j, offset, pad
568 |         const char *quote = b"'\x00\x00\x00"
569 |         const char *comma = b",\x00\x00\x00"
570 |         char c
571 | 
572 |     for i in range(length):
573 |         memcpy(output, quote, kind)
574 |         output += kind
575 |         offset = stride * i
576 |         if stride != itemsize and wmemchr(<const wchar_t *> data + offset, b"'", itemsize >> 2) != NULL:
577 |             return False
578 |         pad = itemsize
579 |         if kind == 4:
580 |             for j in range(0, itemsize, 4):
581 |                 c = data[offset + j]
582 |                 if c == 0:
583 |                     pad = j
584 |                     break
585 |                 output[j] = c
586 |                 output[j + 1] = data[offset + j + 1]
587 |                 output[j + 2] = data[offset + j + 2]
588 |                 output[j + 3] = data[offset + j + 3]
589 |             output += pad
590 |         elif kind == 2:
591 |             for j in range(0, itemsize, 4):
592 |                 c = data[offset + j]
593 |                 if c == 0:
594 |                     pad = j
595 |                     break
596 |                 output[j >> 1] = c
597 |                 output[(j >> 1) + 1] = data[offset + j + 1]
598 |             output += (pad >> 1)
599 |         else:
600 |             for j in range(0, itemsize, 4):
601 |                 c = data[offset + j]
602 |                 if c == 0:
603 |                     pad = j
604 |                     break
605 |                 output[j >> 2] = c
606 |             output += (pad >> 2)
607 |         memcpy(output, quote, kind)
608 |         output += kind
609 |         for j in range(pad, itemsize, 4):
610 |             output[0] = b" "
611 |             if kind > 1:
612 |                 output[1] = 0
613 |                 if kind > 2:
614 |                     output[2] = 0
615 |                     output[3] = 0
616 |             output += kind
617 |         if i < length - 1:
618 |             memcpy(output, comma, kind)
619 |             output += kind
620 | 
621 |     return True
622 | 
623 | 
624 | cdef void _in_int64(const char *data,
625 |                     int stride,
626 |                     int length,
627 |                     int alignment,
628 |                     char *output) nogil:
629 |     cdef:
630 |         int i, pos = 0, valstart
631 |         lldiv_t qr
632 | 
633 |     for i in range(length):
634 |         valstart = pos
635 |         pos += alignment - 1
636 |         qr.quot = (<const int64_t *>(data + i * stride))[0]
637 |         while True:
638 |             qr = lldiv(qr.quot, 10)
639 |             output[pos] = (<char>b'0') + (<char>qr.rem)
640 |             pos -= 1
641 |             if qr.quot == 0:
642 |                 break
643 |         while pos >= valstart:
644 |             output[pos] = b" "
645 |             pos -= 1
646 |         pos = valstart + alignment
647 |         if i < length - 1:
648 |             output[pos] = b","
649 |             pos += 1
650 | 
651 | 
652 | cdef str in_inline_list(list values):
653 |     cdef:
654 |         Py_ssize_t i, length = len(values)
655 |         PyObject *item
656 |         bint with_null = length == 0
657 | 
658 |     for i in range(length):
659 |         item = PyList_GET_ITEM(<PyObject *> values, i)
660 |         if item == Py_None:
661 |             with_null = True
662 |             continue
663 |         if PyUnicode_Check(item):
664 |             return _in_inline_list_u(values)
665 |         elif PyBytes_Check(item):
666 |             return _in_inline_list_s(values)
667 |         elif PyLong_CheckExact(item) or PyObject_TypeCheck(item, &PyIntegerArrType_Type):
668 |             return _in_inline_list_int64(values)
669 |         else:
670 |             raise ValueError(f"Unsupported type of list item #{i}: {type(values[i])}")
671 | 
672 |     assert with_null
673 |     return "null"
674 | 
675 | 
676 | cdef str _in_inline_list_u(list values):
677 |     cdef:
678 |         Py_ssize_t i, j, length = len(values), size1 = -1, size2 = 0, size4 = 0, str_len, result_len
679 |         int kind, effective_kind
680 |         Py_UCS4 max_char
681 |         str result
682 |         char *output
683 |         char *border
684 |         char *str_data
685 |         PyObject *item
686 |         bint quote_found = False
687 | 
688 |     for i in range(length):
689 |         item = PyList_GET_ITEM(<PyObject *> values, i)
690 |         if item == Py_None:
691 |             continue
692 |         if not PyUnicode_Check(item):
693 |             raise ValueError(f"Mixed types in list: expected str, got {type(values[i])}")
694 |         kind = PyUnicode_KIND(item)
695 |         str_len = PyUnicode_GET_LENGTH(item)
696 |         if kind == 1:
697 |             size1 += str_len
698 |         elif kind == 2:
699 |             size2 += str_len
700 |         else:
701 |             size4 += str_len
702 |         size1 += 3
703 |     if size4 > 0:
704 |         effective_kind = 4
705 |     elif size2 > 0:
706 |         effective_kind = 2
707 |     else:
708 |         effective_kind = 1
709 |     result_len = size1 + size2 + size4
710 |     max_char = (1 << (8 * effective_kind) - 1) if effective_kind != 4 else 1114111
711 |     result = PyUnicode_New(result_len, max_char)
712 |     output = <char *> PyUnicode_DATA(<PyObject *> result)
713 | 
714 |     with nogil:
715 |         border = output + (result_len - 1) * effective_kind
716 | 
717 |         for i in range(length):
718 |             item = PyList_GET_ITEM(<PyObject *> values, i)
719 |             if item == Py_None:
720 |                 continue
721 |             str_len = PyUnicode_GET_LENGTH(item)
722 |             if PyUnicode_FindChar(item, ord(b"'"), 0, str_len, 1) >= 0:
723 |                 quote_found = True
724 |                 break
725 |             kind = PyUnicode_KIND(item)
726 |             str_data = <char *> PyUnicode_DATA(item)
727 |             if effective_kind == 1:
728 |                 output[0] = b"'"
729 |                 output += 1
730 |                 memcpy(output, str_data, str_len)
731 |                 output += str_len
732 |                 str_len = 2 if output != border else 1
733 |                 memcpy(output, b"',", str_len)
734 |                 output += str_len
735 |             elif effective_kind == 2:
736 |                 memcpy(output, b"'\x00", 2)
737 |                 output += 2
738 |                 if kind == 1:
739 |                     for j in range(str_len):
740 |                         output[j * 2] = str_data[j]
741 |                         output[j * 2 + 1] = 0
742 |                 else:
743 |                     memcpy(output, str_data, str_len * 2)
744 |                 output += str_len * 2
745 |                 str_len = 4 if output != border else 2
746 |                 memcpy(output, b"'\x00,\x00", str_len)
747 |                 output += str_len
748 |             else:
749 |                 memcpy(output, b"'\x00\x00\x00", 4)
750 |                 output += 4
751 |                 if kind == 1:
752 |                     for j in range(str_len):
753 |                         output[j * 4] = str_data[j]
754 |                         output[j * 4 + 1] = 0
755 |                         output[j * 4 + 2] = 0
756 |                         output[j * 4 + 3] = 0
757 |                 elif kind == 2:
758 |                     for j in range(str_len):
759 |                         output[j * 4] = str_data[j * 2]
760 |                         output[j * 4 + 1] = str_data[j * 2 + 1]
761 |                         output[j * 4 + 2] = 0
762 |                         output[j * 4 + 3] = 0
763 |                 else:
764 |                     memcpy(output, str_data, str_len * 4)
765 |                 output += str_len * 4
766 |                 str_len = 8 if output != border else 4
767 |                 memcpy(output, b"'\x00\x00\x00,\x00\x00\x00", str_len)
768 |                 output += str_len
769 | 
770 |     if quote_found:
771 |         raise NotImplementedError("One of the strings requires escaping the single quote")
772 |     return result
773 | 
774 | 
775 | cdef str _in_inline_list_s(list values):
776 |     cdef:
777 |         Py_ssize_t i, length = len(values), str_len, result_len = 0
778 |         str result
779 |         char *output
780 |         char *border
781 |         PyObject *item
782 |         bint quote_found = False
783 | 
784 |     for i in range(length):
785 |         item = PyList_GET_ITEM(<PyObject *> values, i)
786 |         if item == Py_None:
787 |             continue
788 |         if not PyBytes_Check(item):
789 |             raise ValueError(f"Mixed types in list: expected bytes, got {type(values[i])}")
790 |         result_len += PyBytes_GET_SIZE(item) + 3
791 | 
792 |     result_len -= 1
793 |     result = PyUnicode_New(result_len, 255)
794 |     output = <char *> PyUnicode_DATA(<PyObject *> result)
795 |     with nogil:
796 |         border = output + result_len - 2
797 | 
798 |         for i in range(length):
799 |             item = PyList_GET_ITEM(<PyObject *> values, i)
800 |             if item == Py_None:
801 |                 continue
802 |             str_len = PyBytes_GET_SIZE(item)
803 |             str_data = PyBytes_AS_STRING(item)
804 |             output[0] = b"'"
805 |             output += 1
806 |             if memchr(str_data, b"'", str_len) != NULL:
807 |                 quote_found = True
808 |                 break
809 |             memcpy(output, str_data, str_len)
810 |             output += str_len
811 |             str_len = 2 if output != border else 1
812 |             memcpy(output, b"',", str_len)
813 |             output += str_len
814 | 
815 |     if quote_found:
816 |         raise NotImplementedError("One of the strings requires escaping the single quote")
817 |     return result
818 | 
819 | 
820 | @cython.cdivision(True)
821 | cdef str _in_inline_list_int64(list values):
822 |     cdef:
823 |         Py_ssize_t i, length = len(values), pos = 0, nulls = 0, result_len, valstart
824 |         long val = 0, max_val = 0
825 |         int digits = 0
826 |         str result
827 |         char *output
828 |         lldiv_t qr
829 |         PyObject *item
830 | 
831 |     for i in range(length):
832 |         item = PyList_GET_ITEM(<PyObject *> values, i)
833 |         if item == Py_None:
834 |             nulls += 1
835 |             continue
836 |         if PyLong_CheckExact(item):
837 |             val = PyLong_AsLong(item)
838 |         else:
839 |             PyArray_ScalarAsCtype(item, &val)
840 |         if val > max_val:
841 |             max_val = val
842 | 
843 |     while max_val:
844 |         max_val //= 10
845 |         digits += 1
846 | 
847 |     result_len = (1 + digits) * (length - nulls) - 1
848 |     result = PyUnicode_New(result_len, 255)
849 |     output = <char *> PyUnicode_DATA(<PyObject *> result)
850 |     with nogil:
851 |         for i in range(length):
852 |             item = PyList_GET_ITEM(<PyObject *> values, i)
853 |             if item == Py_None:
854 |                 continue
855 |             if PyLong_CheckExact(item):
856 |                 val = PyLong_AsLong(item)
857 |             else:
858 |                 PyArray_ScalarAsCtype(item, &val)
859 |             valstart = pos
860 |             pos += digits - 1
861 |             qr.quot = val
862 |             while True:
863 |                 qr = lldiv(qr.quot, 10)
864 |                 output[pos] = (<char> b'0') + (<char> qr.rem)
865 |                 pos -= 1
866 |                 if qr.quot == 0:
867 |                     break
868 |             while pos >= valstart:
869 |                 output[pos] = b" "
870 |                 pos -= 1
871 |             pos = valstart + digits
872 |             if pos < result_len:
873 |                 output[pos] = b","
874 |                 pos += 1
875 |     return result
876 | 


--------------------------------------------------------------------------------
/2023/pyx/types_accelerated.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: extra_compile_args = -std=c++17
  5 | # distutils: libraries = mimalloc
  6 | # distutils: runtime_library_dirs = /usr/local/lib
  7 | 
  8 | from cpython cimport PyDict_GetItem, PyDict_New, PyObject
  9 | from cpython.memoryview cimport PyMemoryView_Check, PyMemoryView_GET_BUFFER
 10 | from cython.operator cimport dereference, postincrement
 11 | from libc.stdint cimport int64_t
 12 | 
 13 | from athenian.api.native.cpython cimport (
 14 |     Py_INCREF,
 15 |     Py_None,
 16 |     Py_TYPE,
 17 |     PyByteArray_AS_STRING,
 18 |     PyByteArray_CheckExact,
 19 |     PyBytes_AS_STRING,
 20 |     PyBytes_Check,
 21 |     PyLong_AsLong,
 22 |     PyMemberDef,
 23 |     PyTypeObject,
 24 | )
 25 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport (
 26 |     mi_heap_allocator_from_capsule,
 27 |     mi_heap_destroy_stl_allocator,
 28 |     mi_unordered_set,
 29 | )
 30 | from athenian.api.native.numpy cimport (
 31 |     NPY_FR_s,
 32 |     PyArray_DATA,
 33 |     PyArray_DIM,
 34 |     PyDatetimeScalarObject,
 35 |     npy_intp,
 36 | )
 37 | from athenian.api.native.optional cimport optional
 38 | 
 39 | from medvedi import DataFrame
 40 | 
 41 | from athenian.api.internal.miners.participation import PRParticipationKind
 42 | from athenian.api.models.metadata.github import (
 43 |     PullRequest,
 44 |     PullRequestComment,
 45 |     PullRequestCommit,
 46 |     PullRequestReview,
 47 |     Release,
 48 | )
 49 | 
 50 | pr_user_node_id_col = PullRequest.user_node_id.name
 51 | pr_merged_by_id_col = PullRequest.merged_by_id.name
 52 | release_author_node_id_col = Release.author_node_id.name
 53 | review_user_node_id_col = PullRequestReview.user_node_id.name
 54 | comment_user_node_id_col = PullRequestComment.user_node_id.name
 55 | commit_committer_user_id_col = PullRequestCommit.committer_user_id.name
 56 | commit_author_user_id_col = PullRequestCommit.author_user_id.name
 57 | values_attr = "values"
 58 | 
 59 | # these must be Python integers
 60 | PRParticipationKind_AUTHOR = PRParticipationKind.AUTHOR
 61 | PRParticipationKind_REVIEWER = PRParticipationKind.REVIEWER
 62 | PRParticipationKind_COMMENTER = PRParticipationKind.COMMENTER
 63 | PRParticipationKind_COMMIT_AUTHOR = PRParticipationKind.COMMIT_AUTHOR
 64 | PRParticipationKind_COMMIT_COMMITTER = PRParticipationKind.COMMIT_COMMITTER
 65 | PRParticipationKind_MERGER = PRParticipationKind.MERGER
 66 | PRParticipationKind_RELEASER = PRParticipationKind.RELEASER
 67 | 
 68 | 
 69 | cdef enum MinedPullRequestFields:
 70 |     MinedPullRequest_check_run = 0
 71 |     MinedPullRequest_comments = 1
 72 |     MinedPullRequest_commits = 2
 73 |     MinedPullRequest_deployments = 3
 74 |     MinedPullRequest_jiras = 4
 75 |     MinedPullRequest_labels = 5
 76 |     MinedPullRequest_pr = 6
 77 |     MinedPullRequest_release = 7
 78 |     MinedPullRequest_review_comments = 8
 79 |     MinedPullRequest_review_requests = 9
 80 |     MinedPullRequest_reviews = 10
 81 | 
 82 | 
 83 | cdef :
 84 |     set empty_set = set()
 85 |     Py_ssize_t df_columns_offset = (<PyTypeObject *> DataFrame).tp_members[0].offset
 86 | 
 87 | 
 88 | def extract_participant_nodes(mpr, alloc_capsule=None) -> dict:
 89 |     # the slot indexes correspond to the alphabetical order of fields
 90 |     cdef:
 91 |         optional[mi_heap_destroy_stl_allocator[char]] alloc
 92 |         optional[mi_unordered_set[int64_t]] boilerplate
 93 |         PyMemberDef *mpr_slots = Py_TYPE(<PyObject *> mpr).tp_members
 94 |         PyObject *pr = dereference(
 95 |             <PyObject **> ((<char *><PyObject *> mpr) + mpr_slots[<int>MinedPullRequest_pr].offset)
 96 |         )
 97 |         PyObject *commits = dereference(
 98 |             <PyObject **> ((<char *><PyObject *> mpr) + mpr_slots[<int>MinedPullRequest_commits].offset)
 99 |         )
100 |         PyObject *commits_dict = dereference(
101 |             <PyObject **> ((<char *><PyObject *> commits) + df_columns_offset)
102 |         )
103 |         PyObject *commit_committers = PyDict_GetItem(
104 |             <object>commits_dict, commit_committer_user_id_col,
105 |         )
106 |         PyObject *commit_authors = PyDict_GetItem(
107 |             <object>commits_dict,  commit_author_user_id_col,
108 |         )
109 | 
110 |         PyObject *reviews = PyDict_GetItem(
111 |             <object> dereference(
112 |                 <PyObject **> ((<char *> dereference(
113 |                     <PyObject **> ((<char *><PyObject *> mpr) + mpr_slots[<int>MinedPullRequest_reviews].offset)
114 |                 )) + df_columns_offset)
115 |             ),
116 |             review_user_node_id_col,
117 |         )
118 |         PyObject *comments = PyDict_GetItem(
119 |             <object> dereference(
120 |                 <PyObject **> ((<char *> dereference(
121 |                         <PyObject **> ((<char *><PyObject *> mpr) + mpr_slots[<int>MinedPullRequest_comments].offset)
122 |                 )) + df_columns_offset)
123 |             ),
124 |             comment_user_node_id_col,
125 |         )
126 |         PyObject *release = dereference(
127 |             <PyObject **> ((<char *> <PyObject *> mpr) + mpr_slots[<int>MinedPullRequest_release].offset)
128 |         )
129 |         PyObject *author = PyDict_GetItem(<object> pr, pr_user_node_id_col)
130 |         PyObject *merger = PyDict_GetItem(<object> pr, pr_merged_by_id_col)
131 |         PyObject *releaser = PyDict_GetItem(
132 |             <object> dereference(
133 |                 <PyObject **> ((<char *><PyObject *> mpr) + mpr_slots[<int>MinedPullRequest_release].offset)
134 |             ),
135 |             release_author_node_id_col,
136 |         )
137 |         dict participants = PyDict_New()
138 |         set py_boilerplate
139 |         int64_t *data
140 |         npy_intp i
141 |         mi_unordered_set[int64_t].const_iterator it
142 | 
143 |     if merger == Py_None or PyLong_AsLong(merger) == 0:
144 |         participants[PRParticipationKind_MERGER] = empty_set
145 |     else:
146 |         Py_INCREF(merger)
147 |         participants[PRParticipationKind_MERGER] = {<object> merger}
148 | 
149 |     if releaser == Py_None or PyLong_AsLong(releaser) == 0:
150 |         participants[PRParticipationKind_RELEASER] = empty_set
151 |     else:
152 |         Py_INCREF(releaser)
153 |         participants[PRParticipationKind_RELEASER] = {<object> releaser}
154 | 
155 |     if alloc_capsule is not None:
156 |         alloc.emplace(dereference(mi_heap_allocator_from_capsule(alloc_capsule)))
157 |     else:
158 |         alloc.emplace()
159 |     boilerplate.emplace(dereference(alloc))
160 | 
161 |     data = <int64_t *> PyArray_DATA(reviews)
162 |     for i in range(PyArray_DIM(reviews, 0)):
163 |         dereference(boilerplate).emplace(data[i])
164 |     dereference(boilerplate).erase(0)
165 | 
166 |     if author == Py_None or PyLong_AsLong(author) == 0:
167 |         participants[PRParticipationKind_AUTHOR] = empty_set
168 |     else:
169 |         Py_INCREF(author)
170 |         participants[PRParticipationKind_AUTHOR] = {<object> author}
171 |         dereference(boilerplate).erase(PyLong_AsLong(author))
172 | 
173 |     py_boilerplate = set()
174 |     it = dereference(boilerplate).const_begin()
175 |     while it != dereference(boilerplate).const_end():
176 |         py_boilerplate.add(dereference(it))
177 |         postincrement(it)
178 |     participants[PRParticipationKind_REVIEWER] = py_boilerplate
179 | 
180 |     dereference(boilerplate).clear()
181 |     data = <int64_t *> PyArray_DATA(comments)
182 |     for i in range(PyArray_DIM(comments, 0)):
183 |         dereference(boilerplate).emplace(data[i])
184 |     dereference(boilerplate).erase(0)
185 |     py_boilerplate = set()
186 |     it = dereference(boilerplate).const_begin()
187 |     while it != dereference(boilerplate).const_end():
188 |         py_boilerplate.add(dereference(it))
189 |         postincrement(it)
190 |     participants[PRParticipationKind_COMMENTER] = py_boilerplate
191 | 
192 |     dereference(boilerplate).clear()
193 |     data = <int64_t *> PyArray_DATA(commit_committers)
194 |     for i in range(PyArray_DIM(commit_committers, 0)):
195 |         dereference(boilerplate).emplace(data[i])
196 |     dereference(boilerplate).erase(0)
197 |     py_boilerplate = set()
198 |     it = dereference(boilerplate).const_begin()
199 |     while it != dereference(boilerplate).const_end():
200 |         py_boilerplate.add(dereference(it))
201 |         postincrement(it)
202 |     participants[PRParticipationKind_COMMIT_COMMITTER] = py_boilerplate
203 | 
204 |     dereference(boilerplate).clear()
205 |     data = <int64_t *> PyArray_DATA(commit_authors)
206 |     for i in range(PyArray_DIM(commit_authors, 0)):
207 |         dereference(boilerplate).emplace(data[i])
208 |     dereference(boilerplate).erase(0)
209 |     py_boilerplate = set()
210 |     it = dereference(boilerplate).const_begin()
211 |     while it != dereference(boilerplate).const_end():
212 |         py_boilerplate.add(dereference(it))
213 |         postincrement(it)
214 |     participants[PRParticipationKind_COMMIT_AUTHOR] = py_boilerplate
215 | 
216 |     return participants
217 | 
218 | 
219 | # for field_name, (field_dtype, _) in self.dtype.fields.items():
220 | #     if np.issubdtype(field_dtype, np.datetime64):
221 | #         if (dt := getattr(self, field_name)) is not None and dt >= after_dt:
222 | #             changed.append(field_name)
223 | 
224 | def find_truncated_datetime(facts, offsets, time_from) -> list:
225 |     cdef:
226 |         list result = []
227 |         int64_t time_from_i64 = (<PyDatetimeScalarObject *> time_from).obval
228 |         PyMemberDef *slots = Py_TYPE(<PyObject *> facts).tp_members
229 |         PyObject *data_obj = dereference(
230 |             <PyObject **> ((<char *> <PyObject *> facts) + slots[1].offset)
231 |         )
232 |         int64_t *offsets_i64 = <int64_t *> PyArray_DATA(<PyObject *> offsets)
233 |         npy_intp size = PyArray_DIM(<PyObject *> offsets, 0), i
234 |         char *data
235 | 
236 |     if PyBytes_Check(data_obj):
237 |         data = PyBytes_AS_STRING(data_obj)
238 |     elif PyMemoryView_Check(<object> data_obj):
239 |         data = <char *>PyMemoryView_GET_BUFFER(<object> data_obj).buf
240 |     elif PyByteArray_CheckExact(data_obj):
241 |         data = PyByteArray_AS_STRING(data_obj)
242 |     else:
243 |         raise AssertionError(f"unsupported buffer type: {type(facts.data)}")
244 | 
245 |     assert (<PyDatetimeScalarObject *> time_from).obmeta.base == NPY_FR_s
246 |     for i in range(size):
247 |         if dereference(<int64_t *>(data + offsets_i64[i])) >= time_from_i64:
248 |             result.append(i)
249 |     return result
250 | 


--------------------------------------------------------------------------------
/2023/pyx/unordered_unique.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: extra_compile_args = -mavx2 -ftree-vectorize -std=c++17
  5 | # distutils: libraries = mimalloc
  6 | # distutils: runtime_library_dirs = /usr/local/lib
  7 | 
  8 | import cython
  9 | 
 10 | from cpython cimport PyObject
 11 | from cython.operator cimport dereference as deref
 12 | from libc.stddef cimport wchar_t
 13 | from libc.stdint cimport int32_t, int64_t
 14 | from libc.string cimport memchr, memcpy
 15 | from numpy cimport (
 16 |     PyArray_DATA,
 17 |     PyArray_DESCR,
 18 |     PyArray_DIM,
 19 |     PyArray_IS_C_CONTIGUOUS,
 20 |     PyArray_NDIM,
 21 |     PyArray_STRIDE,
 22 |     dtype as np_dtype,
 23 |     ndarray,
 24 | )
 25 | 
 26 | from athenian.api.internal.miners.github.dag_accelerated import searchsorted_inrange
 27 | from athenian.api.native.cpython cimport (
 28 |     Py_INCREF,
 29 |     Py_None,
 30 |     PyUnicode_DATA,
 31 |     PyUnicode_GET_LENGTH,
 32 |     PyUnicode_KIND,
 33 | )
 34 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport (
 35 |     mi_heap_destroy_stl_allocator,
 36 |     mi_unordered_map,
 37 |     mi_unordered_set,
 38 |     pair,
 39 | )
 40 | from athenian.api.native.optional cimport optional
 41 | from athenian.api.native.string_view cimport string_view
 42 | 
 43 | import numpy as np
 44 | 
 45 | 
 46 | cdef extern from "wchar.h" nogil:
 47 |     wchar_t *wmemchr(const wchar_t *, wchar_t, size_t)
 48 | 
 49 | 
 50 | def unordered_unique(ndarray arr not None) -> np.ndarray:
 51 |     cdef:
 52 |         np_dtype dtype = <np_dtype>PyArray_DESCR(arr)
 53 |     assert PyArray_NDIM(arr) == 1
 54 |     assert PyArray_IS_C_CONTIGUOUS(arr)
 55 |     if dtype.kind == b"S" or dtype.kind == b"U":
 56 |         return _unordered_unique_str(arr, dtype)
 57 |     elif dtype.kind == b"i" or dtype.kind == b"u":
 58 |         if dtype.itemsize == 8:
 59 |             return _unordered_unique_int[int64_t](arr, dtype, 0)
 60 |         elif dtype.itemsize == 4:
 61 |             return _unordered_unique_int[int64_t](arr, dtype, 4)
 62 |         else:
 63 |             raise AssertionError(f"dtype {dtype} is not supported")
 64 |     elif dtype.kind == b"O":
 65 |         return _unordered_unique_pystr(arr)
 66 |     else:
 67 |         raise AssertionError(f"dtype {dtype} is not supported")
 68 | 
 69 | 
 70 | @cython.cdivision(True)
 71 | cdef ndarray _unordered_unique_pystr(ndarray arr):
 72 |     cdef:
 73 |         PyObject **data_in = <PyObject **>PyArray_DATA(arr)
 74 |         PyObject **data_out
 75 |         PyObject *str_obj
 76 |         char *str_data
 77 |         unsigned int str_kind
 78 |         Py_ssize_t str_len
 79 |         int64_t i, \
 80 |             length = PyArray_DIM(arr, 0), \
 81 |             stride = PyArray_STRIDE(arr, 0) >> 3
 82 |         optional[mi_heap_destroy_stl_allocator[char]] alloc
 83 |         optional[mi_unordered_map[string_view, int64_t]] hashtable
 84 |         pair[string_view, int64_t] it
 85 |         ndarray result
 86 | 
 87 |     with nogil:
 88 |         alloc.emplace()
 89 |         hashtable.emplace(deref(alloc))
 90 |         deref(hashtable).reserve(length // 16)
 91 |         for i in range(length):
 92 |             str_obj = data_in[i * stride]
 93 |             if str_obj == Py_None:
 94 |                 continue
 95 |             str_data = <char *> PyUnicode_DATA(str_obj)
 96 |             str_len = PyUnicode_GET_LENGTH(str_obj)
 97 |             str_kind = PyUnicode_KIND(str_obj)
 98 |             deref(hashtable).try_emplace(string_view(str_data, str_len * str_kind), i)
 99 | 
100 |     result = np.empty(deref(hashtable).size(), dtype=object)
101 |     data_out = <PyObject **>PyArray_DATA(result)
102 |     i = 0
103 |     for it in deref(hashtable):
104 |         str_obj = data_in[it.second]
105 |         data_out[i] = str_obj
106 |         Py_INCREF(str_obj)
107 |         i += 1
108 |     return result
109 | 
110 | 
111 | @cython.cdivision(True)
112 | cdef ndarray _unordered_unique_str(ndarray arr, np_dtype dtype):
113 |     cdef:
114 |         char *data = <char *>PyArray_DATA(arr)
115 |         int64_t i, \
116 |             itemsize = dtype.itemsize, \
117 |             length = PyArray_DIM(arr, 0), \
118 |             stride = PyArray_STRIDE(arr, 0)
119 |         optional[mi_heap_destroy_stl_allocator[string_view]] alloc
120 |         optional[mi_unordered_set[string_view]] hashtable
121 |         string_view it
122 |         ndarray result
123 | 
124 |     with nogil:
125 |         alloc.emplace()
126 |         hashtable.emplace(deref(alloc))
127 |         deref(hashtable).reserve(length // 16)
128 |         for i in range(length):
129 |             deref(hashtable).emplace(data + i * stride, itemsize)
130 | 
131 |     result = np.empty(deref(hashtable).size(), dtype=dtype)
132 | 
133 |     with nogil:
134 |         data = <char *>PyArray_DATA(result)
135 |         i = 0
136 |         for it in deref(hashtable):
137 |             memcpy(data + i * itemsize, it.data(), itemsize)
138 |             i += 1
139 |     return result
140 | 
141 | 
142 | ctypedef fused varint:
143 |     int64_t
144 |     int32_t
145 | 
146 | 
147 | @cython.cdivision(True)
148 | cdef ndarray _unordered_unique_int(ndarray arr, np_dtype dtype, varint _):
149 |     cdef:
150 |         char *data = <char *> PyArray_DATA(arr)
151 |         int64_t i, \
152 |             itemsize = dtype.itemsize, \
153 |             length = PyArray_DIM(arr, 0), \
154 |             stride = PyArray_STRIDE(arr, 0)
155 |         optional[mi_heap_destroy_stl_allocator[varint]] alloc
156 |         optional[mi_unordered_set[varint]] hashtable
157 |         varint it
158 |         ndarray result
159 | 
160 |     with nogil:
161 |         alloc.emplace()
162 |         hashtable.emplace(deref(alloc))
163 |         deref(hashtable).reserve(length // 16)
164 |         for i in range(length):
165 |             deref(hashtable).emplace((<varint *>(data + i * stride))[0])
166 | 
167 |     result = np.empty(deref(hashtable).size(), dtype=dtype)
168 | 
169 |     with nogil:
170 |         data = <char *> PyArray_DATA(result)
171 |         i = 0
172 |         for it in deref(hashtable):
173 |             (<varint *>(data + i * itemsize))[0] = it
174 |             i += 1
175 |     return result
176 | 
177 | 
178 | def in1d_str(
179 |     ndarray trial not None,
180 |     ndarray dictionary not None,
181 |     bint skip_leading_zeros = False,
182 | ) -> np.ndarray:
183 |     cdef:
184 |         np_dtype dtype_trial = <np_dtype>PyArray_DESCR(trial)
185 |         np_dtype dtype_dict = <np_dtype>PyArray_DESCR(dictionary)
186 |     assert PyArray_NDIM(trial) == 1
187 |     assert PyArray_NDIM(dictionary) == 1
188 |     assert dtype_trial.kind == b"S" or dtype_trial.kind == b"U"
189 |     assert dtype_trial.kind == dtype_dict.kind
190 |     return _in1d_str(trial, dictionary, dtype_trial.kind == b"S", skip_leading_zeros)
191 | 
192 | 
193 | cdef ndarray _in1d_str(ndarray trial, ndarray dictionary, bint is_char, int skip_leading_zeros):
194 |     cdef:
195 |         char *data_trial = <char *>PyArray_DATA(trial)
196 |         char *data_dictionary = <char *> PyArray_DATA(dictionary)
197 |         char *output
198 |         char *s
199 |         char *nullptr
200 |         np_dtype dtype_trial = <np_dtype>PyArray_DESCR(trial)
201 |         np_dtype dtype_dict = <np_dtype>PyArray_DESCR(dictionary)
202 |         int64_t i, size, \
203 |             itemsize = dtype_dict.itemsize, \
204 |             length = PyArray_DIM(dictionary, 0), \
205 |             stride = PyArray_STRIDE(dictionary, 0)
206 |         optional[mi_heap_destroy_stl_allocator[string_view]] alloc
207 |         optional[mi_unordered_set[string_view]] hashtable
208 |         mi_unordered_set[string_view].iterator end
209 |         ndarray result
210 | 
211 |     with nogil:
212 |         alloc.emplace()
213 |         hashtable.emplace(deref(alloc))
214 |         deref(hashtable).reserve(length * 4)
215 |         if is_char:
216 |             for i in range(length):
217 |                 s = data_dictionary + i * stride
218 |                 nullptr = s
219 |                 if skip_leading_zeros:
220 |                     while nullptr < (s + itemsize) and nullptr[0] == 0:
221 |                         nullptr += 1
222 |                 nullptr = <char *> memchr(nullptr, 0, itemsize + (s - nullptr))
223 |                 if nullptr:
224 |                     size = nullptr - s
225 |                 else:
226 |                     size = itemsize
227 |                 deref(hashtable).emplace(s, size)
228 |         else:
229 |             for i in range(length):
230 |                 s = data_dictionary + i * stride
231 |                 nullptr = <char *> wmemchr(<wchar_t *>s, 0, itemsize >> 2)
232 |                 if nullptr:
233 |                     size = nullptr - s
234 |                 else:
235 |                     size = itemsize
236 |                 deref(hashtable).emplace(s, size)
237 | 
238 |         itemsize = dtype_trial.itemsize
239 |         length = PyArray_DIM(trial, 0)
240 |         stride = PyArray_STRIDE(trial, 0)
241 | 
242 |     result = np.empty(length, dtype=bool)
243 | 
244 |     with nogil:
245 |         output = <char *>PyArray_DATA(result)
246 |         end = deref(hashtable).end()
247 |         if is_char:
248 |             for i in range(length):
249 |                 s = data_trial + i * stride
250 |                 nullptr = s
251 |                 if skip_leading_zeros:
252 |                     while nullptr < (s + itemsize) and nullptr[0] == 0:
253 |                         nullptr += 1
254 |                 nullptr = <char *> memchr(nullptr, 0, itemsize + (s - nullptr))
255 |                 if nullptr:
256 |                     size = nullptr - s
257 |                 else:
258 |                     size = itemsize
259 |                 output[i] = deref(hashtable).find(string_view(s, size)) != end
260 |         else:
261 |             for i in range(length):
262 |                 s = data_trial + i * stride
263 |                 nullptr = <char *> wmemchr(<wchar_t *> s, 0, itemsize >> 2)
264 |                 if nullptr:
265 |                     size = nullptr - s
266 |                 else:
267 |                     size = itemsize
268 |                 output[i] = deref(hashtable).find(string_view(s, size)) != end
269 |     return result
270 | 
271 | 
272 | def map_array_values(
273 |     ndarray arr not None,
274 |     ndarray map_keys not None,
275 |     ndarray map_values not None,
276 |     miss_value,
277 | ) -> np.ndarray:
278 |     """Map the values in the array `arr` using the dictionary expressed by two arrays.
279 | 
280 |     `map_keys` and `map_values` must have the same length and together represent the translation
281 |     array.
282 |     `map_keys` must be sorted.
283 |     Values in `arr` not found in `map_keys` will be mapped to `miss_value`, which datatype
284 |     should be compatible with `map_values` datatype.
285 |     """
286 |     cdef:
287 |         ndarray found_keys_indexes, mapped, non_matching_keys
288 | 
289 |     assert len(map_keys) == len(map_values)
290 |     if len(map_keys) == 0:
291 |         return np.full(len(arr), miss_value)
292 |     # indexes selecting from map_key in the same order as ar
293 |     found_keys_indexes = searchsorted_inrange(map_keys, arr)
294 |     mapped = map_values[found_keys_indexes]
295 | 
296 |     # found_keys_indexes will also have an index for ar elements not present in map_keys
297 |     # these positions must be set to miss_value
298 |     non_matching_keys = map_keys[found_keys_indexes] != arr
299 |     mapped[non_matching_keys] = miss_value
300 |     return mapped
301 | 


--------------------------------------------------------------------------------
/2023/pyx/utils_accelerated.h:
--------------------------------------------------------------------------------
  1 | #include <type_traits>
  2 | 
  3 | #include <immintrin.h>
  4 | 
  5 | #define interleave_lo do { \
  6 |   __m256i tlo = _mm256_permute4x64_epi64(s, 0b00010100); \
  7 |   tlo = _mm256_unpacklo_epi8(tlo, zeros); \
  8 |   if (step == 2) { \
  9 | 	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * step), tlo); \
 10 |   } else { \
 11 | 	__m256i tlolo = _mm256_permute4x64_epi64(tlo, 0b00010100); \
 12 | 	tlolo = _mm256_unpacklo_epi8(tlolo, zeros); \
 13 | 	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * step), tlolo); \
 14 | 	__m256i tlohi = _mm256_permute4x64_epi64(tlo, 0b11101011); \
 15 | 	tlohi = _mm256_unpackhi_epi8(tlohi, zeros); \
 16 | 	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * step + 32), tlohi); \
 17 |   } \
 18 | } while (false)
 19 | 
 20 | #define interleave_hi do { \
 21 |   __m256i thi = _mm256_permute4x64_epi64(s, 0b11101011); \
 22 |   thi = _mm256_unpackhi_epi8(thi, zeros); \
 23 |   if (step == 2) { \
 24 | 	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * step + 32), thi); \
 25 |   } else { \
 26 | 	__m256i thilo = _mm256_permute4x64_epi64(thi, 0b00010100); \
 27 | 	thilo = _mm256_unpacklo_epi8(thilo, zeros); \
 28 | 	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * step + 64), thilo); \
 29 | 	__m256i thihi = _mm256_permute4x64_epi64(thi, 0b11101011); \
 30 | 	thihi = _mm256_unpackhi_epi8(thihi, zeros); \
 31 | 	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * step + 96), thihi); \
 32 |   } \
 33 | } while (false)
 34 | 
 35 | template <int step, typename = std::enable_if_t<step == 2 || step == 4>>
 36 | void interleave_bytes(const char *__restrict__ src, ssize_t length, char *__restrict__ out) {
 37 |   const __m256i zeros = _mm256_setzero_si256();
 38 |   __m256i s;
 39 |   ssize_t i;
 40 |   for (i = 0; i < length - 31; i += 32) {
 41 |     s = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(src + i));
 42 |     interleave_lo;
 43 |     interleave_hi;
 44 |   }
 45 |   if (i < length - 15) {
 46 | 	const __m128i *head = reinterpret_cast<const __m128i *>(src + i);
 47 | 	s = _mm256_loadu2_m128i(head, head);
 48 | 	interleave_lo;
 49 | 	i += 16;
 50 |   }
 51 |   if (step == 2) {
 52 | 	for (; i < length - 3; i += 4) {
 53 | 	  uint64_t quad = *reinterpret_cast<const uint32_t *>(src + i);
 54 | 	  *reinterpret_cast<uint64_t *>(out + i * step) = ((quad & 0xFF000000) << 24) | ((quad & 0xFF0000) << 16) | ((quad & 0xFF00) << 8) | (quad & 0xFF);
 55 | 	}
 56 | 	for (; i < length; i++) {
 57 | 	  *reinterpret_cast<uint16_t *>(out + i * step) = static_cast<const uint16_t>(reinterpret_cast<const uint8_t *>(src)[i]);
 58 | 	}
 59 |   } else {
 60 | 	for (; i < length - 1; i += 2) {
 61 | 	  uint64_t pair = *reinterpret_cast<const uint16_t *>(src + i);
 62 | 	  *reinterpret_cast<uint64_t *>(out + i * step) = ((pair & 0xFF00) << 24) | (pair & 0xFF);
 63 | 	}
 64 | 	if (i < length) {
 65 | 	  *reinterpret_cast<uint32_t *>(out + i * step) = static_cast<const uint32_t>(reinterpret_cast<const uint8_t *>(src)[i]);
 66 | 	}
 67 |   }
 68 | }
 69 | 
 70 | constexpr auto interleave_bytes2 = interleave_bytes<2>;
 71 | constexpr auto interleave_bytes4 = interleave_bytes<4>;
 72 | 
 73 | 
 74 | void interleave_bytes24(const char *__restrict__ src, ssize_t length, char *__restrict__ out) {
 75 |   const __m256i zeros = _mm256_setzero_si256();
 76 |   __m256i s;
 77 |   ssize_t i;
 78 |   for (i = 0; i < length - 31; i += 32) {
 79 |     s = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(src + i));
 80 |     __m256i tlo = _mm256_permute4x64_epi64(s, 0b00010100);
 81 |   	tlo = _mm256_unpacklo_epi16(tlo, zeros);
 82 |   	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * 2), tlo);
 83 |   	__m256i thi = _mm256_permute4x64_epi64(s, 0b11101011);
 84 |   	thi = _mm256_unpackhi_epi16(thi, zeros);
 85 |   	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * 2 + 32), thi);
 86 |   }
 87 |   if (i < length - 15) {
 88 | 	const __m128i *head = reinterpret_cast<const __m128i *>(src + i);
 89 | 	s = _mm256_loadu2_m128i(head, head);
 90 | 	__m256i tlo = _mm256_permute4x64_epi64(s, 0b00010100);
 91 |   	tlo = _mm256_unpacklo_epi16(tlo, zeros);
 92 |   	_mm256_storeu_si256(reinterpret_cast<__m256i *>(out + i * 2), tlo);
 93 | 	i += 16;
 94 |   }
 95 |   for (; i < length - 3; i += 4) {
 96 |   	uint64_t pair = *reinterpret_cast<const uint32_t *>(src + i);
 97 | 	*reinterpret_cast<uint64_t *>(out + i * 2) = (pair & 0xFFFF) | ((pair & 0xFFFF0000) << 16);
 98 |   }
 99 |   if (i < length) {
100 |     *reinterpret_cast<uint32_t *>(out + i * 2) = static_cast<const uint32_t>(*reinterpret_cast<const uint16_t *>(src + i));
101 |   }
102 | }


--------------------------------------------------------------------------------
/2023/pyx/utils_accelerated.pyx:
--------------------------------------------------------------------------------
  1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
  2 | # cython: warn.maybe_uninitialized=True
  3 | # distutils: language = c++
  4 | # distutils: extra_compile_args = -std=c++17 -mavx2
  5 | 
  6 | from typing import Any, Optional
  7 | 
  8 | from athenian.api.internal.settings import ReleaseMatch, ReleaseSettings, default_branch_alias
  9 | 
 10 | from cpython cimport PyObject
 11 | from cpython.bytes cimport PyBytes_FromStringAndSize
 12 | from cpython.dict cimport PyDict_GetItem
 13 | from cpython.unicode cimport PyUnicode_GET_LENGTH
 14 | from cython.operator cimport dereference as deref
 15 | from libc.string cimport memchr, memcmp, memcpy
 16 | from libcpp.memory cimport allocator, unique_ptr
 17 | 
 18 | from athenian.api.native.cpython cimport PyUnicode_DATA, PyUnicode_KIND
 19 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport (
 20 |     empty_deleter,
 21 |     mi_heap_destroy_stl_allocator,
 22 |     mi_unordered_map,
 23 |     mi_vector,
 24 | )
 25 | from athenian.api.native.optional cimport optional
 26 | from athenian.api.native.string_view cimport string_view
 27 | 
 28 | 
 29 | cdef extern from "string.h" nogil:
 30 |     void *memmem(
 31 |         const void *haystack, size_t haystacklen,
 32 |         const void *needle, size_t needlelen,
 33 |     )
 34 | 
 35 | cdef extern from "utils_accelerated.h" nogil:
 36 |     void interleave_bytes2(const char *src, size_t length, char *out)
 37 |     void interleave_bytes4(const char *src, size_t length, char *out)
 38 |     void interleave_bytes24(const char *src, size_t length, char *out)
 39 | 
 40 | cdef:
 41 |     mi_heap_destroy_stl_allocator[char] alloc
 42 |     optional[mi_vector[unique_ptr[char[], empty_deleter]]] _str_storage
 43 | 
 44 | _str_storage.emplace(alloc)
 45 | 
 46 | cdef:
 47 |     Py_ssize_t default_branch_alias_len = len(default_branch_alias)
 48 |     const char *default_branch_alias_data1 = <const char *> PyUnicode_DATA(
 49 |         <PyObject *> default_branch_alias
 50 |     )
 51 |     char *default_branch_alias_data2 = <char *> deref(_str_storage).emplace_back(alloc.allocate(default_branch_alias_len * 2)).get()
 52 |     char *default_branch_alias_data4 = <char *> deref(_str_storage).emplace_back(alloc.allocate(default_branch_alias_len * 4)).get()
 53 |     long branch = ReleaseMatch.branch
 54 |     long tag = ReleaseMatch.tag
 55 |     long tag_or_branch = ReleaseMatch.tag_or_branch
 56 |     long event = ReleaseMatch.event
 57 |     const char *tag_or_branch_data = <const char *>PyUnicode_DATA(<PyObject *> ReleaseMatch.tag_or_branch.name)
 58 |     Py_ssize_t tag_or_branch_name_len = len(ReleaseMatch.tag_or_branch.name)
 59 |     const char *rejected_data = <const char *>PyUnicode_DATA(<PyObject *> ReleaseMatch.rejected.name)
 60 |     Py_ssize_t rejected_name_len = len(ReleaseMatch.rejected.name)
 61 |     const char *force_push_drop_data = <const char *>PyUnicode_DATA(<PyObject *> ReleaseMatch.force_push_drop.name)
 62 |     Py_ssize_t force_push_drop_name_len = len(ReleaseMatch.force_push_drop.name)
 63 |     optional[mi_unordered_map[string_view, long]] release_match_name_to_enum
 64 | 
 65 | interleave_bytes2(
 66 |     default_branch_alias_data1,
 67 |     default_branch_alias_len,
 68 |     default_branch_alias_data2,
 69 | )
 70 | interleave_bytes4(
 71 |     default_branch_alias_data1,
 72 |     default_branch_alias_len,
 73 |     default_branch_alias_data4,
 74 | )
 75 | 
 76 | release_match_name_to_enum.emplace(alloc)
 77 | for obj in ReleaseMatch:
 78 |     deref(release_match_name_to_enum)[string_view(
 79 |         <char *>PyUnicode_DATA(<PyObject *> obj.name), PyUnicode_GET_LENGTH(obj.name)
 80 |     )] = obj
 81 |     deref(_str_storage).emplace_back(alloc.allocate(PyUnicode_GET_LENGTH(obj.name) * 2))
 82 |     interleave_bytes2(
 83 |         <char *>PyUnicode_DATA(<PyObject *> obj.name),
 84 |         PyUnicode_GET_LENGTH(obj.name),
 85 |         <char *> deref(_str_storage).back().get(),
 86 |     )
 87 |     deref(release_match_name_to_enum)[string_view(
 88 |         <char *> deref(_str_storage).back().get(), PyUnicode_GET_LENGTH(obj.name) * 2
 89 |     )] = obj
 90 |     deref(_str_storage).emplace_back(alloc.allocate(PyUnicode_GET_LENGTH(obj.name) * 4))
 91 |     interleave_bytes4(
 92 |         <char *> PyUnicode_DATA(<PyObject *> obj.name),
 93 |         PyUnicode_GET_LENGTH(obj.name),
 94 |         <char *> deref(_str_storage).back().get(),
 95 |     )
 96 |     deref(release_match_name_to_enum)[string_view(
 97 |         <char *> deref(_str_storage).back().get(), PyUnicode_GET_LENGTH(obj.name) * 4
 98 |     )] = obj
 99 | 
100 | 
101 | def interleave_expand(bytes src, int srckind, int dstkind):
102 |     """
103 |     srckind > dstkind
104 |     dstkind=2: abcd -> a0b0c0d0 (ucs1 -> ucs2)
105 |     srckind=1 dstkind=4: abcd -> a000b000c000d000 (ucs1 -> ucs4)
106 |     srckind=2 dstkind=4: a0b0c0d0 -> a000b000c000d000 (ucs2 -> ucs4)
107 |     """
108 |     assert srckind < dstkind
109 |     cdef:
110 |         bytes output = PyBytes_FromStringAndSize(NULL, len(src) << ((dstkind >> 1) - (srckind >> 1)))
111 |     if dstkind == 2:
112 |         interleave_bytes2(src, len(src), output)
113 |     elif dstkind == 4:
114 |         if srckind == 1:
115 |             interleave_bytes4(src, len(src), output)
116 |         else:
117 |             interleave_bytes24(src, len(src), output)
118 |     return output
119 | 
120 | 
121 | def triage_by_release_match(
122 |     repo: str,
123 |     release_match: str,
124 |     release_settings: ReleaseSettings,
125 |     default_branches: dict[str, str],
126 |     result: Any,
127 |     ambiguous: dict[str, Any],
128 | ) -> Optional[Any]:
129 |     """Check the release match of the specified `repo` and return `None` if it is not effective \
130 |     according to `release_settings`, or decide between `result` and `ambiguous`."""
131 |     cdef:
132 |         Py_ssize_t str_kind = PyUnicode_KIND(<PyObject *> release_match)
133 |         Py_ssize_t release_match_len = PyUnicode_GET_LENGTH(release_match) * str_kind
134 |         const char *release_match_data = <const char *>PyUnicode_DATA(<PyObject *> release_match)
135 |         PyObject *required_release_match
136 |         const char *match_name
137 |         int match_name_len
138 |         const char *match_by
139 |         int match_by_len
140 |         long match
141 |         long required_release_match_match
142 |         const char *target_data
143 |         Py_ssize_t target_len, target_kind
144 |         PyObject *default_branch
145 |         const char *default_branch_data
146 |         Py_ssize_t default_branch_len, default_branch_kind, target_kind_shift
147 |         const char *default_branch_alias_data
148 |         const char *found
149 |         unique_ptr[char] resolved_branch
150 |         Py_ssize_t pos
151 |     if (
152 |         (
153 |             release_match_len == rejected_name_len
154 |             and memcmp(release_match_data, rejected_data, rejected_name_len) == 0
155 |         )
156 |         or
157 |         (
158 |             release_match_len == force_push_drop_name_len
159 |             and memcmp(release_match_data, force_push_drop_data, force_push_drop_name_len) == 0
160 |         )
161 |     ):
162 |         return result
163 | 
164 |     required_release_match = PyDict_GetItem(release_settings.native, repo)
165 |     if required_release_match == NULL:
166 |         # DEV-1451: if we don't have this repository in the release settings, then it is deleted
167 |         raise AssertionError(
168 |             f"You must take care of deleted repositories separately: {repo}",
169 |         ) from None
170 |     match_name = <const char *> memchr(release_match_data, ord(b"|"), release_match_len)
171 |     if match_name == NULL:
172 |         match_name_len = release_match_len
173 |         match_by = release_match_data + release_match_len
174 |         match_by_len = 0
175 |     else:
176 |         match_name_len = match_name - release_match_data
177 |         match_by = match_name + str_kind
178 |         match_by_len = release_match_len - match_name_len - str_kind
179 |     match_name = release_match_data
180 |     match = deref(release_match_name_to_enum)[string_view(match_name, match_name_len)]
181 |     required_release_match_match = (<object>required_release_match).match
182 |     if required_release_match_match != tag_or_branch:
183 |         if match != required_release_match_match:
184 |             return None
185 |         dump = result
186 |     else:
187 |         if memcmp(match_name, b"event", 5) == 0:
188 |             return None
189 |         match_name_len >>= (str_kind >> 1)
190 |         dump = ambiguous[release_match[:match_name_len]]
191 |     if match == tag:
192 |         target = (<object>required_release_match).tags
193 |     elif match == branch:
194 |         target = (<object>required_release_match).branches
195 |     elif match == event:
196 |         target = (<object>required_release_match).events
197 |     else:
198 |         raise AssertionError("Precomputed DB may not contain Match.tag_or_branch")
199 |     target_data = <const char *> PyUnicode_DATA(<PyObject *> target)
200 |     target_kind = PyUnicode_KIND(<PyObject *> target)
201 |     target_len = PyUnicode_GET_LENGTH(target) << (target_kind >> 1)
202 |     if match == branch:
203 |         if target_kind == 1:
204 |             default_branch_alias_data = default_branch_alias_data1
205 |         elif target_kind == 2:
206 |             default_branch_alias_data = default_branch_alias_data2
207 |         else:
208 |             default_branch_alias_data = default_branch_alias_data4
209 |         found = <const char *>memmem(
210 |             target_data,
211 |             target_len,
212 |             default_branch_alias_data,
213 |             default_branch_alias_len << (target_kind >> 1),
214 |         )
215 |         if found != NULL:
216 |             target_len -= default_branch_alias_len * target_kind
217 |             default_branch = PyDict_GetItem(default_branches, repo)
218 |             default_branch_len = PyUnicode_GET_LENGTH(<object> default_branch)
219 |             default_branch_kind = PyUnicode_KIND(default_branch)
220 |             if target_kind == default_branch_kind:
221 |                 target_kind_shift = 0
222 |                 target_len += default_branch_len << (default_branch_kind >> 1)
223 |             elif target_kind > default_branch_kind:
224 |                 target_kind_shift = 0
225 |                 target_len += default_branch_len << ((target_kind >> 1) - (default_branch_kind >> 1))
226 |             else:
227 |                 target_kind_shift = (default_branch_kind >> 1) - (target_kind >> 1)
228 |                 target_len <<= target_kind_shift
229 |                 target_len += default_branch_len << (default_branch_kind >> 1)
230 |             if target_len != match_by_len:
231 |                 return None
232 |             default_branch_data = <const char *> PyUnicode_DATA(default_branch)
233 |             resolved_branch.reset(allocator[char]().allocate(target_len))
234 |             pos = found - target_data
235 |             if target_kind >= default_branch_kind:
236 |                 memcpy(resolved_branch.get(), target_data, pos)
237 |             elif target_kind == 1:
238 |                 if default_branch_kind == 2:
239 |                     interleave_bytes2(target_data, pos, resolved_branch.get())
240 |                 else:
241 |                     interleave_bytes4(target_data, pos, resolved_branch.get())
242 |             else:
243 |                 # target_kind == 2
244 |                 interleave_bytes24(target_data, pos, resolved_branch.get())
245 |             pos <<= target_kind_shift
246 |             if default_branch_kind >= target_kind:
247 |                 memcpy(resolved_branch.get() + pos, default_branch_data, default_branch_len << (default_branch_kind >> 1))
248 |                 pos += default_branch_len << (default_branch_kind >> 1)
249 |             else:
250 |                 if default_branch_kind == 1:
251 |                     if target_kind == 2:
252 |                         interleave_bytes2(default_branch_data, default_branch_len, resolved_branch.get() + pos)
253 |                     else:
254 |                         interleave_bytes4(default_branch_data, default_branch_len, resolved_branch.get() + pos)
255 |                 else:
256 |                     # default_branch_kind = 2
257 |                     interleave_bytes24(default_branch_data, default_branch_len * 2, resolved_branch.get() + pos)
258 |                 pos += default_branch_len << (target_kind >> 1)
259 |             if target_kind >= default_branch_kind:
260 |                 memcpy(resolved_branch.get() + pos, found + default_branch_alias_len * target_kind, target_len - pos)
261 |             elif target_kind == 1:
262 |                 if default_branch_kind == 2:
263 |                     interleave_bytes2(found + default_branch_alias_len, target_len - pos, resolved_branch.get() + pos)
264 |                 else:
265 |                     interleave_bytes4(found + default_branch_alias_len, target_len - pos, resolved_branch.get() + pos)
266 |             else:
267 |                 # target_kind == 2
268 |                 interleave_bytes24(found + default_branch_alias_len * 2, (target_len - pos) >> 1, resolved_branch.get() + pos)
269 |             target_data = resolved_branch.get()
270 | 
271 |     if target_len != match_by_len or memcmp(target_data, match_by, match_by_len):
272 |         return None
273 |     return dump
274 | 


--------------------------------------------------------------------------------
/2023/pyx/web_model_io.h:
--------------------------------------------------------------------------------
 1 | // copied private functions from numpy, licensed under BSD-3
 2 | 
 3 | namespace {
 4 | /*
 5 |  * Computes the python `ret, d = divmod(d, unit)`.
 6 |  *
 7 |  * Note that GCC is smart enough at -O2 to eliminate the `if(*d < 0)` branch
 8 |  * for subsequent calls to this command - it is able to deduce that `*d >= 0`.
 9 |  */
10 | inline npy_int64 extract_unit_64(npy_int64 *d, npy_int64 unit) {
11 |     assert(unit > 0);
12 |     npy_int64 div = *d / unit;
13 |     npy_int64 mod = *d % unit;
14 |     if (mod < 0) {
15 |         mod += unit;
16 |         div -= 1;
17 |     }
18 |     assert(mod >= 0);
19 |     *d = mod;
20 |     return div;
21 | }
22 | 
23 | inline int is_leapyear(npy_int64 year) {
24 |     return (year & 0x3) == 0 && /* year % 4 == 0 */
25 |            ((year % 100) != 0 ||
26 |             (year % 400) == 0);
27 | }
28 | 
29 | int _days_per_month_table[2][12] = {
30 |     { 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 },
31 |     { 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 },
32 | };
33 | 
34 | /*
35 |  * Modifies '*days_' to be the day offset within the year,
36 |  * and returns the year.
37 |  */
38 | inline npy_int64 days_to_yearsdays(npy_int64 *days_) {
39 |     const npy_int64 days_per_400years = (400*365 + 100 - 4 + 1);
40 |     /* Adjust so it's relative to the year 2000 (divisible by 400) */
41 |     npy_int64 days = (*days_) - (365*30 + 7);
42 | 
43 |     /* Break down the 400 year cycle to get the year and day within the year */
44 |     npy_int64 year = 400 * extract_unit_64(&days, days_per_400years);
45 | 
46 |     /* Work out the year/day within the 400 year cycle */
47 |     if (days >= 366) {
48 |         year += 100 * ((days-1) / (100*365 + 25 - 1));
49 |         days = (days-1) % (100*365 + 25 - 1);
50 |         if (days >= 365) {
51 |             year += 4 * ((days+1) / (4*365 + 1));
52 |             days = (days+1) % (4*365 + 1);
53 |             if (days >= 366) {
54 |                 year += (days-1) / 365;
55 |                 days = (days-1) % 365;
56 |             }
57 |         }
58 |     }
59 | 
60 |     *days_ = days;
61 |     return year + 2000;
62 | }
63 | 
64 | /*
65 |  * Fills in the year, month, day in 'dts' based on the days
66 |  * offset from 1970.
67 |  */
68 | inline void set_datetimestruct_days(npy_int64 days, int *year, int *month, int *day) {
69 |     *year = days_to_yearsdays(&days);
70 |     int *month_lengths = _days_per_month_table[is_leapyear(*year)];
71 | 
72 |     for (int i = 0; i < 12; ++i) {
73 |         if (days < month_lengths[i]) {
74 |             *month = i + 1;
75 |             *day = (int)days + 1;
76 |             return;
77 |         }
78 |         else {
79 |             days -= month_lengths[i];
80 |         }
81 |     }
82 | }
83 | 
84 | }
85 | 


--------------------------------------------------------------------------------
/2023/pyx/web_model_io.pyx:
--------------------------------------------------------------------------------
   1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
   2 | # cython: warn.maybe_uninitialized=False
   3 | # distutils: language = c++
   4 | # distutils: extra_compile_args = -mavx2 -ftree-vectorize -std=c++20
   5 | # distutils: libraries = mimalloc
   6 | # distutils: runtime_library_dirs = /usr/local/lib
   7 | 
   8 | cimport cython
   9 | 
  10 | from cython.operator import dereference
  11 | 
  12 | from cpython cimport (
  13 |     Py_INCREF,
  14 |     PyBytes_FromStringAndSize,
  15 |     PyDict_New,
  16 |     PyDict_SetItem,
  17 |     PyFloat_FromDouble,
  18 |     PyList_New,
  19 |     PyList_SET_ITEM,
  20 |     PyLong_FromLong,
  21 |     PyObject,
  22 |     PyTuple_New,
  23 |     PyTuple_SET_ITEM,
  24 | )
  25 | from cpython.datetime cimport PyDateTimeAPI, import_datetime
  26 | from cpython.dict cimport PyDict_GetItemString
  27 | from libc.stdint cimport int32_t, uint8_t, uint16_t, uint32_t
  28 | from libc.stdio cimport FILE, SEEK_CUR, fclose, fread, fseek, ftell
  29 | from libc.string cimport memcpy, strlen
  30 | from numpy cimport import_array, npy_int64
  31 | 
  32 | from athenian.api.native.chunked_stream cimport chunked_stream
  33 | from athenian.api.native.cpython cimport (
  34 |     Py_False,
  35 |     Py_None,
  36 |     Py_True,
  37 |     Py_TYPE,
  38 |     PyBaseObject_Type,
  39 |     PyBool_Type,
  40 |     PyBytes_AS_STRING,
  41 |     PyBytes_Check,
  42 |     PyBytes_GET_SIZE,
  43 |     PyDateTime_CAPI,
  44 |     PyDateTime_Check,
  45 |     PyDateTime_DATE_GET_HOUR,
  46 |     PyDateTime_DATE_GET_MINUTE,
  47 |     PyDateTime_DATE_GET_SECOND,
  48 |     PyDateTime_DATE_GET_TZINFO,
  49 |     PyDateTime_DELTA_GET_DAYS,
  50 |     PyDateTime_DELTA_GET_SECONDS,
  51 |     PyDateTime_GET_DAY,
  52 |     PyDateTime_GET_MONTH,
  53 |     PyDateTime_GET_YEAR,
  54 |     PyDelta_Check,
  55 |     PyDict_CheckExact,
  56 |     PyDict_Next,
  57 |     PyDict_Size,
  58 |     PyDict_Type,
  59 |     PyFloat_AS_DOUBLE,
  60 |     PyFloat_CheckExact,
  61 |     PyFloat_Type,
  62 |     PyList_CheckExact,
  63 |     PyList_GET_ITEM,
  64 |     PyList_GET_SIZE,
  65 |     PyList_Type,
  66 |     PyLong_AsLong,
  67 |     PyLong_CheckExact,
  68 |     PyLong_Type,
  69 |     PyMemberDef,
  70 |     PyObject_TypeCheck,
  71 |     PyTuple_GET_ITEM,
  72 |     PyTypeObject,
  73 |     PyUnicode_1BYTE_KIND,
  74 |     PyUnicode_2BYTE_KIND,
  75 |     PyUnicode_4BYTE_KIND,
  76 |     PyUnicode_Check,
  77 |     PyUnicode_DATA,
  78 |     PyUnicode_FromKindAndData,
  79 |     PyUnicode_FromString,
  80 |     PyUnicode_GET_LENGTH,
  81 |     PyUnicode_KIND,
  82 |     PyUnicode_Type,
  83 | )
  84 | from athenian.api.native.mi_heap_destroy_stl_allocator cimport (
  85 |     mi_heap_allocator_from_capsule,
  86 |     mi_heap_destroy_stl_allocator,
  87 |     mi_vector,
  88 | )
  89 | from athenian.api.native.numpy cimport (
  90 |     NPY_DATETIMEUNIT,
  91 |     NPY_FR_ns,
  92 |     NPY_FR_s,
  93 |     NPY_FR_us,
  94 |     PyArray_CheckExact,
  95 |     PyArray_DATA,
  96 |     PyArray_DIM,
  97 |     PyArray_IS_C_CONTIGUOUS,
  98 |     PyArray_NDIM,
  99 |     PyArray_ScalarAsCtype,
 100 |     PyDatetimeArrType_Type,
 101 |     PyDatetimeScalarObject,
 102 |     PyDoubleArrType_Type,
 103 |     PyFloatArrType_Type,
 104 |     PyIntegerArrType_Type,
 105 |     PyTimedeltaArrType_Type,
 106 | )
 107 | from athenian.api.native.optional cimport optional
 108 | from athenian.api.native.utf8 cimport ucs4_to_utf8_json
 109 | 
 110 | import pickle
 111 | from types import GenericAlias
 112 | 
 113 | from athenian.api.typing_utils import is_generic, is_optional
 114 | 
 115 | 
 116 | cdef extern from "stdio.h" nogil:
 117 |     FILE *fmemopen(void *buf, size_t size, const char *mode)
 118 | 
 119 | 
 120 | cdef extern from "<stdlib.h>" nogil:
 121 |     char *gcvt(double number, int ndigit, char *buf)
 122 | 
 123 | 
 124 | cdef extern from "web_model_io.h" nogil:
 125 |     void set_datetimestruct_days(npy_int64 days, int *year, int *month, int *day)
 126 | 
 127 | 
 128 | import_datetime()
 129 | import_array()
 130 | 
 131 | 
 132 | cdef enum DataType:
 133 |     DT_INVALID = 0
 134 |     DT_MODEL = 1
 135 |     DT_LIST = 2
 136 |     DT_DICT = 3
 137 |     DT_LONG = 4
 138 |     DT_FLOAT = 5
 139 |     DT_STRING = 6
 140 |     DT_DT = 7
 141 |     DT_TD = 8
 142 |     DT_BOOL = 9
 143 |     DT_FREEFORM = 10
 144 | 
 145 | 
 146 | cdef enum DataFlags:
 147 |     DF_KEY_UNMAPPED = 1
 148 |     DF_OPTIONAL = 2
 149 |     DF_VERBATIM = 4
 150 | 
 151 | 
 152 | ctypedef struct SpecNode:
 153 |     DataType type
 154 |     uint8_t flags
 155 |     Py_ssize_t offset
 156 |     const void *key
 157 |     PyTypeObject *model
 158 |     optional[mi_vector[SpecNode]] nested
 159 | 
 160 | 
 161 | cdef inline DataType _discover_data_type(
 162 |     PyTypeObject *obj,
 163 |     const char *key,
 164 |     PyTypeObject **deref,
 165 |     bint *optional,
 166 |     bint *verbatim,
 167 | ) except DT_INVALID:
 168 |     if is_optional(<object> obj):
 169 |         optional[0] = 1
 170 |         verbatim[0] = hasattr((<object> obj).__origin__, "__verbatim__")
 171 |         args = (<object> obj).__args__
 172 |         obj = <PyTypeObject *> PyTuple_GET_ITEM(<PyObject *> args, 0)
 173 |     if obj == &PyLong_Type:
 174 |         return DT_LONG
 175 |     elif obj == &PyFloat_Type:
 176 |         return DT_FLOAT
 177 |     elif obj == &PyUnicode_Type:
 178 |         return DT_STRING
 179 |     elif obj == PyDateTimeAPI.DateTimeType:
 180 |         return DT_DT
 181 |     elif obj == PyDateTimeAPI.DeltaType:
 182 |         return DT_TD
 183 |     elif obj == &PyBool_Type:
 184 |         return DT_BOOL
 185 |     elif is_generic(<object> obj):
 186 |         origin = (<object> obj).__origin__
 187 |         args = (<object> obj).__args__
 188 |         deref[0] = obj
 189 |         if <PyTypeObject *> origin == &PyList_Type:
 190 |             return DT_LIST
 191 |         elif <PyTypeObject *> origin == &PyDict_Type:
 192 |             return DT_DICT
 193 |         else:
 194 |             return DT_INVALID
 195 |     elif hasattr(<object> obj, "attribute_types"):
 196 |         deref[0] = obj
 197 |         return DT_MODEL
 198 |     elif obj == &PyDict_Type or obj == &PyBaseObject_Type:
 199 |         return DT_FREEFORM
 200 |     else:
 201 |         raise AssertionError(f"{'Optional f' if optional[0] else 'F'}ield `{PyUnicode_FromString(key)}` type is not supported: {<object> obj}")
 202 | 
 203 | 
 204 | cdef inline void _apply_data_type(
 205 |     Py_ssize_t offset,
 206 |     const char *key,
 207 |     PyTypeObject *member_type,
 208 |     SpecNode *fields,
 209 |     mi_heap_destroy_stl_allocator[char] &alloc,
 210 | ) except *:
 211 |     cdef:
 212 |         PyTypeObject *deref = NULL
 213 |         bint optional = 0, verbatim = 0
 214 |         DataType dtype = _discover_data_type(member_type, key, &deref, &optional, &verbatim)
 215 |         SpecNode *back = &dereference(fields.nested).emplace_back()
 216 |     back.type = dtype
 217 |     back.offset = offset
 218 |     if optional:
 219 |         back.flags |= DF_OPTIONAL
 220 |     if verbatim:
 221 |         back.flags |= DF_VERBATIM
 222 |     back.nested.emplace(alloc)
 223 |     if deref != NULL:
 224 |         _discover_fields(deref, back, alloc)
 225 | 
 226 | 
 227 | cdef void _discover_fields(
 228 |     PyTypeObject *model,
 229 |     SpecNode *fields,
 230 |     mi_heap_destroy_stl_allocator[char] &alloc,
 231 | ) except *:
 232 |     cdef:
 233 |         object attribute_types
 234 |         object attribute_map
 235 |         PyTypeObject *member_type
 236 |         PyMemberDef *members
 237 |         PyObject *key
 238 |         SpecNode *back
 239 | 
 240 |     if fields.type == DT_MODEL:
 241 |         attribute_types = (<object> model).attribute_types
 242 |         attribute_map = (<object> model).attribute_map
 243 |         fields.model = model
 244 |         members = model.tp_members
 245 |         i = 0
 246 |         while members[i].name != NULL:
 247 |             member_type = <PyTypeObject *> PyDict_GetItemString(attribute_types, members[i].name + 1)
 248 |             _apply_data_type(members[i].offset, members[i].name + 1, member_type, fields, alloc)
 249 |             back = &dereference(fields.nested).back()
 250 |             key = PyDict_GetItemString(attribute_map, members[i].name + 1)
 251 |             if key != NULL:
 252 |                 back.flags &= ~DF_KEY_UNMAPPED
 253 |                 back.key = key
 254 |             else:
 255 |                 back.flags |= DF_KEY_UNMAPPED
 256 |                 back.key = members[i].name + 1
 257 |             i += 1
 258 |     elif fields.type == DT_LIST:
 259 |         attribute_types = (<object> model).__args__
 260 |         _apply_data_type(
 261 |             0,
 262 |             b"list",
 263 |             <PyTypeObject *> PyTuple_GET_ITEM(<PyObject *> attribute_types, 0),
 264 |             fields,
 265 |             alloc,
 266 |         )
 267 |     elif fields.type == DT_DICT:
 268 |         attribute_types = (<object> model).__args__
 269 |         _apply_data_type(
 270 |             0,
 271 |             b"dict key",
 272 |             <PyTypeObject *> PyTuple_GET_ITEM(<PyObject *> attribute_types, 0),
 273 |             fields,
 274 |             alloc,
 275 |         )
 276 |         _apply_data_type(
 277 |             0,
 278 |             b"dict value",
 279 |             <PyTypeObject *> PyTuple_GET_ITEM(<PyObject *> attribute_types, 1),
 280 |             fields,
 281 |             alloc,
 282 |         )
 283 |     else:
 284 |         raise AssertionError(f"Cannot recurse in dtype {fields.type}")
 285 | 
 286 | 
 287 | @cython.cdivision(True)
 288 | cdef PyObject *_write_object(PyObject *obj, SpecNode *spec, chunked_stream &stream) nogil:
 289 |     cdef:
 290 |         char dtype = spec.type, bool
 291 |         long val_long
 292 |         double val_double
 293 |         float val_float
 294 |         uint32_t str_length, val32, i
 295 |         uint16_t val16[4]
 296 |         int32_t vali32
 297 |         PyObject *exc
 298 |         bint is_unicode, is_float
 299 |         NPY_DATETIMEUNIT npy_unit
 300 |         npy_int64 obval
 301 |         Py_ssize_t dict_pos = 0
 302 |         PyObject *dict_key = NULL
 303 |         PyObject *dict_val = NULL
 304 |         PyObject **npdata
 305 |         SpecNode *field
 306 |     if obj == Py_None:
 307 |         dtype = 0
 308 |         stream.write(<char *> &dtype, 1)
 309 |         return NULL
 310 |     stream.write(<char *> &dtype, 1)
 311 |     if dtype == DT_LONG:
 312 |         if PyLong_CheckExact(obj):
 313 |             val_long = PyLong_AsLong(obj)
 314 |         elif PyObject_TypeCheck(obj, &PyIntegerArrType_Type):
 315 |             val_long = 0
 316 |             PyArray_ScalarAsCtype(obj, &val_long)
 317 |         else:
 318 |             return obj
 319 |         stream.write(<char *> &val_long, sizeof(long))
 320 |     elif dtype == DT_FLOAT:
 321 |         if PyFloat_CheckExact(obj):
 322 |             val_double = PyFloat_AS_DOUBLE(obj)
 323 |         elif PyLong_CheckExact(obj):
 324 |             val_double = PyLong_AsLong(obj)
 325 |         elif PyObject_TypeCheck(obj, &PyDoubleArrType_Type):
 326 |             PyArray_ScalarAsCtype(obj, &val_double)
 327 |         elif PyObject_TypeCheck(obj, &PyFloatArrType_Type):
 328 |             PyArray_ScalarAsCtype(obj, &val_float)
 329 |             val_double = val_float
 330 |         elif PyObject_TypeCheck(obj, &PyIntegerArrType_Type):
 331 |             val_long = 0
 332 |             PyArray_ScalarAsCtype(obj, &val_long)
 333 |             val_double = val_long
 334 |         else:
 335 |             return obj
 336 |         stream.write(<char *> &val_double, sizeof(double))
 337 |     elif dtype == DT_STRING:
 338 |         is_unicode = PyUnicode_Check(obj)
 339 |         if not is_unicode and not PyBytes_Check(obj):
 340 |             return obj
 341 |         if is_unicode:
 342 |             str_length = PyUnicode_GET_LENGTH(obj)
 343 |             val32 = str_length | ((PyUnicode_KIND(obj) - 1) << 30)
 344 |             stream.write(<char *> &val32, 4)
 345 |             # each code point in PyUnicode_DATA buffer has PyUnicode_KIND(obj) bytes
 346 |             stream.write(<char *> PyUnicode_DATA(obj), PyUnicode_KIND(obj) * str_length)
 347 |         else:
 348 |             val32 = PyBytes_GET_SIZE(obj)
 349 |             stream.write(<char *> &val32, 4)
 350 |             stream.write(PyBytes_AS_STRING(obj), val32)
 351 |     elif dtype == DT_DT:
 352 |         if not PyDateTime_Check(obj):
 353 |             if PyObject_TypeCheck(obj, &PyDatetimeArrType_Type):
 354 |                 npy_unit = (<PyDatetimeScalarObject *> obj).obmeta.base
 355 |                 obval = (<PyDatetimeScalarObject *> obj).obval
 356 |                 if npy_unit == NPY_FR_ns:
 357 |                     obval //= 1000000000
 358 |                 elif npy_unit == NPY_FR_us:
 359 |                     obval //= 1000000
 360 |                 elif npy_unit != NPY_FR_s:
 361 |                     return obj
 362 |                 memcpy(val16, &obval, 8)   # little-endian
 363 |             else:
 364 |                 return obj
 365 |         else:
 366 |             val16[0] = PyDateTime_GET_YEAR(obj) << 4
 367 |             val16[0] |= PyDateTime_GET_MONTH(obj)
 368 |             val16[1] = PyDateTime_GET_DAY(obj) << 7
 369 |             val16[1] |= PyDateTime_DATE_GET_HOUR(obj)
 370 |             val16[2] = (PyDateTime_DATE_GET_MINUTE(obj) << 8) | 0x8000
 371 |             val16[2] |= PyDateTime_DATE_GET_SECOND(obj)
 372 |         stream.write(<char *> val16, 2 * 3)
 373 |     elif dtype == DT_TD:
 374 |         if not PyDelta_Check(obj):
 375 |             if PyObject_TypeCheck(obj, &PyTimedeltaArrType_Type):
 376 |                 npy_unit = (<PyDatetimeScalarObject *> obj).obmeta.base
 377 |                 obval = (<PyDatetimeScalarObject *> obj).obval
 378 |                 if npy_unit == NPY_FR_ns:
 379 |                     obval //= 1000000000
 380 |                 elif npy_unit == NPY_FR_us:
 381 |                     obval //= 1000000
 382 |                 elif npy_unit != NPY_FR_s:
 383 |                     return obj
 384 |                 if obval >= 0:
 385 |                     vali32 = obval // (24 * 3600)
 386 |                     var_long = obval % (24 * 3600)
 387 |                 else:
 388 |                     vali32 = -1 -(obval // (24 * 3600))
 389 |                     var_long = 24 * 3600 + obval % (24 * 3600)
 390 |             else:
 391 |                 return obj
 392 |         else:
 393 |             vali32 = PyDateTime_DELTA_GET_DAYS(obj)
 394 |             var_long = PyDateTime_DELTA_GET_SECONDS(obj)
 395 |         vali32 <<= 1
 396 |         if var_long >= 1 << 16:
 397 |             vali32 |= 1
 398 |         val16[0] = var_long & 0xFFFF
 399 |         stream.write(<char *> &vali32, 4)
 400 |         stream.write(<char *> val16, 2)
 401 |     elif dtype == DT_BOOL:
 402 |         bool = obj == Py_True
 403 |         if not bool and obj != Py_False:
 404 |             return obj
 405 |         stream.write(<char *> &bool, 1)
 406 |     elif dtype == DT_LIST:
 407 |         if not PyList_CheckExact(obj):
 408 |             if not PyArray_CheckExact(obj) or not PyArray_IS_C_CONTIGUOUS(obj) or PyArray_NDIM(obj) != 1:
 409 |                 return obj
 410 |             val32 = PyArray_DIM(obj, 0)
 411 |             stream.write(<char *> &val32, 4)
 412 |             npdata = <PyObject **> PyArray_DATA(obj)
 413 |             for i in range(val32):
 414 |                 exc = _write_object(npdata[i], &dereference(spec.nested)[0], stream)
 415 |                 if exc != NULL:
 416 |                     return exc
 417 |         else:
 418 |             val32 = PyList_GET_SIZE(obj)
 419 |             stream.write(<char *> &val32, 4)
 420 |             for i in range(val32):
 421 |                 exc = _write_object(PyList_GET_ITEM(obj, i), &dereference(spec.nested)[0], stream)
 422 |                 if exc != NULL:
 423 |                     return exc
 424 |     elif dtype == DT_DICT:
 425 |         if not PyDict_CheckExact(obj):
 426 |             return obj
 427 |         val32 = PyDict_Size(obj)
 428 |         stream.write(<char *> &val32, 4)
 429 |         while PyDict_Next(obj, &dict_pos, &dict_key, &dict_val):
 430 |             exc = _write_object(dict_key, &dereference(spec.nested)[0], stream)
 431 |             if exc != NULL:
 432 |                 return exc
 433 |             exc = _write_object(dict_val, &dereference(spec.nested)[1], stream)
 434 |             if exc != NULL:
 435 |                 return exc
 436 |     elif dtype == DT_MODEL:
 437 |         val32 = dereference(spec.nested).size()
 438 |         stream.write(<char *> &val32, 4)
 439 |         for i in range(val32):
 440 |             field = &dereference(spec.nested)[i]
 441 |             exc = _write_object(
 442 |                 dereference(<PyObject **>((<char *> obj) + field.offset)),
 443 |                 field,
 444 |                 stream,
 445 |             )
 446 |             if exc != NULL:
 447 |                 return exc
 448 |     else:
 449 |         return obj
 450 |     return NULL
 451 | 
 452 | 
 453 | cdef void _serialize_list_of_models(
 454 |     list models,
 455 |     chunked_stream &stream,
 456 |     mi_heap_destroy_stl_allocator[char] &alloc,
 457 | ) except *:
 458 |     cdef:
 459 |         uint32_t size
 460 |         SpecNode spec
 461 |         type item_type
 462 |         PyObject *exc
 463 | 
 464 |     spec.type = DT_LIST
 465 |     spec.nested.emplace(alloc)
 466 |     if len(models) == 0:
 467 |         size = 0
 468 |         stream.write(<char *> &size, 4)
 469 |         return
 470 |     item_type = type(models[0])
 471 |     result = pickle.dumps(GenericAlias(list, (item_type,)))
 472 |     _apply_data_type(0, b"root", <PyTypeObject *> item_type, &spec, alloc)
 473 |     with nogil:
 474 |         size = PyBytes_GET_SIZE(<PyObject *> result)
 475 |         stream.write(<char *> &size, 4)
 476 |         stream.write(PyBytes_AS_STRING(<PyObject *> result), size)
 477 |         exc = _write_object(<PyObject *> models, &spec, stream)
 478 |     if exc != NULL:
 479 |         raise ValueError(f"Could not serialize `{<object> exc}` of type {type(<object> exc)} in {item_type.__qualname__}")
 480 | 
 481 | 
 482 | cdef void _serialize_generic(model, chunked_stream &stream) except *:
 483 |     cdef:
 484 |         bytes buf = pickle.dumps(model)
 485 |         uint32_t size = len(buf)
 486 |     stream.write(<char *> &size, 4)
 487 |     stream.write(PyBytes_AS_STRING(<PyObject *> buf), size)
 488 | 
 489 | 
 490 | def serialize_models(tuple models not None, alloc_capsule=None) -> bytes:
 491 |     cdef:
 492 |         optional[chunked_stream] stream
 493 |         bytes result
 494 |         char count
 495 |         optional[mi_heap_destroy_stl_allocator[char]] alloc
 496 |         size_t size
 497 |     assert len(models) < 255
 498 |     if alloc_capsule is not None:
 499 |         alloc.emplace(dereference(mi_heap_allocator_from_capsule(alloc_capsule)))
 500 |     else:
 501 |         alloc.emplace()
 502 |     stream.emplace(dereference(alloc))
 503 |     count = len(models)
 504 |     dereference(stream).write(&count, 1)
 505 |     for model in models:
 506 |         if PyList_CheckExact(<PyObject *> model):
 507 |             _serialize_list_of_models(model, dereference(stream), dereference(alloc))
 508 |         else:
 509 |             _serialize_generic(model, dereference(stream))
 510 |     size = dereference(stream).size()
 511 |     result = PyBytes_FromStringAndSize(NULL, size)
 512 |     dereference(stream).dump(PyBytes_AS_STRING(<PyObject *> result), size)
 513 |     return result
 514 | 
 515 | 
 516 | def deserialize_models(bytes buffer not None, alloc_capsule=None) -> tuple[list[object], ...]:
 517 |     cdef:
 518 |         char *input = PyBytes_AS_STRING(<PyObject *> buffer)
 519 |         uint32_t aux = 0, tuple_pos
 520 |         str corrupted_msg = "Corrupted buffer at position %d: %s"
 521 |         FILE *stream
 522 |         tuple result
 523 |         long pos
 524 |         bytes type_buf
 525 |         object model_type
 526 |         SpecNode spec
 527 |         optional[mi_heap_destroy_stl_allocator[char]] alloc
 528 | 
 529 |     if alloc_capsule is not None:
 530 |         alloc.emplace(dereference(mi_heap_allocator_from_capsule(alloc_capsule)))
 531 |     else:
 532 |         alloc.emplace()
 533 |     stream = fmemopen(input, PyBytes_GET_SIZE(<PyObject *> buffer), b"r")
 534 |     if fread(&aux, 1, 1, stream) != 1:
 535 |         raise ValueError(corrupted_msg % (ftell(stream), "tuple"))
 536 |     result = PyTuple_New(aux)
 537 |     for tuple_pos in range(aux):
 538 |         if fread(&aux, 4, 1, stream) != 1:
 539 |             raise ValueError(corrupted_msg % (ftell(stream), "pickle/header"))
 540 |         if aux == 0:
 541 |             model = []
 542 |         else:
 543 |             pos = ftell(stream)
 544 |             if fseek(stream, aux, SEEK_CUR):
 545 |                 raise ValueError(corrupted_msg % (ftell(stream), "pickle/body"))
 546 |             type_buf = PyBytes_FromStringAndSize(input + pos, aux)
 547 |             model_type = pickle.loads(type_buf)
 548 |             if not isinstance(model_type, (type, GenericAlias)):
 549 |                 model = model_type
 550 |             else:
 551 |                 spec.type = DT_LIST
 552 |                 spec.nested.emplace(dereference(alloc))
 553 |                 _discover_fields(<PyTypeObject *> model_type, &spec, dereference(alloc))
 554 |                 model = _read_model(&spec, stream, input, corrupted_msg)
 555 |         Py_INCREF(model)
 556 |         PyTuple_SET_ITEM(result, tuple_pos, model)
 557 |     fclose(stream)
 558 |     return result
 559 | 
 560 | 
 561 | cdef object _read_model(SpecNode *spec, FILE *stream, const char *raw, str corrupted_msg):
 562 |     cdef:
 563 |         char dtype, bool
 564 |         long val_long
 565 |         double val_double
 566 |         uint32_t aux32, i
 567 |         int32_t auxi32
 568 |         uint16_t aux16[4]
 569 |         unsigned int kind
 570 |         int year, month, day, hour, minute, second
 571 |         PyObject *utctz = (<PyDateTime_CAPI *> PyDateTimeAPI).TimeZone_UTC
 572 |         PyObject *obj_val
 573 |         SpecNode *field
 574 | 
 575 |     if fread(&dtype, 1, 1, stream) != 1:
 576 |         raise ValueError(corrupted_msg % (ftell(stream), "dtype"))
 577 |     if dtype == DT_INVALID:
 578 |         return None
 579 |     if dtype != spec.type:
 580 |         raise ValueError(corrupted_msg % (ftell(stream), f"dtype {dtype} != {spec.type}"))
 581 |     if dtype == DT_LONG:
 582 |         if fread(&val_long, sizeof(long), 1, stream) != 1:
 583 |             raise ValueError(corrupted_msg % (ftell(stream), "long"))
 584 |         return PyLong_FromLong(val_long)
 585 |     elif dtype == DT_FLOAT:
 586 |         if fread(&val_double, sizeof(double), 1, stream) != 1:
 587 |             raise ValueError(corrupted_msg % (ftell(stream), "float"))
 588 |         return PyFloat_FromDouble(val_double)
 589 |     elif dtype == DT_STRING:
 590 |         if fread(&aux32, 4, 1, stream) != 1:
 591 |             raise ValueError(corrupted_msg % (ftell(stream), "str/header"))
 592 |         kind = (aux32 >> 30) + 1
 593 |         aux32 &= 0x3FFFFFFF
 594 |         val_long = ftell(stream)
 595 |         # move stream forward of the number of bytes we are about to read from raw
 596 |         if fseek(stream, aux32 * kind, SEEK_CUR):
 597 |             raise ValueError(corrupted_msg % (ftell(stream), "str/body"))
 598 |         return PyUnicode_FromKindAndData(kind, raw + val_long, aux32)
 599 |     elif dtype == DT_DT:
 600 |         if fread(aux16, 2, 3, stream) != 3:
 601 |             raise ValueError(corrupted_msg % (ftell(stream), "dt"))
 602 |         if aux16[2] & 0x8000:
 603 |             year = aux16[0] >> 4
 604 |             month = aux16[0] & 0xF
 605 |             day = aux16[1] >> 7
 606 |             hour = aux16[1] & 0x7F
 607 |             minute = (aux16[2] >> 8) & 0x7F
 608 |             second = aux16[2] & 0xFF
 609 |         else:
 610 |             aux16[3] = 0
 611 |             obj_val = PyDatetimeArrType_Type.tp_alloc(&PyDatetimeArrType_Type, 0)
 612 |             memcpy(&(<PyDatetimeScalarObject *>obj_val).obval, aux16, 8)
 613 |             (<PyDatetimeScalarObject *> obj_val).obmeta.base = NPY_FR_s
 614 |             (<PyDatetimeScalarObject *> obj_val).obmeta.num = 1
 615 |             return <object> obj_val
 616 |         return PyDateTimeAPI.DateTime_FromDateAndTime(
 617 |             year, month, day, hour, minute, second, 0, <object> utctz, PyDateTimeAPI.DateTimeType,
 618 |         )
 619 |     elif dtype == DT_TD:
 620 |         if fread(&auxi32, 4, 1, stream) != 1:
 621 |             raise ValueError(corrupted_msg % (ftell(stream), "td"))
 622 |         if fread(aux16, 2, 1, stream) != 1:
 623 |             raise ValueError(corrupted_msg % (ftell(stream), "td"))
 624 |         day = auxi32 >> 1
 625 |         second = aux16[0] + ((auxi32 & 1) << 16)
 626 |         return PyDateTimeAPI.Delta_FromDelta(day, second, 0, 1, PyDateTimeAPI.DeltaType)
 627 |     elif dtype == DT_BOOL:
 628 |         if fread(&bool, 1, 1, stream) != 1:
 629 |             raise ValueError(corrupted_msg % (ftell(stream), "bool"))
 630 |         if bool:
 631 |             return True
 632 |         return False
 633 |     elif dtype == DT_MODEL:
 634 |         obj = <object> spec.model
 635 |         if fread(&aux32, 4, 1, stream) != 1:
 636 |             raise ValueError(corrupted_msg % (ftell(stream), "model"))
 637 |         if aux32 != dereference(spec.nested).size():
 638 |             raise ValueError(corrupted_msg % (ftell(stream), f"{obj} has changed"))
 639 |         obj = obj.__new__(obj)
 640 |         for i in range(aux32):
 641 |             field = &dereference(spec.nested)[i]
 642 |             val = _read_model(field, stream, raw, corrupted_msg)
 643 |             Py_INCREF(val)
 644 |             (<PyObject **>((<char *><PyObject *> obj) + field.offset))[0] = <PyObject *> val
 645 |         return obj
 646 |     elif dtype == DT_LIST:
 647 |         if fread(&aux32, 4, 1, stream) != 1:
 648 |             raise ValueError(corrupted_msg % (ftell(stream), "list"))
 649 |         obj = PyList_New(aux32)
 650 |         for i in range(aux32):
 651 |             val = _read_model(&dereference(spec.nested)[0], stream, raw, corrupted_msg)
 652 |             Py_INCREF(val)
 653 |             PyList_SET_ITEM(obj, i, val)
 654 |         return obj
 655 |     elif dtype == DT_DICT:
 656 |         if fread(&aux32, 4, 1, stream) != 1:
 657 |             raise ValueError(corrupted_msg % (ftell(stream), "dict"))
 658 |         obj = PyDict_New()
 659 |         for i in range(aux32):
 660 |             key = _read_model(&dereference(spec.nested)[0], stream, raw, corrupted_msg)
 661 |             val = _read_model(&dereference(spec.nested)[1], stream, raw, corrupted_msg)
 662 |             PyDict_SetItem(obj, key, val)
 663 |         return obj
 664 |     else:
 665 |         raise AssertionError(f"Unsupported dtype: {dtype}")
 666 | 
 667 | 
 668 | def model_to_json(model, alloc_capsule=None) -> bytes:
 669 |     cdef:
 670 |         bytes result
 671 |         optional[mi_heap_destroy_stl_allocator[char]] alloc
 672 |         optional[chunked_stream] stream
 673 |         type root_type
 674 |         SpecNode spec
 675 |         PyObject *error
 676 | 
 677 |     if model is None:
 678 |         return b"null"
 679 |     if PyList_CheckExact(<PyObject *> model):
 680 |         if PyList_GET_SIZE(<PyObject *> model) == 0:
 681 |             return b"[]"
 682 |         spec.type = DT_LIST
 683 |         root_type = type(model[0])
 684 |     else:
 685 |         spec.type = DT_MODEL
 686 |         root_type = type(model)
 687 | 
 688 |     if alloc_capsule is not None:
 689 |         alloc.emplace(dereference(mi_heap_allocator_from_capsule(alloc_capsule)))
 690 |     else:
 691 |         alloc.emplace()
 692 |     stream.emplace(dereference(alloc))
 693 | 
 694 |     spec.nested.emplace(dereference(alloc))
 695 |     _apply_data_type(0, b"root", <PyTypeObject *> root_type, &spec, dereference(alloc))
 696 |     if spec.type == DT_MODEL:
 697 |         spec = dereference(spec.nested)[0]
 698 | 
 699 |     with nogil:
 700 |         error = _write_json(<PyObject *> model, spec, dereference(stream))
 701 |     if error != NULL:
 702 |         raise AssertionError(
 703 |             f"failed to serialize to JSON: {<object> error} of type {type(<object> error).__name__}"
 704 |         )
 705 | 
 706 |     size = dereference(stream).size()
 707 |     result = PyBytes_FromStringAndSize(NULL, size)
 708 |     dereference(stream).dump(PyBytes_AS_STRING(<PyObject *> result), size)
 709 |     return result
 710 | 
 711 | 
 712 | cdef SpecNode fake_str_model
 713 | fake_str_model.type = DT_STRING
 714 | 
 715 | 
 716 | @cython.cdivision(True)
 717 | cdef PyObject *_write_json(PyObject *obj, SpecNode &spec, chunked_stream &stream) nogil:
 718 |     cdef:
 719 |         PyObject *key = NULL
 720 |         PyObject *value = NULL
 721 |         PyObject *r
 722 |         Py_ssize_t pos = 0, size, i, j, item_len, char_len
 723 |         unsigned int kind
 724 |         char sym
 725 |         char *data
 726 |         int aux, auxdiv, rem, year, month, day
 727 |         long val_long, div
 728 |         npy_int64 obval
 729 |         double val_double
 730 |         float val_float
 731 |         char buffer[24]
 732 |         SpecNode *nested
 733 | 
 734 |     if obj == Py_None:
 735 |         stream.write(b"null", 4)
 736 |         return NULL
 737 |     if spec.type == DT_MODEL:
 738 |         if Py_TYPE(obj).tp_members == NULL:
 739 |             # this is just a check for __slots__, it's hard to validate better without GIL
 740 |             return obj
 741 |         stream.write(b"{", 1)
 742 |         kind = 0
 743 |         for i in range(<Py_ssize_t> dereference(spec.nested).size()):
 744 |             nested = &dereference(spec.nested)[i]
 745 |             value = dereference(<PyObject **>((<char *> obj) + nested.offset))
 746 |             if (nested.flags & DF_OPTIONAL) and not (nested.flags & DF_VERBATIM):
 747 |                 if value == NULL or value == Py_None:
 748 |                     continue
 749 |                 if PyList_CheckExact(value) and PyList_GET_SIZE(value) == 0:
 750 |                     continue
 751 |                 if PyDict_CheckExact(value) and PyDict_Size(value) == 0:
 752 |                     continue
 753 |                 if PyArray_CheckExact(value) and PyArray_NDIM(value) == 1 and PyArray_DIM(value, 0) == 0:
 754 |                     continue
 755 |                 if nested.type == DT_FLOAT:
 756 |                     if PyFloat_CheckExact(value):
 757 |                         val_double = PyFloat_AS_DOUBLE(value)
 758 |                     elif PyObject_TypeCheck(value, &PyDoubleArrType_Type):
 759 |                         PyArray_ScalarAsCtype(value, &val_double)
 760 |                     elif PyObject_TypeCheck(value, &PyFloatArrType_Type):
 761 |                         PyArray_ScalarAsCtype(value, &val_float)
 762 |                         val_double = val_float
 763 |                     if val_double != val_double:
 764 |                         continue
 765 |             if kind:
 766 |                 stream.write(b",", 1)
 767 |             else:
 768 |                 kind = 1
 769 |             if nested.flags & DF_KEY_UNMAPPED:
 770 |                 stream.write(b'"', 1)
 771 |                 stream.write(<const char *> nested.key, strlen(<const char *> nested.key))
 772 |                 stream.write(b'"', 1)
 773 |             else:
 774 |                 r = _write_json(<PyObject *> nested.key, fake_str_model, stream)
 775 |                 if r != NULL:
 776 |                     return r
 777 |             stream.write(b":", 1)
 778 |             r = _write_json(value, dereference(nested), stream)
 779 |             if r != NULL:
 780 |                 return r
 781 |         stream.write(b"}", 1)
 782 |     elif spec.type == DT_DICT:
 783 |         stream.write(b"{", 1)
 784 |         if not PyDict_CheckExact(obj):
 785 |             return obj
 786 |         while PyDict_Next(obj, &pos, &key, &value):
 787 |             if pos != 1:
 788 |                 stream.write(b",", 1)
 789 |             r = _write_json(key, dereference(spec.nested)[0], stream)
 790 |             if r != NULL:
 791 |                 return r
 792 |             stream.write(b":", 1)
 793 |             r = _write_json(value, dereference(spec.nested)[1], stream)
 794 |             if r != NULL:
 795 |                 return r
 796 |         stream.write(b"}", 1)
 797 |     elif spec.type == DT_LIST:
 798 |         stream.write(b"[", 1)
 799 |         nested = &dereference(spec.nested)[0]
 800 |         if not PyList_CheckExact(obj):
 801 |             if not PyArray_CheckExact(obj) or not PyArray_IS_C_CONTIGUOUS(obj) or PyArray_NDIM(obj) != 1:
 802 |                 return obj
 803 |             npdata = <PyObject **> PyArray_DATA(obj)
 804 |             for i in range(PyArray_DIM(obj, 0)):
 805 |                 if i != 0:
 806 |                     stream.write(b",", 1)
 807 |                 r = _write_json(npdata[i], dereference(nested), stream)
 808 |                 if r != NULL:
 809 |                     return r
 810 |         else:
 811 |             for i in range(PyList_GET_SIZE(obj)):
 812 |                 if i != 0:
 813 |                     stream.write(b",", 1)
 814 |                 r = _write_json(PyList_GET_ITEM(obj, i), dereference(nested), stream)
 815 |                 if r != NULL:
 816 |                     return r
 817 |         stream.write(b"]", 1)
 818 |     elif spec.type == DT_STRING:
 819 |         stream.write(b'"', 1)
 820 | 
 821 |         if PyUnicode_Check(obj):
 822 |             data = <char *> PyUnicode_DATA(obj)
 823 |             kind = PyUnicode_KIND(obj)
 824 |             item_len = PyUnicode_GET_LENGTH(obj)
 825 |             if kind == PyUnicode_1BYTE_KIND:
 826 |                 for i in range(item_len):
 827 |                     stream.write(buffer, ucs4_to_utf8_json((<uint8_t *> data)[i], buffer))
 828 |             elif kind == PyUnicode_2BYTE_KIND:
 829 |                 for i in range(item_len):
 830 |                     stream.write(buffer, ucs4_to_utf8_json((<uint16_t *> data)[i], buffer))
 831 |             elif kind == PyUnicode_4BYTE_KIND:
 832 |                 for i in range(item_len):
 833 |                     stream.write(buffer, ucs4_to_utf8_json((<uint32_t *> data)[i], buffer))
 834 |         elif PyBytes_Check(obj):
 835 |             data = PyBytes_AS_STRING(obj)
 836 |             item_len = PyBytes_GET_SIZE(obj)
 837 |             for i in range(item_len):
 838 |                 stream.write(buffer, ucs4_to_utf8_json((<uint8_t *> data)[i], buffer))
 839 |         else:
 840 |             return obj
 841 | 
 842 |         stream.write(b'"', 1)
 843 |     elif spec.type == DT_DT:
 844 |         buffer[0] = buffer[21] = b'"'
 845 |         if not PyDateTime_Check(obj):
 846 |             if PyObject_TypeCheck(obj, &PyDatetimeArrType_Type):
 847 |                 npy_unit = (<PyDatetimeScalarObject *> obj).obmeta.base
 848 |                 obval = (<PyDatetimeScalarObject *> obj).obval
 849 |                 if npy_unit == NPY_FR_ns:
 850 |                     obval //= 1000000000
 851 |                 elif npy_unit == NPY_FR_us:
 852 |                     obval //= 1000000
 853 |                 elif npy_unit != NPY_FR_s:
 854 |                     return obj
 855 |                 val_long = obval // (60 * 60 * 24)
 856 |                 obval = obval - val_long * 60 * 60 * 24
 857 | 
 858 |                 year = month = day = 0
 859 |                 set_datetimestruct_days(val_long, &year, &month, &day)
 860 |                 aux = year
 861 |                 pos = 4
 862 |                 while pos > 0:
 863 |                     auxdiv = aux
 864 |                     aux = aux // 10
 865 |                     buffer[pos] = auxdiv - aux * 10 + ord(b"0")
 866 |                     pos -= 1
 867 |                 buffer[5] = b"-"
 868 | 
 869 |                 aux = month
 870 |                 if aux < 10:
 871 |                     buffer[6] = b"0"
 872 |                     buffer[7] = ord(b"0") + aux
 873 |                 else:
 874 |                     buffer[6] = b"1"
 875 |                     buffer[7] = ord(b"0") + aux - 10
 876 |                 buffer[8] = b"-"
 877 | 
 878 |                 aux = day
 879 |                 auxdiv = aux // 10
 880 |                 buffer[9] = ord(b"0") + auxdiv
 881 |                 buffer[10] = ord(b"0") + aux - auxdiv * 10
 882 |                 buffer[11] = b"T"
 883 | 
 884 |                 auxdiv = obval // 60
 885 |                 aux = obval - auxdiv * 60
 886 |                 rem = auxdiv
 887 |                 auxdiv = aux // 10
 888 |                 buffer[18] = ord(b"0") + auxdiv
 889 |                 buffer[19] = ord(b"0") + aux - auxdiv * 10
 890 |                 buffer[20] = b"Z"
 891 | 
 892 |                 auxdiv = rem // 60
 893 |                 aux = rem - auxdiv * 60
 894 |                 rem = auxdiv
 895 |                 auxdiv = aux // 10
 896 |                 buffer[15] = ord(b"0") + auxdiv
 897 |                 buffer[16] = ord(b"0") + aux - auxdiv * 10
 898 |                 buffer[17] = b":"
 899 | 
 900 |                 aux = rem
 901 |                 auxdiv = aux // 10
 902 |                 buffer[12] = ord(b"0") + auxdiv
 903 |                 buffer[13] = ord(b"0") + aux - auxdiv * 10
 904 |                 buffer[14] = b":"
 905 |             else:
 906 |                 return obj
 907 |         else:
 908 |             if (<PyDateTime_CAPI *> PyDateTimeAPI).TimeZone_UTC != PyDateTime_DATE_GET_TZINFO(obj):
 909 |                 return obj
 910 |             aux = PyDateTime_GET_YEAR(obj)
 911 |             pos = 4
 912 |             while pos > 0:
 913 |                 auxdiv = aux
 914 |                 aux = aux // 10
 915 |                 buffer[pos] = auxdiv - aux * 10 + ord(b"0")
 916 |                 pos -= 1
 917 |             buffer[5] = b"-"
 918 |             aux = PyDateTime_GET_MONTH(obj)
 919 |             if aux < 10:
 920 |                 buffer[6] = b"0"
 921 |                 buffer[7] = ord(b"0") + aux
 922 |             else:
 923 |                 buffer[6] = b"1"
 924 |                 buffer[7] = ord(b"0") + aux - 10
 925 |             buffer[8] = b"-"
 926 |             aux = PyDateTime_GET_DAY(obj)
 927 |             auxdiv = aux // 10
 928 |             buffer[9] = ord(b"0") + auxdiv
 929 |             buffer[10] = ord(b"0") + aux - auxdiv * 10
 930 |             buffer[11] = b"T"
 931 |             aux = PyDateTime_DATE_GET_HOUR(obj)
 932 |             auxdiv = aux // 10
 933 |             buffer[12] = ord(b"0") + auxdiv
 934 |             buffer[13] = ord(b"0") + aux - auxdiv * 10
 935 |             buffer[14] = b":"
 936 |             aux = PyDateTime_DATE_GET_MINUTE(obj)
 937 |             auxdiv = aux // 10
 938 |             buffer[15] = ord(b"0") + auxdiv
 939 |             buffer[16] = ord(b"0") + aux - auxdiv * 10
 940 |             buffer[17] = b":"
 941 |             aux = PyDateTime_DATE_GET_SECOND(obj)
 942 |             auxdiv = aux // 10
 943 |             buffer[18] = ord(b"0") + auxdiv
 944 |             buffer[19] = ord(b"0") + aux - auxdiv * 10
 945 |             buffer[20] = b"Z"
 946 |         stream.write(buffer, 22)
 947 |     elif spec.type == DT_TD:
 948 |         stream.write(b'"', 1)
 949 |         if not PyDelta_Check(obj):
 950 |             if PyObject_TypeCheck(obj, &PyTimedeltaArrType_Type):
 951 |                 npy_unit = (<PyDatetimeScalarObject *> obj).obmeta.base
 952 |                 val_long = (<PyDatetimeScalarObject *> obj).obval
 953 |                 if npy_unit == NPY_FR_ns:
 954 |                     val_long //= 1000000000
 955 |                 elif npy_unit == NPY_FR_us:
 956 |                     val_long //= 1000000
 957 |                 elif npy_unit != NPY_FR_s:
 958 |                     return obj
 959 |             else:
 960 |                 return obj
 961 |         else:
 962 |             val_long = PyDateTime_DELTA_GET_DAYS(obj)
 963 |             val_long *= 24 * 3600
 964 |             val_long += PyDateTime_DELTA_GET_SECONDS(obj)
 965 |         if val_long < 0:
 966 |             stream.write(b"-", 1)
 967 |             val_long = -val_long
 968 |         if val_long == 0:
 969 |             stream.write(b"0", 1)
 970 |         else:
 971 |             pos = 0
 972 |             while val_long:
 973 |                 div = val_long
 974 |                 val_long = val_long // 10
 975 |                 buffer[pos] = div - val_long * 10 + ord(b"0")
 976 |                 pos += 1
 977 |             for i in range(pos // 2):
 978 |                 sym = buffer[i]
 979 |                 div = pos - i - 1
 980 |                 buffer[i] = buffer[div]
 981 |                 buffer[div] = sym
 982 |             stream.write(buffer, pos)
 983 |         stream.write(b's"', 2)
 984 |     elif spec.type == DT_LONG:
 985 |         if PyLong_CheckExact(obj):
 986 |             val_long = PyLong_AsLong(obj)
 987 |         elif PyObject_TypeCheck(obj, &PyIntegerArrType_Type):
 988 |             val_long = 0
 989 |             PyArray_ScalarAsCtype(obj, &val_long)
 990 |         else:
 991 |             return obj
 992 |         if val_long < 0:
 993 |             stream.write(b"-", 1)
 994 |             val_long = -val_long
 995 |         if val_long == 0:
 996 |             stream.write(b"0", 1)
 997 |         else:
 998 |             pos = 0
 999 |             while val_long:
1000 |                 div = val_long
1001 |                 val_long = val_long // 10
1002 |                 buffer[pos] = div - val_long * 10 + ord(b"0")
1003 |                 pos += 1
1004 |             for i in range(pos // 2):
1005 |                 sym = buffer[i]
1006 |                 div = pos - i - 1
1007 |                 buffer[i] = buffer[div]
1008 |                 buffer[div] = sym
1009 |             stream.write(buffer, pos)
1010 |     elif spec.type == DT_BOOL:
1011 |         if obj == Py_True:
1012 |             stream.write(b"true", 4)
1013 |         elif obj == Py_False:
1014 |             stream.write(b"false", 5)
1015 |         else:
1016 |             return obj
1017 |     elif spec.type == DT_FLOAT:
1018 |         if PyFloat_CheckExact(obj):
1019 |             val_double = PyFloat_AS_DOUBLE(obj)
1020 |         elif PyLong_CheckExact(obj):
1021 |             val_double = PyLong_AsLong(obj)
1022 |         elif PyObject_TypeCheck(obj, &PyDoubleArrType_Type):
1023 |             PyArray_ScalarAsCtype(obj, &val_double)
1024 |         elif PyObject_TypeCheck(obj, &PyFloatArrType_Type):
1025 |             PyArray_ScalarAsCtype(obj, &val_float)
1026 |             val_double = val_float
1027 |         elif PyObject_TypeCheck(obj, &PyIntegerArrType_Type):
1028 |             val_long = 0
1029 |             PyArray_ScalarAsCtype(obj, &val_long)
1030 |             val_double = val_long
1031 |         else:
1032 |             return obj
1033 |         gcvt(val_double, 24, buffer)
1034 |         stream.write(buffer, strlen(buffer))
1035 |     elif spec.type == DT_FREEFORM:
1036 |         r = _write_freeform_json(obj, stream)
1037 |         if r != NULL:
1038 |             return r
1039 |     else:
1040 |         return obj
1041 |     return NULL
1042 | 
1043 | 
1044 | @cython.cdivision(True)
1045 | cdef PyObject *_write_freeform_json(PyObject *node, chunked_stream &stream) nogil:
1046 |     cdef:
1047 |         PyObject *key = NULL
1048 |         PyObject *value = NULL
1049 |         PyObject *r
1050 |         Py_ssize_t pos = 0, size = 0, i, j, item_len, char_len
1051 |         char *data
1052 |         unsigned int kind
1053 |         char sym
1054 |         long val_long, div
1055 |         double float_val
1056 |         char buffer[24]
1057 | 
1058 |     if PyDict_CheckExact(node):
1059 |         stream.write(b"{", 1)
1060 |         while PyDict_Next(node, &pos, &key, &value):
1061 |             if pos != 1:
1062 |                 stream.write(b",", 1)
1063 |             r = _write_freeform_json(key, stream)
1064 |             if r != NULL:
1065 |                 return r
1066 |             stream.write(b":", 1)
1067 |             r = _write_freeform_json(value, stream)
1068 |             if r != NULL:
1069 |                 return r
1070 |         stream.write(b"}", 1)
1071 |     elif PyList_CheckExact(node):
1072 |         stream.write(b"[", 1)
1073 |         for i in range(PyList_GET_SIZE(node)):
1074 |             if i != 0:
1075 |                 stream.write(b",", 1)
1076 |             r = _write_freeform_json(PyList_GET_ITEM(node, i), stream)
1077 |             if r != NULL:
1078 |                 return r
1079 |         stream.write(b"]", 1)
1080 |     elif PyUnicode_Check(node):
1081 |         stream.write(b'"', 1)
1082 | 
1083 |         data = <char *>PyUnicode_DATA(node)
1084 |         kind = PyUnicode_KIND(node)
1085 |         item_len = PyUnicode_GET_LENGTH(node)
1086 |         if kind == PyUnicode_1BYTE_KIND:
1087 |             for i in range(item_len):
1088 |                 stream.write(buffer, ucs4_to_utf8_json((<uint8_t *> data)[i], buffer))
1089 |         elif kind == PyUnicode_2BYTE_KIND:
1090 |             for i in range(item_len):
1091 |                 stream.write(buffer, ucs4_to_utf8_json((<uint16_t *> data)[i], buffer))
1092 |         elif kind == PyUnicode_4BYTE_KIND:
1093 |             for i in range(item_len):
1094 |                 stream.write(buffer, ucs4_to_utf8_json((<uint32_t *> data)[i], buffer))
1095 | 
1096 |         stream.write(b'"', 1)
1097 |     elif PyLong_CheckExact(node):
1098 |         val_long = PyLong_AsLong(node)
1099 |         if val_long < 0:
1100 |             stream.write(b"-", 1)
1101 |             val_long = -val_long
1102 |         if val_long == 0:
1103 |             stream.write(b"0", 1)
1104 |         else:
1105 |             pos = 0
1106 |             while val_long:
1107 |                 div = val_long
1108 |                 val_long = val_long // 10
1109 |                 buffer[pos] = div - val_long * 10 + ord(b"0")
1110 |                 pos += 1
1111 |             for i in range(pos // 2):
1112 |                 sym = buffer[i]
1113 |                 div = pos - i - 1
1114 |                 buffer[i] = buffer[div]
1115 |                 buffer[div] = sym
1116 |             stream.write(buffer, pos)
1117 |     elif node == Py_True:
1118 |         stream.write(b"true", 4)
1119 |     elif node == Py_False:
1120 |         stream.write(b"false", 5)
1121 |     elif PyFloat_CheckExact(node):
1122 |         gcvt(PyFloat_AS_DOUBLE(node), 24, buffer)
1123 |         stream.write(buffer, strlen(buffer))
1124 |     elif node == Py_None:
1125 |         stream.write(b"null", 4)
1126 |     else:
1127 |         return node
1128 |     return NULL
1129 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2022 Athenian
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Athenian API - open
2 | 
3 | Open source bits of [athenianco/athenian-api](https://github.com/athenianco/athenian-api). 
4 | 
5 | License: MIT.
6 | 


--------------------------------------------------------------------------------
/async_utils.py:
--------------------------------------------------------------------------------
  1 | import asyncio
  2 | from datetime import datetime, timezone
  3 | import logging
  4 | import textwrap
  5 | from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Type, Union
  6 | 
  7 | import numpy as np
  8 | import pandas as pd
  9 | from pandas.core.dtypes.cast import OutOfBoundsDatetime, tslib
 10 | from pandas.core.internals import make_block
 11 | from pandas.core.internals.managers import BlockManager, form_blocks
 12 | import sentry_sdk
 13 | from sqlalchemy import Boolean, Column, DateTime, Integer
 14 | from sqlalchemy.orm.attributes import InstrumentedAttribute
 15 | from sqlalchemy.sql import ClauseElement
 16 | from sqlalchemy.sql.elements import Label
 17 | 
 18 | from athenian.api import metadata
 19 | from athenian.api.db import Database, DatabaseLike
 20 | from athenian.api.models.metadata.github import Base as MetadataBase
 21 | from athenian.api.models.persistentdata.models import Base as PerdataBase
 22 | from athenian.api.models.precomputed.models import GitHubBase as PrecomputedBase
 23 | from athenian.api.models.state.models import Base as StateBase
 24 | from athenian.api.to_object_arrays import to_object_arrays_split
 25 | from athenian.api.tracing import MAX_SENTRY_STRING_LENGTH
 26 | 
 27 | 
 28 | async def read_sql_query(sql: ClauseElement,
 29 |                          con: DatabaseLike,
 30 |                          columns: Union[Sequence[str], Sequence[InstrumentedAttribute],
 31 |                                         MetadataBase, PerdataBase, PrecomputedBase, StateBase],
 32 |                          index: Optional[Union[str, Sequence[str]]] = None,
 33 |                          soft_limit: Optional[int] = None,
 34 |                          ) -> pd.DataFrame:
 35 |     """Read SQL query into a DataFrame.
 36 | 
 37 |     Returns a DataFrame corresponding to the result set of the query.
 38 |     Optionally provide an `index_col` parameter to use one of the
 39 |     columns as the index, otherwise default integer index will be used.
 40 | 
 41 |     Parameters
 42 |     ----------
 43 |     sql     : SQLAlchemy query object to be executed.
 44 |     con     : async SQLAlchemy database engine.
 45 |     columns : list of the resulting columns names, column objects or the model if SELECT *
 46 |     index   : Name(s) of the index column(s).
 47 |     soft_limit
 48 |             : Load this number of rows at maximum.
 49 | 
 50 |     Returns
 51 |     -------
 52 |     DataFrame
 53 | 
 54 |     Notes
 55 |     -----
 56 |     Any datetime values with time zone information parsed via the `parse_dates`
 57 |     parameter will be converted to UTC.
 58 |     """
 59 |     try:
 60 |         data = await con.fetch_all(query=sql)
 61 |     except Exception as e:
 62 |         try:
 63 |             sql = str(sql)
 64 |         except Exception:
 65 |             sql = repr(sql)
 66 |         sql = textwrap.shorten(sql, MAX_SENTRY_STRING_LENGTH - 500)
 67 |         logging.getLogger("%s.read_sql_query" % metadata.__package__).error(
 68 |             "%s: %s; %s", type(e).__name__, e, sql)
 69 |         raise e from None
 70 |     if soft_limit is not None and len(data) > soft_limit:
 71 |         data = data[:soft_limit]
 72 |     return wrap_sql_query(data, columns, index)
 73 | 
 74 | 
 75 | def _create_block_manager_from_arrays(
 76 |     arrays_typed: List[np.ndarray],
 77 |     arrays_obj: np.ndarray,
 78 |     names_typed: List[str],
 79 |     names_obj: List[str],
 80 |     size: int,
 81 | ) -> BlockManager:
 82 |     assert len(arrays_typed) == len(names_typed)
 83 |     assert len(arrays_obj) == len(names_obj)
 84 |     range_index = pd.RangeIndex(stop=size)
 85 |     typed_index = pd.Index(names_typed)
 86 |     blocks = form_blocks(arrays_typed, typed_index, [typed_index, range_index])
 87 |     blocks.append(make_block(arrays_obj, placement=np.arange(len(arrays_obj)) + len(arrays_typed)))
 88 |     return BlockManager(blocks, [pd.Index(names_typed + names_obj), range_index])
 89 | 
 90 | 
 91 | def wrap_sql_query(data: List[Sequence[Any]],
 92 |                    columns: Union[Sequence[str], Sequence[InstrumentedAttribute],
 93 |                                   MetadataBase, StateBase],
 94 |                    index: Optional[Union[str, Sequence[str]]] = None,
 95 |                    ) -> pd.DataFrame:
 96 |     """Turn the fetched DB records to a pandas DataFrame."""
 97 |     try:
 98 |         columns[0]
 99 |     except TypeError:
100 |         dt_columns = _extract_datetime_columns(columns.__table__.columns)
101 |         int_columns = _extract_integer_columns(columns.__table__.columns)
102 |         bool_columns = _extract_boolean_columns(columns.__table__.columns)
103 |         columns = [c.name for c in columns.__table__.columns]
104 |     else:
105 |         dt_columns = _extract_datetime_columns(columns)
106 |         int_columns = _extract_integer_columns(columns)
107 |         bool_columns = _extract_boolean_columns(columns)
108 |         columns = [(c.name if not isinstance(c, str) else c) for c in columns]
109 |     typed_cols_indexes = []
110 |     typed_cols_names = []
111 |     obj_cols_indexes = []
112 |     obj_cols_names = []
113 |     for i, column in enumerate(columns):
114 |         if column in dt_columns or column in int_columns or column in bool_columns:
115 |             cols_indexes = typed_cols_indexes
116 |             cols_names = typed_cols_names
117 |         else:
118 |             cols_indexes = obj_cols_indexes
119 |             cols_names = obj_cols_names
120 |         cols_indexes.append(i)
121 |         cols_names.append(column)
122 |     log = logging.getLogger(f"{metadata.__package__}.wrap_sql_query")
123 |     # we used to have pd.DataFrame.from_records + bunch of convert_*() in relevant columns
124 |     # the current approach is faster for several reasons:
125 |     # 1. avoid an expensive copy of the object dtype columns in the BlockManager construction
126 |     # 2. call tslib.array_to_datetime directly without surrounding Pandas bloat
127 |     # 3. convert to int in the numpy domain and thus do not have to mess with indexes
128 |     #
129 |     # an ideal conversion would be loading columns directly from asyncpg but that requires
130 |     # quite some changes in their internals
131 |     with sentry_sdk.start_span(op="wrap_sql_query/convert", description=str(size := len(data))):
132 |         data_typed, data_obj = to_object_arrays_split(data, typed_cols_indexes, obj_cols_indexes)
133 |         converted_typed = []
134 |         discard_mask = None
135 |         for column, values in zip(typed_cols_names, data_typed):
136 |             if column in dt_columns:
137 |                 converted_typed.append(_convert_datetime(values))
138 |             elif column in int_columns:
139 |                 values, discarded = _convert_integer(values, column, int_columns[column], log)
140 |                 converted_typed.append(values)
141 |                 if discarded is not None:
142 |                     if discard_mask is None:
143 |                         discard_mask = np.zeros(len(data), dtype=bool)
144 |                     discard_mask[discarded] = True
145 |             elif column in bool_columns:
146 |                 converted_typed.append(values.astype(bool))
147 |             else:
148 |                 raise AssertionError("impossible: typed columns are either dt or int")
149 |         if discard_mask is not None:
150 |             left = ~discard_mask
151 |             size = left.sum()
152 |             converted_typed = [arr[left] for arr in converted_typed]
153 |             data_obj = data_obj[:, left]
154 |     with sentry_sdk.start_span(op="wrap_sql_query/pd.DataFrame()", description=str(size)):
155 |         block_mgr = _create_block_manager_from_arrays(
156 |             converted_typed, data_obj, typed_cols_names, obj_cols_names, size)
157 |         frame = pd.DataFrame(block_mgr, columns=typed_cols_names + obj_cols_names, copy=False)
158 |         for column in dt_columns:
159 |             try:
160 |                 frame[column] = frame[column].dt.tz_localize(timezone.utc)
161 |             except (AttributeError, TypeError):
162 |                 continue
163 |         if index is not None:
164 |             frame.set_index(index, inplace=True)
165 |     return frame
166 | 
167 | 
168 | def _extract_datetime_columns(columns: Iterable[Union[Column, str]]) -> Set[str]:
169 |     return {
170 |         c.name for c in columns
171 |         if not isinstance(c, str) and (
172 |             isinstance(c.type, DateTime) or
173 |             (isinstance(c.type, type) and issubclass(c.type, DateTime))
174 |         )
175 |     }
176 | 
177 | 
178 | def _extract_boolean_columns(columns: Iterable[Union[Column, str]]) -> Set[str]:
179 |     return {
180 |         c.name for c in columns
181 |         if not isinstance(c, str) and (
182 |             isinstance(c.type, Boolean) or
183 |             (isinstance(c.type, type) and issubclass(c.type, Boolean))
184 |         )
185 |     }
186 | 
187 | 
188 | def _extract_integer_columns(columns: Iterable[Union[Column, str]],
189 |                              ) -> Dict[str, bool]:
190 |     return {
191 |         c.name: getattr(
192 |             c, "info", {} if not isinstance(c, Label) else getattr(c.element, "info", {}),
193 |         ).get("erase_nulls", False)
194 |         for c in columns
195 |         if not isinstance(c, str) and (
196 |             isinstance(c.type, Integer) or
197 |             (isinstance(c.type, type) and issubclass(c.type, Integer))
198 |         )
199 |         and not getattr(c, "nullable", False)
200 |         and (not isinstance(c, Label) or (
201 |             (not getattr(c.element, "nullable", False))
202 |             and (not getattr(c, "nullable", False))
203 |         ))
204 |     }
205 | 
206 | 
207 | def _convert_datetime(arr: np.ndarray) -> np.ndarray:
208 |     # None converts to NaT
209 |     try:
210 |         ts, offset = tslib.array_to_datetime(arr, utc=True, errors="raise")
211 |         assert offset is None
212 |     except OutOfBoundsDatetime:
213 |         # TODO(vmarkovtsev): copy the function and set OOB values to NaT
214 |         # this comparison is very slow but still faster than removing tzinfo and taking np.array()
215 |         arr[arr == datetime(1, 1, 1)] = None
216 |         arr[arr == datetime(1, 1, 1, tzinfo=timezone.utc)] = None
217 |         try:
218 |             return _convert_datetime(arr)
219 |         except OutOfBoundsDatetime as e:
220 |             raise e from None
221 |     # 0 converts to 1970-01-01T00:00:00
222 |     ts[ts == np.zeros(1, ts.dtype)[0]] = None
223 |     return ts
224 | 
225 | 
226 | def postprocess_datetime(frame: pd.DataFrame,
227 |                          columns: Optional[Iterable[str]] = None,
228 |                          ) -> pd.DataFrame:
229 |     """Ensure *inplace* that all the timestamps inside the dataframe are valid UTC or NaT.
230 | 
231 |     :return: Fixed dataframe - the same instance as `frame`.
232 |     """
233 |     utc_dt1 = datetime(1, 1, 1, tzinfo=timezone.utc)
234 |     dt1 = datetime(1, 1, 1)
235 |     if columns is not None:
236 |         obj_cols = dt_cols = columns
237 |     else:
238 |         obj_cols = frame.select_dtypes(include=[object])
239 |         dt_cols = frame.select_dtypes(include=["datetime"])
240 |     for col in obj_cols:
241 |         fc = frame[col]
242 |         if utc_dt1 in fc:
243 |             fc.replace(utc_dt1, pd.NaT, inplace=True)
244 |         if dt1 in fc:
245 |             fc.replace(dt1, pd.NaT, inplace=True)
246 |     for col in dt_cols:
247 |         fc = frame[col]
248 |         if 0 in fc:
249 |             fc.replace(0, pd.NaT, inplace=True)
250 |         try:
251 |             frame[col] = fc.dt.tz_localize(timezone.utc)
252 |         except (AttributeError, TypeError):
253 |             continue
254 |     return frame
255 | 
256 | 
257 | def _convert_integer(arr: np.ndarray,
258 |                      name: str,
259 |                      erase_null: bool,
260 |                      log: logging.Logger,
261 |                      ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
262 |     nulls = None
263 |     while True:
264 |         try:
265 |             return arr.astype(int), nulls
266 |         except TypeError as e:
267 |             nulls = np.equal(arr, None)
268 |             if not nulls.any() or not erase_null:
269 |                 raise ValueError(f"Column {name} is not all-integer") from e
270 |             log.error("fetched nulls instead of integers in %s", name)
271 |             arr[nulls] = 0
272 | 
273 | 
274 | def postprocess_integer(frame: pd.DataFrame, columns: Iterable[Tuple[str, int]]) -> pd.DataFrame:
275 |     """Ensure *inplace* that all the integers inside the dataframe are not objects.
276 | 
277 |     :return: Fixed dataframe, a potentially different instance.
278 |     """
279 |     dirty_index = False
280 |     log = None
281 |     for col, erase_null in columns:
282 |         while True:
283 |             try:
284 |                 frame[col] = frame[col].astype(int, copy=False)
285 |                 break
286 |             except TypeError as e:
287 |                 nulls = frame[col].isnull().values
288 |                 if not nulls.any():
289 |                     raise ValueError(f"Column {col} is not all-integer") from e
290 |                 if not erase_null:
291 |                     raise ValueError(f"Column {col} is not all-integer\n"
292 |                                      f"{frame.loc[nulls].to_dict('records')}") from e
293 |                 if log is None:
294 |                     log = logging.getLogger(f"{metadata.__package__}.read_sql_query")
295 |                 log.error("fetched nulls instead of integers in %s: %s",
296 |                           col, frame.loc[nulls].to_dict("records"))
297 |                 frame = frame.take(np.flatnonzero(~nulls))
298 |                 dirty_index = True
299 |     if dirty_index:
300 |         frame.reset_index(drop=True, inplace=True)
301 |     return frame
302 | 
303 | 
304 | async def gather(*coros_or_futures,
305 |                  op: Optional[str] = None,
306 |                  description: Optional[str] = None,
307 |                  catch: Type[BaseException] = Exception,
308 |                  ) -> Tuple[Any, ...]:
309 |     """Return a future aggregating results/exceptions from the given coroutines/futures.
310 | 
311 |     This is equivalent to `asyncio.gather(*coros_or_futures, return_exceptions=True)` with
312 |     subsequent exception forwarding.
313 | 
314 |     :param op: Wrap the execution in a Sentry span with this `op`.
315 |     :param description: Sentry span description.
316 |     :param catch: Forward exceptions of this type.
317 |     """
318 |     async def body():
319 |         if len(coros_or_futures) == 0:
320 |             return tuple()
321 |         if len(coros_or_futures) == 1:
322 |             return (await coros_or_futures[0],)
323 |         results = await asyncio.gather(*coros_or_futures, return_exceptions=True)
324 |         for r in results:
325 |             if isinstance(r, catch):
326 |                 raise r from None
327 |         return results
328 | 
329 |     if op is not None:
330 |         with sentry_sdk.start_span(op=op, description=description):
331 |             return await body()
332 |     return await body()
333 | 
334 | 
335 | async def read_sql_query_with_join_collapse(
336 |         query: ClauseElement,
337 |         db: Database,
338 |         columns: Union[Sequence[str], Sequence[InstrumentedAttribute],
339 |                        MetadataBase, PerdataBase, PrecomputedBase, StateBase],
340 |         index: Optional[Union[str, Sequence[str]]] = None,
341 |         soft_limit: Optional[int] = None,
342 | ) -> pd.DataFrame:
343 |     """Enforce the predefined JOIN order in read_sql_query()."""
344 |     query = query.with_statement_hint("Set(join_collapse_limit 1)")
345 |     return await read_sql_query(query, db, columns=columns, index=index, soft_limit=soft_limit)
346 | 
347 | 
348 | # Allow other coroutines to execute every Nth iteration in long loops
349 | COROUTINE_YIELD_EVERY_ITER = 250
350 | 
351 | 
352 | async def list_with_yield(iterable: Iterable[Any], sentry_op: str) -> List[Any]:
353 |     """Drain an iterable to a list, tracing the loop in Sentry and respecting other coroutines."""
354 |     with sentry_sdk.start_span(op=sentry_op) as span:
355 |         things = []
356 |         for i, thing in enumerate(iterable):
357 |             if (i + 1) % COROUTINE_YIELD_EVERY_ITER == 0:
358 |                 await asyncio.sleep(0)
359 |             things.append(thing)
360 |         try:
361 |             span.description = str(i)
362 |         except UnboundLocalError:
363 |             pass
364 |     return things
365 | 


--------------------------------------------------------------------------------
/asyncpg_recordobj.h:
--------------------------------------------------------------------------------
 1 | // no need to #include anything, this file is used internally by to_object_arrays.pyx
 2 | 
 3 | typedef struct {
 4 |     PyObject_VAR_HEAD
 5 | 
 6 |     // asyncpg specifics begin here
 7 |     // if they add another field, we will break spectacularly
 8 |     Py_hash_t self_hash;
 9 |     PyObject *desc;  // we don't care of the actual type
10 |     PyObject *ob_item[1];  // embedded in the tail, the count matches len()
11 | } ApgRecordObject;
12 | 
13 | #define ApgRecord_GET_ITEM(op, i) (((ApgRecordObject *)(op))->ob_item[i])
14 | #define ApgRecord_SET_ITEM(op, i, v) (((ApgRecordObject *)(op))->ob_item[i] = v)
15 | #define ApgRecord_GET_DESC(op) (((ApgRecordObject *)(op))->desc)
16 | 


--------------------------------------------------------------------------------
/to_object_arrays.pyx:
--------------------------------------------------------------------------------
 1 | # cython: language_level=3, boundscheck=False, nonecheck=False, optimize.unpack_method_calls=True
 2 | # cython: warn.maybe_uninitialized=True
 3 | # distutils: language = c++
 4 | 
 5 | from typing import Any, List, Sequence, Tuple
 6 | 
 7 | import asyncpg
 8 | import numpy as np
 9 | 
10 | cimport cython
11 | from cpython cimport PyObject
12 | from numpy cimport ndarray
13 | 
14 | 
15 | cdef extern from "asyncpg_recordobj.h":
16 |     PyObject *ApgRecord_GET_ITEM(PyObject *, int)
17 | 
18 | 
19 | cdef extern from "Python.h":
20 |     # added nogil -> from cpython cimport ...
21 |     # these are the macros that read directly from the internal ob_items
22 |     PyObject *PyList_GET_ITEM(PyObject *, Py_ssize_t) nogil
23 |     PyObject *PyTuple_GET_ITEM(PyObject *, Py_ssize_t) nogil
24 | 
25 | 
26 | @cython.boundscheck(False)
27 | def to_object_arrays_split(rows: List[Sequence[Any]],
28 |                            typed_indexes: Sequence[int],
29 |                            obj_indexes: Sequence[int],
30 |                            ) -> Tuple[np.ndarray, np.ndarray]:
31 |     """
32 |     Convert a list of tuples into an object array. Any subclass of
33 |     tuple in `rows` will be casted to tuple.
34 | 
35 |     Parameters
36 |     ----------
37 |     rows : 2-d array (N, K)
38 |         List of tuples to be converted into an array. Each tuple must be of equal length,
39 |         otherwise, the results are undefined.
40 |     typed_indexes : array of integers
41 |         Sequence of integer indexes in each tuple in `rows` that select the first result.
42 |     obj_indexes : array of integers
43 |         Sequence of integer indexes in each tuple in `rows` that select the second result.
44 | 
45 |     Returns
46 |     -------
47 |     (np.ndarray[object, ndim=2], np.ndarray[object, ndim=2])
48 |     The first array is the concatenation of columns in `rows` chosen by `typed_indexes`.
49 |     The second array is the concatenation of columns in `rows` chosen by `object_indexes`.
50 |     """
51 |     cdef:
52 |         Py_ssize_t i, j, size, cols_typed, cols_obj
53 |         ndarray[object, ndim=2] result_typed
54 |         ndarray[object, ndim=2] result_obj
55 |         PyObject *record
56 |         long[:] typed_indexes_arr
57 |         long[:] obj_indexes_arr
58 | 
59 |     assert isinstance(rows, list)
60 |     typed_indexes_arr = np.asarray(typed_indexes, dtype=int)
61 |     obj_indexes_arr = np.asarray(obj_indexes, dtype=int)
62 |     size = len(rows)
63 |     cols_typed = len(typed_indexes_arr)
64 |     cols_obj = len(obj_indexes_arr)
65 | 
66 |     result_typed = np.empty((cols_typed, size), dtype=object)
67 |     result_obj = np.empty((cols_obj, size), dtype=object)
68 |     if size == 0:
69 |         return result_typed, result_obj
70 | 
71 |     if isinstance(rows[0], asyncpg.Record):
72 |         for i in range(size):
73 |             record = PyList_GET_ITEM(<PyObject *>rows, i)
74 |             for j in range(cols_typed):
75 |                 result_typed[j, i] = <object>ApgRecord_GET_ITEM(record, typed_indexes_arr[j])
76 |             for j in range(cols_obj):
77 |                 result_obj[j, i] = <object>ApgRecord_GET_ITEM(record, obj_indexes_arr[j])
78 |     elif isinstance(rows[0], tuple):
79 |         for i in range(size):
80 |             record = PyList_GET_ITEM(<PyObject *>rows, i)
81 |             for j in range(cols_typed):
82 |                 result_typed[j, i] = <object>PyTuple_GET_ITEM(record, typed_indexes_arr[j])
83 |             for j in range(cols_obj):
84 |                 result_obj[j, i] = <object>PyTuple_GET_ITEM(record, obj_indexes_arr[j])
85 |     else:
86 |         # convert to tuple
87 |         for i in range(size):
88 |             row = tuple(rows[i])
89 |             for j in range(cols_typed):
90 |                 result_typed[j, i] = row[typed_indexes_arr[j]]
91 |             for j in range(cols_obj):
92 |                 result_obj[j, i] = row[obj_indexes_arr[j]]
93 | 
94 |     return result_typed, result_obj


--------------------------------------------------------------------------------
/typing_utils.py:
--------------------------------------------------------------------------------
  1 | from contextlib import contextmanager
  2 | from contextvars import ContextVar
  3 | import dataclasses
  4 | from datetime import datetime, timedelta
  5 | from itertools import chain
  6 | from typing import Any, Callable, Dict, Iterable, Iterator, List, Mapping, NamedTuple, \
  7 |     Optional, Tuple, Type, TypeVar, Union
  8 | 
  9 | import numpy as np
 10 | import pandas as pd
 11 | from pandas._libs import tslib
 12 | import sentry_sdk
 13 | import xxhash
 14 | 
 15 | from athenian.api.tracing import sentry_span
 16 | 
 17 | 
 18 | def is_generic(klass: type):
 19 |     """Determine whether klass is a generic class."""
 20 |     return hasattr(klass, "__origin__")
 21 | 
 22 | 
 23 | def is_dict(klass: type):
 24 |     """Determine whether klass is a Dict."""
 25 |     return getattr(klass, "__origin__", None) == dict
 26 | 
 27 | 
 28 | def is_list(klass: type):
 29 |     """Determine whether klass is a List."""
 30 |     return getattr(klass, "__origin__", None) == list
 31 | 
 32 | 
 33 | def is_union(klass: type):
 34 |     """Determine whether klass is a Union."""
 35 |     return getattr(klass, "__origin__", None) == Union
 36 | 
 37 | 
 38 | def is_optional(klass: type):
 39 |     """Determine whether klass is an Optional."""
 40 |     return is_union(klass) and \
 41 |         len(klass.__args__) == 2 and issubclass(klass.__args__[1], type(None))
 42 | 
 43 | 
 44 | def wraps(wrapper, wrappee):
 45 |     """Alternative to functools.wraps() for async functions."""  # noqa: D402
 46 |     wrapper.__name__ = wrappee.__name__
 47 |     wrapper.__qualname__ = wrappee.__qualname__
 48 |     wrapper.__module__ = wrappee.__module__
 49 |     wrapper.__doc__ = wrappee.__doc__
 50 |     wrapper.__annotations__ = wrappee.__annotations__
 51 |     wrapper.__wrapped__ = wrappee
 52 |     return wrapper
 53 | 
 54 | 
 55 | T = TypeVar("T")
 56 | 
 57 | 
 58 | def dataclass(cls: Optional[T] = None,
 59 |               /, *,
 60 |               slots=False,
 61 |               first_mutable: Optional[str] = None,
 62 |               **kwargs,
 63 |               ) -> Union[T, Type[Mapping[str, Any]]]:
 64 |     """
 65 |     Generate a dataclasses.dataclass with optional __slots__.
 66 | 
 67 |     :param slots: Define __slots__ according to the declared dataclass fields.
 68 |     :param first_mutable: First mutable field name. This and all the following fields will be \
 69 |                           considered mutable and optional. Such fields are not pickled and can be \
 70 |                           changed even though the instance is frozen.
 71 |     """
 72 |     def wrap(cls):
 73 |         cls = dataclasses.dataclass(cls, **kwargs)
 74 |         if slots:
 75 |             cls = _add_slots_to_dataclass(cls, first_mutable)
 76 |         return cls
 77 | 
 78 |     # See if we're being called as @dataclass or @dataclass().
 79 |     if cls is None:
 80 |         # We're called with parens.
 81 |         return wrap
 82 | 
 83 |     # We're called as @dataclass without parens.
 84 |     return wrap(cls)
 85 | 
 86 | 
 87 | # Caching context indicator. By default, we don't save the mutable optional fields.
 88 | _serialize_mutable_fields_in_dataclasses = ContextVar(
 89 |     "serialize_mutable_fields_in_dataclasses", default=False)
 90 | 
 91 | 
 92 | @contextmanager
 93 | def serialize_mutable_fields_in_dataclasses():
 94 |     """Provide a context manager to enable the serialization of mutable optional fields in our \
 95 |     dataclasses."""
 96 |     _serialize_mutable_fields_in_dataclasses.set(True)
 97 |     try:
 98 |         yield
 99 |     finally:
100 |         _serialize_mutable_fields_in_dataclasses.set(False)
101 | 
102 | 
103 | def _add_slots_to_dataclass(cls: T,
104 |                             first_mutable: Optional[str],
105 |                             ) -> Union[T, Type[Mapping[str, Any]]]:
106 |     """Set __slots__ of a dataclass, and make it a Mapping to compensate for a missing __dict__."""
107 |     # Need to create a new class, since we can't set __slots__ after a class has been created.
108 | 
109 |     # Make sure __slots__ isn't already set.
110 |     if "__slots__" in cls.__dict__:
111 |         raise TypeError(f"{cls.__name__} already specifies __slots__")
112 | 
113 |     # Create a new dict for our new class.
114 |     cls_dict = dict(cls.__dict__)
115 |     field_names = tuple(f.name for f in dataclasses.fields(cls))
116 |     cls_dict["__slots__"] = field_names
117 |     for field_name in field_names:
118 |         # Remove our attributes, if present. They"ll still be available in _MARKER.
119 |         cls_dict.pop(field_name, None)
120 |     # Remove __dict__ itself.
121 |     cls_dict.pop("__dict__", None)
122 |     # __hash__ cannot be inherited from SlotsMapping, IDK why.
123 |     if (hash_method := cls_dict.pop("__hash__", None)) is not None:
124 |         cls_dict["__hash__"] = hash_method
125 |     else:
126 |         def __hash__(self) -> int:
127 |             """Implement hash() over the immutable fields."""
128 |             return hash(tuple(
129 |                 (xxhash.xxh64_intdigest(x.view(np.uint8).data) if isinstance(x, np.ndarray) else x)
130 |                 for x in self.__getstate__()
131 |             ))
132 | 
133 |         cls_dict["__hash__"] = __hash__
134 |     qualname = getattr(cls, "__qualname__", None)
135 |     # Record the mutable fields.
136 |     if first_mutable is None:
137 |         first_mutable_index = len(field_names)
138 |     else:
139 |         first_mutable_index = field_names.index(first_mutable)
140 |     mutable_fields = set(field_names[first_mutable_index:])
141 |     if first_mutable is not None:
142 |         def __setattr__(self, attr: str, val: Any) -> None:
143 |             """Alternative to __setattr__ that works with mutable optional fields."""
144 |             assert attr in mutable_fields, "You can only change mutable optional fields."
145 |             object.__setattr__(self, attr, val)
146 | 
147 |         def make_with_attr(attr):
148 |             def with_attr(self, value) -> cls:
149 |                 """Chain __setattr__ to return `self`."""
150 |                 setattr(self, attr, value)
151 |                 return self
152 | 
153 |             return with_attr
154 | 
155 |         cls_dict["__setattr__"] = __setattr__
156 |         for attr in mutable_fields:
157 |             cls_dict["with_" + attr] = make_with_attr(attr)
158 | 
159 |     class SlotsMapping(Mapping[str, Any]):
160 |         """Satisfy Mapping abstractions by relying on the __slots__."""
161 | 
162 |         __slots__ = field_names
163 | 
164 |         def __getitem__(self, item: str) -> Any:
165 |             """Implement []."""
166 |             return getattr(self, item)
167 | 
168 |         def __len__(self) -> int:
169 |             """Implement len()."""
170 |             return len(self.__slots__)
171 | 
172 |         def __iter__(self) -> Iterator[str]:
173 |             """Implement iter()."""
174 |             return iter(self.__slots__)
175 | 
176 |         def __getstate__(self) -> Any:
177 |             """Support pickling back, we lost it when we deleted __dict__."""
178 |             include_mutable = _serialize_mutable_fields_in_dataclasses.get()
179 |             limit = len(self.__slots__) if include_mutable else first_mutable_index
180 |             return tuple(getattr(self, attr) for attr in self.__slots__[:limit])
181 | 
182 |         def __setstate__(self, state: Tuple[Any]) -> None:
183 |             """Construct a new class instance from the given `state`."""
184 |             for attr, val in zip(self.__slots__, state):
185 |                 object.__setattr__(self, attr, val)
186 |             # Fields with a default value.
187 |             if len(self.__slots__) > len(state):
188 |                 for field in dataclasses.fields(self)[len(state):]:
189 |                     object.__setattr__(self, field.name, field.default)
190 | 
191 |     # And finally create the class.
192 |     cls = type(cls)(cls.__name__, (SlotsMapping, *cls.__bases__), cls_dict)
193 |     if qualname is not None:
194 |         cls.__qualname__ = qualname
195 |     return cls
196 | 
197 | 
198 | NST = TypeVar("NST")
199 | 
200 | 
201 | class NumpyStruct(Mapping[str, Any]):
202 |     """
203 |     Constrained dataclass based on numpy structured array.
204 | 
205 |     We divide the fields into two groups: mutable and immutable.
206 |     The mutable fields are stored as regular class members and discarded from serialization.
207 |     The immutable fields are not materialized explicitly. Instead, they are taken from numpy
208 |     structured array (`_arr`) that references an arbitrary memory buffer (`_data`).
209 |     Serialization of the class is as simple as exposing the underlying memory buffer outside.
210 | 
211 |     We support variable-length sub-arrays using the special notation `[<array dtype>]`. That way
212 |     the arrays are appended to `_data`, and `_arr` points to them by pairs (offset, length).
213 |     """
214 | 
215 |     dtype: np.dtype
216 |     nested_dtypes: Mapping[str, np.dtype]
217 | 
218 |     def __init__(self, data: Union[bytes, bytearray, memoryview, np.ndarray], **optional: Any):
219 |         """Initialize a new instance of NumpyStruct from raw memory and the (perhaps incomplete) \
220 |         mapping of mutable field values."""
221 |         if isinstance(data, (np.ndarray, np.void)):
222 |             assert data.shape == () or data.shape == (1,)
223 |             data = data.reshape(1)
224 |             self._data = data.view(np.uint8).data
225 |             self._arr = data
226 |         else:
227 |             self._data = data
228 |             self._arr = None
229 |         for attr in self.__slots__[2:]:
230 |             setattr(self, attr, optional.get(attr))
231 | 
232 |     @classmethod
233 |     def from_fields(cls: NST, **kwargs: Any) -> NST:
234 |         """Initialize a new instance of NumpyStruct from the mapping of immutable field \
235 |         values."""
236 |         arr = np.zeros(1, cls.dtype)
237 |         extra_bytes = []
238 |         offset = cls.dtype.itemsize
239 |         for field_name, (field_dtype, _) in cls.dtype.fields.items():
240 |             value = kwargs.pop(field_name)
241 |             try:
242 |                 nested_dtype = cls.nested_dtypes[field_name]
243 |             except KeyError:
244 |                 if value is None and field_dtype.char in ("S", "U"):
245 |                     value = ""
246 |                 if field_dtype.char == "M" and isinstance(value, datetime):
247 |                     value = value.replace(tzinfo=None)
248 |                 arr[field_name] = np.asarray(value, field_dtype)
249 |             else:
250 |                 if is_str := ((is_ascii := _dtype_is_ascii(nested_dtype)) or
251 |                               nested_dtype.char in ("S", "U")):
252 |                     if isinstance(value, np.ndarray):
253 |                         if value.dtype == np.dtype(object):
254 |                             nan_mask = value == np.array([None])
255 |                         else:
256 |                             nan_mask = np.full(len(value), False)
257 |                     else:
258 |                         nan_mask = np.fromiter((v is None for v in value),
259 |                                                dtype=np.bool_, count=len(value))
260 |                     if is_ascii:
261 |                         nested_dtype = np.dtype("S")
262 |                 value = np.asarray(value, nested_dtype)
263 |                 assert len(value.shape) == 1, "we don't support arrays of more than 1 dimension"
264 |                 if is_str and nan_mask.any():
265 |                     if not value.flags.writeable:
266 |                         value = value.copy()
267 |                     value[nan_mask] = ""
268 |                 extra_bytes.append(data := value.view(np.byte).data)
269 |                 pointer = [offset, len(value)]
270 |                 if is_str and (is_ascii or nested_dtype.itemsize == 0):
271 |                     pointer.append(
272 |                         value.dtype.itemsize // np.dtype(nested_dtype.char + "1").itemsize)
273 |                 arr[field_name] = pointer
274 |                 offset += len(data)
275 |         if not extra_bytes:
276 |             return cls(arr.view(np.byte).data)
277 |         return cls(b"".join(chain([arr.view(np.byte).data], extra_bytes)), **kwargs)
278 | 
279 |     @property
280 |     def data(self) -> bytes:
281 |         """Return the underlying memory."""
282 |         return self._data
283 | 
284 |     @property
285 |     def array(self) -> np.ndarray:
286 |         """Return the underlying numpy array that wraps `data`."""
287 |         if self._arr is None:
288 |             self._arr = np.frombuffer(self.data, self.dtype, count=1)
289 |         return self._arr
290 | 
291 |     @property
292 |     def coerced_data(self) -> memoryview:
293 |         """Return prefix of `data` with nested immutable objects excluded."""
294 |         return memoryview(self.data)[:self.dtype.itemsize]
295 | 
296 |     def __getitem__(self, item: str) -> Any:
297 |         """Implement self[]."""
298 |         return getattr(self, item)
299 | 
300 |     def __setitem__(self, key: str, value: Any) -> None:
301 |         """Implement self[] = ..."""
302 |         setattr(self, key, value)
303 | 
304 |     def __len__(self) -> int:
305 |         """Implement len()."""
306 |         return len(self.dtype) + len(self.__slots__) - 2
307 | 
308 |     def __iter__(self) -> Iterator[str]:
309 |         """Implement iter()."""
310 |         return iter(chain(self.dtype.names, self.__slots__[2:]))
311 | 
312 |     def __hash__(self) -> int:
313 |         """Implement hash()."""
314 |         return hash(self._data)
315 | 
316 |     def __str__(self) -> str:
317 |         """Format for human-readability."""
318 |         return "{\n\t%s\n}" % ",\n\t".join("%s: %s" % (k, v) for k, v in self.items())
319 | 
320 |     def __repr__(self) -> str:
321 |         """Implement repr()."""
322 |         kwargs = {k: v for k in self.__slots__[2:] if (v := getattr(self, k)) is not None}
323 |         if kwargs:
324 |             kwargs_str = ", ".join(f"{k}={repr(v)}" for k, v in kwargs.items()) + ", "
325 |         else:
326 |             kwargs_str = ""
327 |         return f"{type(self).__name__}({kwargs_str}data={repr(self._data)})"
328 | 
329 |     def __eq__(self, other) -> bool:
330 |         """Compare this object to another."""
331 |         if self is other:
332 |             return True
333 | 
334 |         if self.__class__ is not other.__class__:
335 |             raise NotImplementedError(
336 |                 f"Cannot compare {self.__class__} and {other.__class__}")
337 | 
338 |         return self.data == other.data
339 | 
340 |     def __getstate__(self) -> Dict[str, Any]:
341 |         """Support pickle.dump()."""
342 |         data = self.data
343 |         return {
344 |             "data": bytes(data) if not isinstance(data, (bytes, bytearray)) else data,
345 |             **{attr: getattr(self, attr) for attr in self.__slots__[2:]},
346 |         }
347 | 
348 |     def __setstate__(self, state: Dict[str, Any]):
349 |         """Support pickle.load()."""
350 |         self.__init__(**state)
351 | 
352 |     def copy(self) -> "NumpyStruct":
353 |         """Clone the instance."""
354 |         return type(self)(self.data, **{attr: getattr(self, attr) for attr in self.__slots__[2:]})
355 | 
356 |     @staticmethod
357 |     def _generate_get(name: str,
358 |                       type_: Union[str, np.dtype, List[Union[str, np.dtype]]],
359 |                       ) -> Callable[["NumpyStruct"], Any]:
360 |         if _dtype_is_ascii(type_):
361 |             type_ = str
362 |         elif isinstance(type_, list):
363 |             type_ = np.ndarray
364 |         elif (char := np.dtype(type_).char) == "U":
365 |             type_ = np.str_
366 |         elif char == "S":
367 |             type_ = np.bytes_
368 |         elif char == "V":
369 |             type_ = np.ndarray
370 | 
371 |         def get_field(self) -> Optional[type_]:
372 |             if self._arr is None:
373 |                 self._arr = np.frombuffer(self.data, self.dtype, count=1)
374 |             value = self._arr[name][0]
375 |             if (nested_dtype := self.nested_dtypes.get(name)) is None:
376 |                 if value != value:
377 |                     return None
378 |                 if type_ is str:
379 |                     value = value.decode()
380 |                 if type_ in (str, np.str_):
381 |                     value = value or None
382 |                 return value
383 |             if (_dtype_is_ascii(nested_dtype) and (char := "S")) or \
384 |                     ((char := nested_dtype.char) in ("S", "U") and nested_dtype.itemsize == 0):
385 |                 offset, count, itemsize = value
386 |                 nested_dtype = f"{char}{itemsize}"
387 |             else:
388 |                 offset, count = value
389 |             return np.frombuffer(self.data, nested_dtype, offset=offset, count=count)
390 | 
391 |         get_field.__name__ = name
392 |         return get_field
393 | 
394 | 
395 | def _dtype_is_ascii(dtype: Union[str, np.dtype]) -> bool:
396 |     return (dtype is ascii) or (isinstance(dtype, str) and dtype.startswith("ascii"))
397 | 
398 | 
399 | def numpy_struct(cls):
400 |     """
401 |     Decorate a class to transform it to a NumpyStruct.
402 | 
403 |     The decorated class must define two sub-classes: `dtype` and `optional`.
404 |     The former annotates numpy-friendly immutable fields. The latter annotates mutable fields.
405 |     """
406 |     dtype = cls.Immutable.__annotations__
407 |     dtype_tuples = []
408 |     nested_dtypes = {}
409 |     for k, v in dtype.items():
410 |         if isinstance(v, list):
411 |             assert len(v) == 1, "Array must be specified as `[dtype]`."
412 |             nested_dtype = v[0]
413 |             if not (is_ascii := _dtype_is_ascii(nested_dtype)):
414 |                 nested_dtype = np.dtype(nested_dtype)
415 |             nested_dtypes[k] = nested_dtype
416 |             if is_ascii or (nested_dtype.char in ("S", "U") and nested_dtype.itemsize == 0):
417 |                 # save the characters count
418 |                 dtype_tuples.append((k, np.int32, 3))
419 |             else:
420 |                 dtype_tuples.append((k, np.int32, 2))
421 |         elif _dtype_is_ascii(v):
422 |             dtype_tuples.append((k, "S" + v[6:-1]))
423 |         else:
424 |             dtype_tuples.append((k, v))
425 |     try:
426 |         optional = cls.Optional.__annotations__
427 |     except AttributeError:
428 |         optional = {}
429 |     field_names = NamedTuple(
430 |         f"{cls.__name__}FieldNames",
431 |         [(k, str) for k in chain(dtype, optional)],
432 |     )(*chain(dtype, optional))
433 |     base = type(cls.__name__ + "Base", (NumpyStruct,),
434 |                 {k: property(NumpyStruct._generate_get(k, v)) for k, v in dtype.items()})
435 |     body = {
436 |         "__slots__": ("_data", "_arr", *optional),
437 |         "dtype": np.dtype(dtype_tuples),
438 |         "nested_dtypes": nested_dtypes,
439 |         "f": field_names,
440 |     }
441 |     struct_cls = type(cls.__name__, (cls, base), body)
442 |     struct_cls.__module__ = cls.__module__
443 |     cls.__name__ += "Origin"
444 |     return struct_cls
445 | 
446 | 
447 | @sentry_span
448 | def df_from_structs(items: Iterable[NumpyStruct],
449 |                     length: Optional[int] = None,
450 |                     ) -> pd.DataFrame:
451 |     """
452 |     Combine several NumpyStruct-s to a Pandas DataFrame.
453 | 
454 |     :param items: A collection, a generator, an iterator - all are accepted.
455 |     :param length: In case `items` does not support `len()`, specify the number of structs \
456 |                    for better performance.
457 |     :return: Pandas DataFrame with columns set to struct fields.
458 |     """
459 |     columns = {}
460 |     try:
461 |         if length is None:
462 |             length = len(items)
463 |     except TypeError:
464 |         # slower branch without pre-allocation
465 |         items_iter = iter(items)
466 |         try:
467 |             first_item = next(items_iter)
468 |         except StopIteration:
469 |             return pd.DataFrame()
470 |         assert isinstance(first_item, NumpyStruct)
471 |         dtype = first_item.dtype
472 |         nested_fields = first_item.nested_dtypes
473 |         coerced_datas = [first_item.coerced_data]
474 |         for k, v in first_item.items():
475 |             if k not in dtype.names or k in nested_fields:
476 |                 columns[k] = [v]
477 |         for item in items_iter:
478 |             coerced_datas.append(item.coerced_data)
479 |             for k in columns:
480 |                 columns[k].append(getattr(item, k))
481 |         table_array = np.frombuffer(b"".join(coerced_datas), dtype=dtype)
482 |         del coerced_datas
483 |     else:
484 |         items_iter = iter(items)
485 |         try:
486 |             first_item = next(items_iter)
487 |         except StopIteration:
488 |             return pd.DataFrame()
489 |         assert isinstance(first_item, NumpyStruct)
490 |         dtype = first_item.dtype
491 |         nested_fields = first_item.nested_dtypes
492 |         itemsize = dtype.itemsize
493 |         coerced_datas = bytearray(itemsize * length)
494 |         coerced_datas[:itemsize] = first_item.coerced_data
495 |         for k, v in first_item.items():
496 |             if k not in dtype.names or k in nested_fields:
497 |                 columns[k] = column = [None] * length
498 |                 column[0] = v
499 |         for i, item in enumerate(items_iter, 1):
500 |             coerced_datas[i * itemsize:(i + 1) * itemsize] = item.coerced_data
501 |             for k in columns:
502 |                 columns[k][i] = item[k]
503 |         table_array = np.frombuffer(coerced_datas, dtype=dtype)
504 |         del coerced_datas
505 |     for field_name in dtype.names:
506 |         if field_name not in nested_fields:
507 |             columns[field_name] = table_array[field_name]
508 |     del table_array
509 |     column_types = {}
510 |     try:
511 |         for k, v in first_item.Optional.__annotations__.items():
512 |             if not isinstance(v, type) or not issubclass(v, (datetime, np.datetime64, float)):
513 |                 # we can only unbox types that have a "NaN" value
514 |                 v = object
515 |             column_types[k] = v
516 |     except AttributeError:
517 |         pass  # no Optional
518 |     for k, v in columns.items():
519 |         column_type = column_types.get(k, object)
520 |         if issubclass(column_type, datetime):
521 |             v = tslib.array_to_datetime(np.array(v, dtype=object), utc=True, errors="raise")[0]
522 |         elif issubclass(column_type, timedelta):
523 |             v = np.array(v, dtype="timedelta64[s]")
524 |         elif np.dtype(column_type) != np.dtype(object):
525 |             v = np.array(v, dtype=column_type)
526 |         columns[k] = v
527 |     df = pd.DataFrame.from_dict(columns)
528 |     sentry_sdk.Hub.current.scope.span.description = str(len(df))
529 |     return df
530 | 


--------------------------------------------------------------------------------