NumPy 用户自定义 DType 的 Partition / Argpartition 支持：基于 ArrayMethod API 实现自定义分区算法-北京尧图网络

NumPy 用户自定义 DType 的 Partition / Argpartition 支持基于 ArrayMethod API 实现自定义分区算法【免费下载链接】numpyThe fundamental package for scientific computing with Python.项目地址: https://gitcode.com/gh_mirrors/nu/numpy本篇技术指南以 NumPy 官方发布说明doc/release/upcoming_changes/31614.new_feature.rst为切入点系统讲解用户自定义 DType 如何通过 ArrayMethod API 为numpy.partition与numpy.argpartition注册自定义实现覆盖注册入口、Spec 契约、参数结构与底层调用链。读完本文你将掌握如何为自己的自定义 DType 编写可被 NumPy 分区操作调用的原生实现并理解其与排序sort/argsort扩展机制的异同。特性概览分区操作正式接入 ArrayMethod APINumPy 的用户自定义 DTypeUser-Defined DType以下简称 UDType机制允许第三方扩展定义全新的数据类型并借助 ArrayMethod API 为各类 ufunc 与数组方法提供原生实现。此前自定义 DType 已经可以通过 ArrayMethod API 注册sort与argsort实现相关细节见 C API 文档中的 Sorting and Argsorting 一节。本次新增特性在此基础上进一步扩展用户自定义 DType 现在可以像实现排序一样为partition就地分区与argpartition分区索引注册自定义实现这些实现通过numpy.partition与numpy.argpartition以及对应的ndarray.partition、ndarray.argpartition方法在传入用户自定义 DType 数组时被自动调用注册入口与排序一致将实现对应操作的 ArrayMethod spec 传递给PyUFunc_AddLoopsFromSpecs函数即可无需额外的专用注册 API。从源码结构看该能力由 numpy/_core/src/multiarray/item_selection.c 中的分区分发逻辑支撑并在 numpy/_core/tests/test_custom_dtypes.py 中有完整的测试覆盖属于可开箱使用的正式功能而非实验性接口。注册入口PyUFunc_AddLoopsFromSpecs 与 PyUFunc_LoopSlot注册自定义 DType 的 partition / argpartition 实现核心入口是PyUFunc_AddLoopsFromSpecs自 NumPy 2.4 起引入见 C API 文档。该函数接收一个以NULL结尾的PyUFunc_LoopSlot数组每个槽位描述为哪一个操作注册哪一个 ArrayMethod spec。PyUFunc_LoopSlot结构体包含两个成员定义见 C API 文档成员类型说明nameconst char *要注册到的操作名称格式与 entry point 类似(module :)? (object .)* name默认模块为numpy。例如sin、strings.str_len、numpy.strings:str_len。注意部分名称并不直接对应 ufunc如sort、argsort、real、imag——它们在内部使用 ufunc 或 ufunc-like 机制实现partition与argpartition同样属于这一类。specPyArrayMethod_Spec *用于创建 loop 的 ArrayMethod spec。分区实现即通过将name设为partition或argpartition、spec设为对应实现来注册。注册示意如下static PyArrayMethod_Spec partition_spec { .nin 2, .nout 1, .dtypes partition_dtypes, // {你的 DType, NPY_INTP, 你的 DType} .slots partition_slots, // 例如 NPY_METH_resolve_descriptors 与 NPY_METH_get_loop .flags NPY_METH_NO_FLOATINGPOINT_ERRORS, }; PyUFunc_LoopSlot loops[] { {partition, partition_spec}, {argpartition, argpartition_spec}, {NULL, NULL}, // NULL 结尾 }; int res PyUFunc_AddLoopsFromSpecs(loops);与 sort / argsort 注册的异同相同点partition、argpartition与sort、argsort一样都是通过PyUFunc_AddLoopsFromSpecs批量注册可与常规 ufunc loops 一起传递不同点sort 系列的 spec 要求nin1, nout1单输入单输出排序就地完成而 partition 系列的 spec 要求nin2, nout1——两个输入分别是被分区数组与分区索引kth数组详见下文契约。Partition Spec 契约输入、输出与就地语义根据 C API 文档 Partitioning and Argpartitioning 一节注册 partition / argpartition 的 ArrayMethod spec 必须遵守以下契约nin2, nout1第一个输入data[0]要分区的数组第二个输入data[1]kth 索引数组即按哪些位置进行分区输出仅 argpartition分区后的索引数组。就地分区约束partition 是就地操作因此强制要求data[0] data[2]第一个输入与输出共享同一块内存/数据指针。kth 数组的形态约束data[1]始终是一个NPY_INTP类型的连续contiguous数组内含分区索引。若传入多个分区索引则数组会对每个索引依次分区多个 kth 值会先被排序以确保分区互不干扰。argpartition 的输出类型argpartition 返回的是新分配的索引数组因此输出必须是NPY_INTP类型。循环维度信息在 strided loop 中dimensions[0]表示要分区的元素个数dimensions[1]表示分区索引的个数。与 Python 层语义的对应上述契约与numpy.partition/numpy.argpartition的既有 Python 语义完全一致kth支持负索引会被转换为shape[axis] kth索引越界会抛出ValueErrorkind参数接受introselect默认与stable两种选择算法并可配合descendingTrue进行降序分区该参数在 NumPy 2.0 引入的稳定选择机制上提供。参数结构PyArrayMethod_PartitionParameters 与 NPY_SELECTKIND与 sort 系列通过PyArrayMethod_SortParameters传递排序标志类似分区操作通过PyArrayMethod_PartitionParameters将分区语义信息传给 looptypedef struct { NPY_SELECTKIND flags; } PyArrayMethod_PartitionParameters;该结构通过 loop 收到的context-parameters字段访问context即PyArrayMethod_Contextflags是NPY_SELECTKIND枚举值的按位或指示本次分区执行的种类——关键语义是是否为降序分区。NPY_SELECTKIND定义在 numpy/_core/include/numpy/ndarraytypes.h#L199-L207typedef enum { _NPY_SELECT_UNDEFINED -1, NPY_INTROSELECT 0, // new style names NPY_SELECT_DEFAULT 0, NPY_SELECT_STABLE 2, NPY_SELECT_DESCENDING 4, } NPY_SELECTKIND; #define NPY_NSELECTS (NPY_SELECT_DESCENDING 1)即合法取值包括NPY_SELECT_DEFAULT等价于NPY_INTROSELECT默认选择算法、NPY_SELECT_STABLE稳定分区与NPY_SELECT_DESCENDING降序分区NPY_NSELECTS定义为 5。编写 loop 的关键建议来自官方文档如果 strided loop 的实现依赖flags例如需要区分升降序最佳实践是只定义NPY_METH_get_loop槽位而不设置其他 loop slots。这样可以在get_strided_loop阶段根据context-parameters中的 flags 动态选择/生成合适的内层循环避免在单一路径中处理所有分支。源码级实现剖析分区分发调用链理解注册契约后再看底层分发逻辑可以印证上述约定的每一个细节。核心实现在 numpy/_core/src/multiarray/item_selection.c。1. kth 数组预处理partition_prep_kth_array无论是 partition 还是 argpartition入口都会先调用partition_prep_kth_arrayitem_selection.c#L1666-L1720对 kth 数组做规范化类型检查拒绝布尔数组ValueError: Booleans unacceptable as partition index非整数数组抛TypeError: Partition index must be integer维度检查kth 数组维度必须 1强制转型将 kth 数组转换为NPY_INTP类型PyArray_Cast(ktharray, NPY_INTP)负索引归一化将负索引加上shape[axis]并做越界检查空数组除外越界抛ValueError: kth(...) out of bounds (...)多索引排序当 kth 值多于一个时先对 kth 数组排序确保多次分区互不干扰——这与文档中对每个索引依次分区的契约对应。2. PyArray_Partition就地分区分发PyArray_Partitionitem_selection.c#L1727-L1808的流程清晰展示了 ArrayMethod 的接入方式校验分区种类which是否在[0, NPY_NSELECTS)范围内非法值抛ValueError: not a valid partition kind预处理 kth 数组见上从 DType 槽位取方法method NPY_DT_SLOTS(NPY_DTYPE(PyArray_DESCR(op)))-part_meth;——这正是注册partitionspec 时被填充的槽位回退机制若method NULLDType 未注册 partition 实现回退到排序实现return PyArray_Sort(op, axis, (NPY_SORTKIND)which);注释明确Use sorting, slower but equivalent用排序代替较慢但结果等价。这意味着未注册自定义分区的 DType 依然可用 partition只是性能退化构造dtypes[3] {dt, kdt, dt}与given_descrs[3] {descr, kdescr, descr}——印证了data[0] data[2]输入与输出同类型同描述符的契约其中kdt是预处理后 kth 数组的NPY_INTPDType依次调用method-resolve_descriptors(...)与method-get_strided_loop(...)将PyArrayMethod_PartitionParameters通过context.parameters传入由于分区对每个轴切片都是连续的strides 直接按loop_descrs[i]-elsize构造注释Arrays are always contiguous for partitioning最终交给_new_sortlike执行多轴循环。3. PyArray_ArgPartition索引分区分发PyArray_ArgPartitionitem_selection.c#L1815-L1909与 partition 几乎对称差异在于方法取自NPY_DT_SLOTS(...)-argpart_meth槽位输出描述符固定为odescr PyArray_DescrFromType(NPY_INTP);dtypes 变为{dt, kdt, odt}——印证了输出必须是NPY_INTP的契约回退逻辑改为调用PyArray_ArgSort(op2, axis, (NPY_SORTKIND)which)最终交给_new_argsortlike执行。4. 槽位与分发的对应关系从 numpy/_core/code_generators/numpy_api.py 与 numpy/_core/src/umath/dispatching.cpp 中PyUFunc_AddLoopsFromSpecs的实现来看注册系统在解析partition/argpartition名称时会将其路由到 DType 的part_meth/argpart_meth槽位与sort/argsort的sort_meth/argsort_meth机制同构。这也解释了为什么文档明确提示以与排序和 argsorting 类似的方式实现。降序分区的处理以 test_custom_dtypes 为参考仓库中的自定义 DType 测试类numpy/_core/tests/test_custom_dtypes.py#L448-L512对 partition 与 argpartition 的注册实现做了全面验证是学习编写实现的极佳范本。其测试要点包括test_partition构造以不同缩放因子表示的 float64 视图 DType 数组正序、逆序、非对齐 stride 等形态调用a.partition(k)后断言a[:k]与a[k:]各自排序后的内容符合预期同时覆盖descendingTrue的降序分区断言分区点前后元素集合互换test_argpartition对相同数组形态调用a.argpartition(k)通过返回索引indices回取元素并验证分区正确性同样覆盖降序模式。这两个测试同时验证了自定义 partition 实现确实被ndarray.partition/ndarray.argpartition分发调用就地语义partition 直接修改原数组与NPY_INTP输出类型约束成立非连续stride 反向与非对齐数组下实现依然正确。版本与适用性说明PyUFunc_AddLoopsFromSpecs自NumPy 2.4起引入.. versionadded:: 2.4本特性为后续迭代中为partition/argpartition扩展的支持官方 C API 文档doc/source/reference/c-api/array.rst是了解 ArrayMethod spec 各字段resolve_descriptors、get_strided_loop、PyArrayMethod_Spec、PyArrayMethod_Context等的一手资料当前仓库的 2.5.0 发布说明doc/source/release/2.5.0-notes.rst还显示类型标注层面已为argpartition提供 shape 相关的返回类型推断shape-typing说明该 API 已被视为稳定的公开能力并持续获得工具链完善。小结DType 的 partition / argpartition ArrayMethod 支持为自定义数据类型在保持 NumPy 统一 API 语义的前提下接入高性能原生分区算法铺平了道路。核心要点可归纳为四条通过PyUFunc_AddLoopsFromSpecs以partition/argpartition为名注册 spec与 sort / argsort 共用同一注册机制严格遵守nin2, nout1、data[0] data[2]就地、data[1]为NPY_INTP连续数组、argpartition 输出为NPY_INTP的契约借助context-parameters中的PyArrayMethod_PartitionParameters.flagsNPY_SELECTKIND判断升降序推荐只定义NPY_METH_get_loop槽位未注册实现的 DType 会自动回退到排序实现结果等价、性能较慢因此该特性是完全向后兼容的可选优化。【免费下载链接】numpyThe fundamental package for scientific computing with Python.项目地址: https://gitcode.com/gh_mirrors/nu/numpy创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

NumPy 用户自定义 DType 的 Partition / Argpartition 支持：基于 ArrayMethod API 实现自定义分区算法

相关推荐

EsDA低代码实现Modbus RTU Master转UDP Client协议转换

大文件分片上传实战：Blob切片、秒传与断点续传五层追问

C# HttpClient跳过HTTPS证书验证的三种写法与踩坑指南

更多资讯

ASTM A640标准文件识别与工程应用核验指南

FPGA新手入门：从点亮LED开始掌握Vivado与Verilog开发全流程

SPSS多元线性回归实例：从截图到语法复现与结果解读

OneUptime Runbook 代理（Agent）完全指南：在自建基础设施内安全执行 Bash 与 JavaScript 步骤

Glances MPP 插件详解：监控 Rockchip 平台硬件视频编解码引擎（RKVENC / RKVDEC / RKJPEGD）

Taro H5 端路由系统解析：从 `@tarojs/router` 看小程序路由规范在 Web 端的落地

今日更新

校园网认证计费模式详解：AAA与代拨架构、原理及排障

Spring Boot+Vue垃圾分类毕设：分层、权限与排错实战

把 Cline 的 Base URL 改到 TaoToken 之后，DeepSeek 仍能修调试配置

本周精选

AI驱动的自动化专利分析平台技术解析

LeetCode-Go 题解：1252. Cells with Odd Values in a Matrix（奇数值单元格计数）双解法剖析

在线答疑系统Java毕设实战：Spring Boot前后端分离与状态流转

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？