/** * @brief The module-init glue — `PyModuleDef`, type registration, and `PyInit_stringzilla`. * @file python/stringzilla/stringzilla.c * @author Ash Vardanian * * - Doesn't use PyBind11, NanoBind, Boost.Python, or any other high-level libs, only CPython API. * - To minimize latency this implementation avoids `PyArg_ParseTupleAndKeywords` calls. * - Reimplements all of the `str` functionality in C as a `Str` type. * - Provides a highly generic `Strs` class for handling collections of strings, Arrow-style or not. * * Pandas doesn't provide a C API, and even in the 2.0 the Apache Arrow representation is opt-in, not default. * PyCapsule protocol in conjunction with @b `__arrow_c_array__` dunder methods can be used to extract strings. * @see https://arrow.apache.org/docs/python/generated/pyarrow.array.html * * This module exports C functions via `PyCapsule` of `PyAPI` for use by other extensions (like `stringzillas-cpus`): * - `sz_py_export_string_like`. * - `sz_py_export_strings_as_sequence`. * - `sz_py_export_strings_as_u32tape`. * - `sz_py_export_strings_as_u64tape`. * - `sz_py_replace_strings_allocator`. * * Function Naming Convention: * - `Str_like_*`: Functions that can be called both as module-level functions AND as member methods. * - `Str_*`: Functions that are member-only methods or have simpler calling conventions. * * This translation unit owns only the module-init glue: `PyModuleDef`, the type-registration table, and * `PyInit_stringzilla`. The `File`/`Str`/`Strs` struct layouts and the `PyTypeObject` forward declarations * every domain file needs live in `stringzilla.h`; the domains themselves are split across `shared.c`, `file.c`, * `str.c`, `strs.c`, `memory.c`, `hash.c`, `cipher.c`, `find.c`, `compare.c`, `sort.c`, `intersect.c`, * and the `utf8_*.c` files. */ #include "stringzilla.h" /** * @brief The function table `stringzillas` reads out of this module's `_sz_py_api` capsule. * * `python/stringzillas.c` carries its own copy of this layout and casts the capsule pointer to it. The two * must stay identical: a field added, removed, or reordered on one side alone makes the other misread memory * with no diagnostic, since the capsule carries no version tag. */ typedef struct PyAPI { sz_bool_t (*sz_py_export_string_like)(PyObject *, sz_cptr_t *, sz_size_t *); sz_bool_t (*sz_py_export_strings_as_sequence)(PyObject *, sz_sequence_t *); sz_bool_t (*sz_py_export_strings_as_u32tape)(PyObject *, sz_cptr_t *, sz_u32_t const **, sz_size_t *); sz_bool_t (*sz_py_export_strings_as_u64tape)(PyObject *, sz_cptr_t *, sz_u64_t const **, sz_size_t *); sz_bool_t (*sz_py_replace_strings_allocator)(PyObject *, sz_memory_allocator_t *); } PyAPI; /** @brief Builds a tuple of the capability names in @p caps, or @c NULL with a Python exception set. */ static PyObject *capabilities_to_tuple(sz_capability_t caps) { sz_cptr_t cap_strings[SZ_CAPABILITIES_COUNT]; sz_size_t cap_count = sz_capabilities_to_strings_implementation_(caps, cap_strings, SZ_CAPABILITIES_COUNT); PyObject *caps_tuple = PyTuple_New(cap_count); if (!caps_tuple) return NULL; for (sz_size_t i = 0; i < cap_count; i++) { PyObject *cap_str = PyUnicode_FromString(cap_strings[i]); if (!cap_str) { Py_DECREF(caps_tuple); return NULL; } PyTuple_SET_ITEM(caps_tuple, i, cap_str); } return caps_tuple; } static int parse_and_intersect_capabilities(PyObject *caps_obj, sz_capability_t *result) { if (!caps_obj) { PyErr_SetString(PyExc_TypeError, "capabilities must be a tuple or list of strings"); return -1; } // A tuple snapshot rather than `PySequence_Fast`, which hands back the caller's list itself and leaves // the walk indexing into storage another thread can resize. PyObject *seq = PySequence_Tuple(caps_obj); if (!seq) { PyErr_SetString(PyExc_TypeError, "capabilities must be a tuple or list of strings"); return -1; } sz_capability_t requested_caps = 0; Py_ssize_t n = PyTuple_GET_SIZE(seq); for (Py_ssize_t i = 0; i < n; i++) { PyObject *item = PyTuple_GET_ITEM(seq, i); if (!PyUnicode_Check(item)) { PyErr_SetString(PyExc_TypeError, "capabilities must be strings"); Py_DECREF(seq); return -1; } char const *cap_str = PyUnicode_AsUTF8(item); if (!cap_str) { Py_DECREF(seq); return -1; } sz_capability_t flag = sz_capability_from_string_implementation_(cap_str); if (flag == sz_caps_none_k) { PyErr_Format(PyExc_ValueError, "Unknown capability: %s", cap_str); Py_DECREF(seq); return -1; } requested_caps |= flag; } Py_DECREF(seq); // An empty request or one entirely outside this machine's reach is a configuration error; silently // degrading to serial would hide it. *result = requested_caps & sz_capabilities(); if (*result == 0) { PyErr_Format(PyExc_ValueError, "No requested capability is available here; available: %s", sz_capabilities_to_string_implementation_(sz_capabilities())); return -1; } return 0; } static char const doc_reset_capabilities[] = // "reset_capabilities(names) -> None\n\n" // "Sets the active SIMD/backend capabilities for this module and updates the\n" // "runtime dispatch table. The provided names are intersected with hardware\n" // "capabilities; raises ValueError if none of them is available.\n\n" // "Side effects: updates stringzilla.__capabilities__ and __capabilities_str__.\n" // "\n" // "Example:\n" // " >>> sz.reset_capabilities(['serial']) # restrict dispatch to the scalar backend"; static PyObject *module_reset_capabilities(PyObject *self, PyObject *args) { PyObject *caps_obj = NULL; if (!PyArg_ParseTuple(args, "O", &caps_obj)) return NULL; sz_capability_t caps = 0; if (parse_and_intersect_capabilities(caps_obj, &caps) != 0) return NULL; // Update the dispatch table sz_dispatch_table_update(caps); // Recompute and set module-level capability exports PyObject *caps_tuple = capabilities_to_tuple(caps); if (!caps_tuple) return NULL; if (PyObject_SetAttrString(self, "__capabilities__", caps_tuple) != 0) { Py_DECREF(caps_tuple); return NULL; } Py_DECREF(caps_tuple); sz_cptr_t caps_str = sz_capabilities_to_string(caps); if (PyObject_SetAttrString(self, "__capabilities_str__", PyUnicode_FromString(caps_str)) != 0) { return NULL; } Py_RETURN_NONE; } static void stringzilla_cleanup(PyObject *m) { // Drain both free-lists, releasing the headers parked for reuse during the interpreter's lifetime. stringzilla_state_t *state = (stringzilla_state_t *)PyModule_GetState(m); if (!state) return; for (Str *node = state->str_freelist_head; node;) { Str *next = (Str *)node->parent; PyObject_Free(node); node = next; } for (Strs *node = state->strs_freelist_head; node;) { Strs *next = *Strs_freelist_next_(node); PyObject_Free(node); node = next; } state->str_freelist_head = NULL; state->str_freelist_count = 0; state->strs_freelist_head = NULL; state->strs_freelist_count = 0; } static PyMethodDef stringzilla_methods[] = { // Basic `str`, `bytes`, and `bytearray`-like functionality {"contains", (PyCFunction)Str_like_contains, SZ_METHOD_FLAGS, doc_contains}, {"count", (PyCFunction)Str_like_count, SZ_METHOD_FLAGS, doc_count}, {"splitlines", (PyCFunction)Str_like_splitlines, SZ_METHOD_FLAGS, doc_splitlines}, {"startswith", (PyCFunction)Str_like_startswith, SZ_METHOD_FLAGS, doc_startswith}, {"endswith", (PyCFunction)Str_like_endswith, SZ_METHOD_FLAGS, doc_endswith}, {"decode", (PyCFunction)Str_like_decode, SZ_METHOD_FLAGS, doc_decode}, {"equal", (PyCFunction)Str_like_equal, SZ_METHOD_FLAGS, doc_like_equal}, {"lstrip", (PyCFunction)Str_like_lstrip, SZ_METHOD_FLAGS, doc_lstrip}, {"rstrip", (PyCFunction)Str_like_rstrip, SZ_METHOD_FLAGS, doc_rstrip}, {"strip", (PyCFunction)Str_like_strip, SZ_METHOD_FLAGS, doc_strip}, // Bidirectional operations {"find", (PyCFunction)Str_like_find, SZ_METHOD_FLAGS, doc_find}, {"index", (PyCFunction)Str_like_index, SZ_METHOD_FLAGS, doc_index}, {"partition", (PyCFunction)Str_like_partition, SZ_METHOD_FLAGS, doc_partition}, {"split", (PyCFunction)Str_like_split, SZ_METHOD_FLAGS, doc_split}, {"rfind", (PyCFunction)Str_like_rfind, SZ_METHOD_FLAGS, doc_rfind}, {"rindex", (PyCFunction)Str_like_rindex, SZ_METHOD_FLAGS, doc_rindex}, {"rpartition", (PyCFunction)Str_like_rpartition, SZ_METHOD_FLAGS, doc_rpartition}, {"rsplit", (PyCFunction)Str_like_rsplit, SZ_METHOD_FLAGS, doc_rsplit}, // Character search extensions {"find_first_of", (PyCFunction)Str_like_find_first_of, SZ_METHOD_FLAGS, doc_find_first_of}, {"find_last_of", (PyCFunction)Str_like_find_last_of, SZ_METHOD_FLAGS, doc_find_last_of}, {"find_first_not_of", (PyCFunction)Str_like_find_first_not_of, SZ_METHOD_FLAGS, doc_find_first_not_of}, {"find_last_not_of", (PyCFunction)Str_like_find_last_not_of, SZ_METHOD_FLAGS, doc_find_last_not_of}, {"count_byteset", (PyCFunction)Str_like_count_byteset, SZ_METHOD_FLAGS, doc_count_byteset}, {"split_byteset", (PyCFunction)Str_like_split_byteset, SZ_METHOD_FLAGS, doc_split_byteset}, {"rsplit_byteset", (PyCFunction)Str_like_rsplit_byteset, SZ_METHOD_FLAGS, doc_rsplit_byteset}, // Lazily evaluated iterators {"split_iter", (PyCFunction)Str_like_split_iter, SZ_METHOD_FLAGS, doc_split_iter}, {"rsplit_iter", (PyCFunction)Str_like_rsplit_iter, SZ_METHOD_FLAGS, doc_rsplit_iter}, {"split_byteset_iter", (PyCFunction)Str_like_split_byteset_iter, SZ_METHOD_FLAGS, doc_split_byteset_iter}, {"rsplit_byteset_iter", (PyCFunction)Str_like_rsplit_byteset_iter, SZ_METHOD_FLAGS, doc_rsplit_byteset_iter}, // UTF-8 aware operations {"utf8_count", (PyCFunction)Str_like_utf8_count, SZ_METHOD_FLAGS, doc_utf8_count}, {"utf8_split_newlines", (PyCFunction)Str_like_utf8_split_newlines, SZ_METHOD_FLAGS, doc_utf8_split_newlines}, {"utf8_newlines", (PyCFunction)Str_like_utf8_newlines, SZ_METHOD_FLAGS, doc_utf8_newlines}, {"utf8_split_whitespaces", (PyCFunction)Str_like_utf8_split_whitespaces, SZ_METHOD_FLAGS, doc_utf8_split_whitespaces}, {"utf8_whitespaces", (PyCFunction)Str_like_utf8_whitespaces, SZ_METHOD_FLAGS, doc_utf8_whitespaces}, {"utf8_split_delimiters", (PyCFunction)Str_like_utf8_split_delimiters, SZ_METHOD_FLAGS, doc_utf8_split_delimiters}, {"utf8_delimiters", (PyCFunction)Str_like_utf8_delimiters, SZ_METHOD_FLAGS, doc_utf8_delimiters}, {"utf8_wordbreaks", (PyCFunction)Str_like_utf8_wordbreaks, SZ_METHOD_FLAGS, doc_utf8_wordbreaks}, {"utf8_codepoints", (PyCFunction)Str_like_utf8_codepoints, SZ_METHOD_FLAGS, doc_utf8_codepoints}, {"utf8_graphemes", (PyCFunction)Str_like_utf8_graphemes, SZ_METHOD_FLAGS, doc_utf8_graphemes}, {"utf8_sentences", (PyCFunction)Str_like_utf8_sentences, SZ_METHOD_FLAGS, doc_utf8_sentences}, {"utf8_linebreaks", (PyCFunction)Str_like_utf8_linebreaks, SZ_METHOD_FLAGS, doc_utf8_linebreaks}, {"utf8_uncased_fold", (PyCFunction)Str_like_utf8_uncased_fold, SZ_METHOD_FLAGS, doc_utf8_uncased_fold}, {"utf8_norm", (PyCFunction)Str_like_utf8_norm, SZ_METHOD_FLAGS, doc_utf8_norm}, {"utf8_find_denormalized", (PyCFunction)Str_like_utf8_find_denormalized, SZ_METHOD_FLAGS, doc_utf8_find_denormalized}, {"utf8_uncased_search", (PyCFunction)Str_like_utf8_uncased_search, SZ_METHOD_FLAGS, doc_utf8_uncased_search}, {"utf8_uncased_matches", (PyCFunction)Str_like_utf8_uncased_matches, SZ_METHOD_FLAGS, doc_utf8_uncased_matches}, {"utf8_uncased_order", (PyCFunction)Str_like_utf8_uncased_order, SZ_METHOD_FLAGS, doc_utf8_uncased_order}, // Dealing with larger-than-memory datasets {"offset_within", (PyCFunction)Str_offset_within, SZ_METHOD_FLAGS, doc_offset_within}, {"write_to", (PyCFunction)Str_write_to, SZ_METHOD_FLAGS, doc_write_to}, // In-place transforms {"translate", (PyCFunction)Str_like_translate, SZ_METHOD_FLAGS, doc_translate}, // Global unary extensions {"hash", (PyCFunction)Str_like_hash, SZ_METHOD_FLAGS, doc_like_hash}, {"hash_multiseed", (PyCFunction)Str_like_hash_multiseed, SZ_METHOD_FLAGS, doc_hash_multiseed}, {"bytesum", (PyCFunction)Str_like_bytesum, SZ_METHOD_FLAGS, doc_like_bytesum}, {"sha256", (PyCFunction)Str_like_sha256, SZ_METHOD_FLAGS, doc_like_sha256}, {"hmac_sha256", (PyCFunction)hmac_sha256, SZ_METHOD_FLAGS, doc_hmac_sha256}, {"fill_random", (PyCFunction)Str_like_fill_random, SZ_METHOD_FLAGS, doc_fill_random}, // Module-level functionality {"random", (PyCFunction)module_random, SZ_METHOD_FLAGS, doc_random}, {"reset_capabilities", (PyCFunction)module_reset_capabilities, METH_VARARGS, doc_reset_capabilities}, {NULL, NULL, 0, NULL}}; PyModuleDef stringzilla_module = { PyModuleDef_HEAD_INIT, "stringzilla", "Search, hash, sort, fingerprint, and fuzzy-match strings faster via SWAR, SIMD, and GPGPU", sizeof(stringzilla_state_t), // Per-interpreter free-list state; also enables `PyState_FindModule`. stringzilla_methods, NULL, NULL, NULL, stringzilla_cleanup, }; /** * @brief Every type the module readies, and the name it is exported under. * * A `NULL` name means the type is readied but not exported: the UTF-8 iterators are returned by methods * and never constructed by name, so they need `tp_dict` filled but no module attribute. */ static struct { char const *name; PyTypeObject *type; } const stringzilla_types[] = { {"Str", &StrType}, {"File", &FileType}, {"Strs", &StrsType}, {"FindSplits", &FindSplitsType}, {"Utf8SplitNewlines", &Utf8SplitNewlinesType}, {"Utf8Newlines", &Utf8NewlinesType}, {"Utf8SplitWhitespaces", &Utf8SplitWhitespacesType}, {"Utf8Whitespaces", &Utf8WhitespacesType}, {"Utf8SplitDelimiters", &Utf8SplitDelimitersType}, {"Utf8Delimiters", &Utf8DelimitersType}, {"Utf8Wordbreaks", &Utf8WordbreaksType}, {NULL, &Utf8CodepointsType}, {NULL, &Utf8GraphemesType}, {NULL, &Utf8SentencesType}, {NULL, &Utf8LinebreaksType}, {NULL, &Utf8UncasedMatchesType}, {"Hasher", &HasherType}, {"Sha256", &Sha256Type}, {"Sha256s", &Sha256sType}, {"Aes256CtrKey", &Aes256CtrKeyType}, {"Aes256GcmKey", &Aes256GcmKeyType}, {"Aes256GcmEncryptor", &Aes256GcmEncryptorType}, {"Aes256GcmDecryptor", &Aes256GcmDecryptorType}, }; PyMODINIT_FUNC PyInit_stringzilla(void) { PyObject *m; sz_size_t const types_count = sizeof(stringzilla_types) / sizeof(stringzilla_types[0]); for (sz_size_t type_index = 0; type_index != types_count; ++type_index) if (PyType_Ready(stringzilla_types[type_index].type) < 0) return NULL; m = PyModule_Create(&stringzilla_module); if (m == NULL) return NULL; #ifdef Py_GIL_DISABLED // Declare that this module is safe for free-threaded Python PyUnstable_Module_SetGIL(m, Py_MOD_GIL_NOT_USED); #endif // Add version metadata { char version_str[50]; sprintf(version_str, "%d.%d.%d", sz_version_major(), sz_version_minor(), sz_version_patch()); PyModule_AddStringConstant(m, "__version__", version_str); } // Publish the digest width on both hasher types, so callers can size an output matrix without // hardcoding it. A class attribute rather than a property, as it describes the algorithm. { PyObject *digest_length = PyLong_FromSize_t(SZ_SHA256_DIGEST_LENGTH); if (!digest_length) goto failed; int const published = PyDict_SetItemString(Sha256Type.tp_dict, "digest_length", digest_length) | PyDict_SetItemString(Sha256sType.tp_dict, "digest_length", digest_length); Py_DECREF(digest_length); if (published < 0) goto failed; } // Define SIMD capabilities as a tuple { sz_capability_t caps = sz_capabilities(); PyObject *caps_tuple = capabilities_to_tuple(caps); if (!caps_tuple) goto failed; if (PyModule_AddObject(m, "__capabilities__", caps_tuple) < 0) { Py_DECREF(caps_tuple); goto failed; } // Also keep the old comma-separated string version for backward compatibility sz_cptr_t caps_str = sz_capabilities_to_string(caps); PyModule_AddStringConstant(m, "__capabilities_str__", caps_str); } for (sz_size_t type_index = 0; type_index != types_count; ++type_index) { char const *const name = stringzilla_types[type_index].name; if (!name) continue; PyTypeObject *const type = stringzilla_types[type_index].type; Py_INCREF(type); // Only a successful `PyModule_AddObject` steals the reference, so undo just this one and let the // dying module release the types it already owns. if (PyModule_AddObject(m, name, (PyObject *)type) < 0) { Py_DECREF(type); goto failed; } } // A refused tag needs a class of its own, created here because a static exception object cannot // be built before the interpreter exists. AuthenticationErrorType = PyErr_NewExceptionWithDoc("stringzilla.AuthenticationError", doc_AuthenticationError, PyExc_ValueError, NULL); if (!AuthenticationErrorType) goto failed; Py_INCREF(AuthenticationErrorType); if (PyModule_AddObject(m, "AuthenticationError", AuthenticationErrorType) < 0) { Py_DECREF(AuthenticationErrorType); // The reference the module refused to take Py_DECREF(AuthenticationErrorType); // The reference this translation unit holds AuthenticationErrorType = NULL; goto failed; } // Export C API functions as a single capsule structure for StringZillas static PyAPI sz_py_api = { .sz_py_export_string_like = sz_py_export_string_like, .sz_py_export_strings_as_sequence = sz_py_export_strings_as_sequence, .sz_py_export_strings_as_u32tape = sz_py_export_strings_as_u32tape, .sz_py_export_strings_as_u64tape = sz_py_export_strings_as_u64tape, .sz_py_replace_strings_allocator = sz_py_replace_strings_allocator, }; if (PyModule_AddObject(m, "_sz_py_api", PyCapsule_New(&sz_py_api, "_sz_py_api", NULL)) < 0) goto failed; return m; failed: Py_DECREF(m); return NULL; }