NEP 55 — 向 NumPy 添加 UTF-8 可变宽度字符串 DType#

作者

Nathan Goldbaum <ngoldbaum@quansight.com>

作者

Warren Weckesser

作者

Marten van Kerkwijk

状态

最终

类型

标准路线

创建时间

2023-06-29

更新时间

2024-01-18

决议

https://mail.python.org/archives/list/numpy-discussion@python.org/thread/Y5CIKBZKMIOWSRYLJ64WV6DKM37QR76B/

摘要#

我们提议向 NumPy 添加一种新的字符串数据类型,其中数组中的每个元素都是任意长度的 UTF-8 编码字符串。这将为 NumPy 用户带来性能、内存使用和可用性方面的改进,包括

  • 对于目前使用固定宽度字符串并在单个 NumPy 数组中主要存储 ASCII 数据或混合短字符串和长字符串的工作流,可节省内存。

  • 下游库和用户将能够摆脱目前用作可变长度字符串数组替代品的 object 数组,通过避免在 NumPy 外部遍历数据,并允许使用快速释放 GIL 的 C 类型转换和字符串 ufuncs 进行字符串操作,从而实现性能改进。

  • 提供更直观的用户级 API,用于处理 Python 字符串数组,而无需考虑内存中的数组表示。

动机和范围#

首先,我们将描述 NumPy 中对字符串或类字符串数据的当前支持状态是如何形成的。接下来,我们将总结上次关于此主题的主要讨论。最后,我们将描述对 NumPy 提出的更改的范围,以及明确超出本提案范围的更改。

NumPy 中字符串支持的历史#

NumPy 中对文本数据的支持是根据早期用户需求以及后来 Python 生态系统的变化而自然演进的。

NumPy 中添加了对字符串的支持,以支持 NumArray chararray 类型的用户。其遗迹在 NumPy API 中仍然可见:字符串相关功能位于 np.char 中,以支持 np.char.chararray 类。此类别尚未正式弃用,但自 NumPy 1.4 以来,模块文档字符串中有一条注释建议改用字符串 dtypes。

NumPy 的 bytes_ DType 最初用于在 NumPy 中添加 Python 3 支持之前表示 Python 2 的 str 类型。bytes DType 在用于表示 Python 2 字符串或其他以 null 结尾的字节序列时最合理。然而,忽略尾随 null 意味着 bytes_ DType 仅适用于不包含尾随 null 的固定宽度字节流,因此对于需要通过 NumPy 字符串进行往返处理的通用字节流,它可能是一个有问题的不匹配项。

C `unicode` DType 的添加是为了支持 Python 2 的 unicode 类型。它以 32 位 UCS-4 码点(例如 UTF-32 编码)存储数据,这使得实现简单直观,但对于使用单字节 ASCII 或 Latin-1 编码可以很好表示的文本来说效率低下。这在 Python 2 中不是问题,因为 ASCII 或大部分 ASCII 文本可以使用 str DType。

随着 NumPy 中对 Python 3 支持的到来,由于向后兼容性问题,字符串 DType 大部分被保留,尽管 unicode DType 成为 str 数据的默认 DType,而旧的 string DType 被重命名为 bytes_ DType。这一变化使得 NumPy 处于一个次优的状态:将最初用于 null 结尾的字节字符串的数据类型作为所有 Python bytes 数据的类型,并且默认字符串类型在内存中的表示方式所消耗的内存是使用单字节 ASCII 或 Latin-1 编码可以很好表示的数据所需内存的四倍。

固定宽度字符串的问题#

现有两种字符串 DType 都表示固定宽度序列,允许将字符串数据存储在数组缓冲区中。这避免了向 NumPy 添加带外存储,但是,这在许多用例中会造成不便的用户界面。特别是,在将数据加载到 NumPy 数组或为字符串操作选择输出 DType 之前,必须由 NumPy 推断或由用户估计最大字符串大小。在最坏的情况下,这需要对整个数据集进行昂贵的遍历以计算数组元素的最大长度。当数组元素长度不同时,它还会浪费内存。数组存储许多短字符串和少数非常长字符串的病理情况尤其会浪费内存。

NumPy 数组中字符串数据的下游使用已证明需要可变宽度字符串数据类型。实际上,许多下游库由于可用性问题而避免使用固定宽度字符串,而是采用 object 数组来存储字符串。特别是,Pandas 已明确弃用对 NumPy 固定宽度字符串的支持,将 NumPy 固定宽度字符串数组强制转换为 object 字符串数组或 PyArrow 支持的字符串数组,并且将来将转换为仅通过 PyArrow 支持字符串数据,后者对 UTF-8 编码的可变宽度字符串数组有原生支持 [1]

此前讨论#

该项目于 2017 年最后一次公开深入讨论此主题,当时 Julian Taylor 提出了一个由编码参数化的固定宽度文本数据类型 [2]。这引发了关于在 NumPy 中处理字符串数据的痛点以及可能的前进方向的广泛讨论。

讨论强调了当前字符串支持处理不佳的两个用例 [3] [4] [5]

  • 加载或内存映射具有未知编码的科学数据集,

  • 以允许 NumPy 数组和 Python 字符串之间透明转换的方式处理“Python 字符串的 NumPy 数组”,包括支持缺失字符串。object DType 部分满足了这一需求,但代价是性能缓慢且没有类型检查。

作为此次讨论的结果,改进对字符串数据的支持被添加到了 NumPy 项目路线图 [6] 中,并明确提出要添加一个更适合内存映射字节(无论有无编码)的 DType,以及一个支持缺失数据的可变宽度字符串 DType,以取代 object 字符串数组的使用。

拟议工作#

本 NEP 提议添加 StringDType,这是一种在 NumPy 数组中存储可变宽度堆分配字符串的 DType,以取代下游使用 object DType 处理字符串数据的情况。这项工作将大量利用 NumPy 近期的改进来提高对用户定义 DType 的支持,因此我们也必然会致力于 NumPy 中的数据类型内部结构。具体来说,我们提议

  • 向 NumPy 添加一种新的可变长度字符串 DType,目标是 NumPy 2.0。

  • 解决将使用实验性 DType API 实现的 DType 添加到 NumPy 本身相关的问题。

  • 支持用户提供的缺失数据哨兵。

  • 在新的 np.strings 命名空间中公开字符串 ufuncs,用于字符串支持相关函数和类型,为未来弃用 np.char 提供了迁移路径。

以下内容超出本次工作的范围

  • 更改字符串数据的 DType 推断。

  • 添加用于内存映射未知编码文本的 DType,或尝试修复 bytes_ DType 问题的 DType。

  • 完全就缺失数据哨兵的语义达成一致,或将缺失数据哨兵添加到 NumPy 本身。

  • 实现字符串操作的 SIMD 优化。

  • 更新 npynpz 文件格式,以允许存储任意长度的伴随数据。

尽管我们明确排除了将这些项目作为本次工作的一部分来实现,但添加新的字符串 DType 有助于为将来实现其中一些项目的工作做好准备。

如果实施,此 NEP 将通过将字符串操作移入长期支持的命名空间并改进 NumPy 内部处理字符串的基础设施,从而在未来更容易添加新的固定宽度文本 DType。我们还提出了一种内存布局,该布局在某些情况下应易于进行 SIMD 优化,从而增加了将来将字符串操作编写为 SIMD 优化 ufuncs 的回报。

尽管我们不提议向 NumPy 添加缺失数据哨兵,但我们提议添加对可选的、用户提供的缺失数据哨兵的支持,因此这确实使 NumPy 更接近于正式支持缺失数据。我们正在努力避免解决 NEP 26 中描述的分歧,并且本提案不要求或排除将来向 ndarray 添加缺失数据哨兵或基于位标志的缺失数据支持。

用法和影响#

该 DType 旨在作为 object 字符串数组的直接替代品。这意味着我们打算支持尽可能多的 object 字符串数组的下游用法,包括所有受支持的 NumPy 功能。Pandas 显然是第一个用户,并且已在 Pandas 的一个分支中进行了大量工作以添加支持。scikit-learn 也使用 object 字符串数组,并且将能够迁移到保证数组只包含字符串的 DType。h5py [7] 和 PyTables [8] 都将能够为 HDF5 中的可变宽度 UTF-8 编码字符串数据集添加一流支持。字符串数据在机器学习工作流中大量使用,下游机器学习库将能够利用这种新的 DType。

希望将字符串数据加载到 NumPy 并利用 NumPy 特性(如高级索引)的用户将有一个自然的选择,与固定宽度 unicode 字符串相比,它能显著节省内存,并提供更好的验证保证以及与 NumPy 的整体集成,优于 object 字符串数组。迁移到一流的字符串 DType 也消除了在字符串操作期间获取 GIL 的需要,从而释放了 object 字符串数组无法实现的未来优化。

性能#

在这里,我们简要描述了我们使用实验性 DType API 在 NumPy 外部实现的 StringDType 原型版本的初步性能测量结果。本节中的所有基准测试均在运行 Ubuntu 22.04 和使用 pyenv 编译的 Python 3.11.3 的戴尔 XPS 13 9380 上进行。NumPy、Pandas 和 StringDType 原型均使用 meson 发布版本编译。

目前,StringDType 原型与 object 数组和固定宽度字符串数组的性能相当。一个例外是从 python 字符串创建数组,其性能比 object 数组稍慢,与固定宽度 unicode 数组相当。

In [1]: from stringdtype import StringDType

In [2]: import numpy as np

In [3]: data = [str(i) * 10 for i in range(100_000)]

In [4]: %timeit arr_object = np.array(data, dtype=object)
3.15 ms ± 74.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [5]: %timeit arr_stringdtype = np.array(data, dtype=StringDType())
8.8 ms ± 12.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [6]: %timeit arr_strdtype = np.array(data, dtype=str)
11.6 ms ± 57.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

在此示例中,object DType 显著更快,因为 data 列表中的对象可以直接在数组中驻留,而 StrDTypeStringDType 需要复制字符串数据,并且 StringDType 需要将数据转换为 UTF-8 并在数组缓冲区之外执行额外的堆分配。未来,如果 Python 转向使用 UTF-8 作为字符串的内部表示,StringDType 的字符串加载性能应该会提高。

字符串操作具有相似的性能

In [7]: %timeit np.array([s.capitalize() for s in data], dtype=object)
31.6 ms ± 728 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [8]: %timeit np.char.capitalize(arr_stringdtype)
41.5 ms ± 84.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [9]: %timeit np.char.capitalize(arr_strdtype)
47.6 ms ± 386 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

此处的性能不佳反映了 np.char 中操作的缓慢的基于迭代器的实现。当我们完成将这些操作重写为 ufuncs 时,我们将实现显著的性能改进。以 add ufunc 为例,我们已为 StringDType 原型实现了它。

In [10]: %timeit arr_object + arr_object
10.1 ms ± 400 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [11]: %timeit arr_stringdtype + arr_stringdtype
3.64 ms ± 258 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [12]: %timeit np.char.add(arr_strdtype, arr_strdtype)
17.7 ms ± 245 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

如下所述,我们已更新了 Pandas 的一个分支,以使用 StringDType 的原型版本。这表明当数据已加载到 NumPy 数组并传递给第三方库时,可获得的性能改进。目前 Pandas 默认尝试将所有 str 数据强制转换为 object DType,并且必须检查和清理传入的现有 object 数组。这需要复制或遍历数据,而 NumPy 和 Pandas 中对可变宽度字符串的一流支持使得这些操作变得不必要。

In [13]: import pandas as pd

In [14]: %timeit pd.Series(arr_stringdtype)
18.8 µs ± 164 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

如果我们强制 Pandas 使用 object 字符串数组(直到最近它一直是默认设置),我们将看到在 NumPy 外部遍历数据的巨大性能开销。

In [15]: %timeit pd.Series(arr_object, dtype='string[python]')
907 µs ± 67 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each

Pandas 默认切换到 PyArrow 支持的字符串数组,就是为了避免这种以及与 object 字符串数组相关的其他性能成本。

向后兼容性#

我们不提议更改 Python 字符串的 DType 推断,也不期望对 NumPy 的现有用法产生任何影响。

详细描述#

在这里,我们提供了我们希望包含在 NumPy 中的 StringDType 版本的详细描述。这与原型基本相同,但有一些差异是 NumPy 外部的 DType 无法实现的。

首先,我们将描述用于实例化 StringDType 实例的 Python API。接下来,我们将描述缺失数据处理支持以及对数组元素进行严格字符串类型检查的支持。然后,我们将讨论我们将定义的类型转换 (cast) 和 ufunc 实现,并讨论我们计划在新的 np.strings 命名空间中直接公开字符串 ufuncs 的方案。最后,我们将概述我们希望公开的 C API 以及为初始实现选择的内存布局和堆分配策略的详细信息。

StringDType 的 Python API#

新的 DType 将通过 np.dtypes 命名空间访问

>>> from numpy.dtypes import StringDType
>>> dt = StringDType()
>>> dt
numpy.dtypes.StringDType()

此外,我们建议将字符 "T"(text 的缩写)保留用于 np.dtype,因此上述用法将与以下用法相同

>>> np.dtype("T")
numpy.dtypes.StringDType()

StringDType 可以直接用于在 NumPy 数组中表示任意长度的字符串

>>> data = ["this is a very long string", "short string"]
>>> arr = np.array(data, dtype=StringDType())
>>> arr
array(['this is a very long string', 'short string'], dtype=StringDType())

请注意,与固定宽度字符串不同,StringDType 不受数组元素最大长度的限制,任意长或短的字符串可以存在于同一个数组中,而无需为短字符串中的填充字节保留存储空间。

StringDType 类作为 NumPy Python API 中的 dtype 参数传递时,它将是默认 StringDType 实例的同义词。我们已经将大部分 API 表面转换为这种工作方式,但仍有一些地方尚未转换,第三方代码也可能尚未转换,因此我们不会在文档中强调这一点。强调 StringDType 是一个类而 StringDType() 是一个实例是一种更具前瞻性的 API,NumPy DType API 的其余部分现在可以向此方向发展,因为 DType 类可以从 np.dtypes 命名空间导入,因此即使并非严格必要,我们仍将在文档中包含 StringDType 对象的显式实例化。

我们建议将 Python 内置类型 str 关联为该 DType 的标量类型

>>> StringDType.type
<class 'str'>

尽管这确实在 API 中造成了一个小问题,即 NumPy 中内置 DType 类到标量的映射将不再是一一对应(unicode DType 的标量类型是 str),但这避免了为此目的定义、优化或维护 str 子类或使用其他技巧来维持这种一一对应关系。为了保持向后兼容性,对 Python 字符串列表检测到的 DType 将仍然是固定宽度的 unicode 字符串。

如下所述,StringDType 支持两个参数,可以调整 DType 的运行时行为。我们不会尝试通过字符代码支持 dtype 的参数。如果用户需要一个不使用默认参数的 DType 实例,他们将需要使用 DType 类实例化一个 DType 实例。

我们还将扩展 C API 中的 NPY_TYPES 枚举,添加一个 NPY_VSTRING 条目(已有 NPY_STRING 条目)。这不应干扰遗留的用户定义 DType,因为这些数据类型的整数类型编号从 256 开始。原则上,NPY_TYPES 枚举中可用的整数范围内仍有数百个内置 DType 的空间。

原则上我们不需要保留字符代码,并且有摆脱字符代码的愿望。然而,大量下游代码依赖于检查 DType 字符代码来区分内置 NumPy DType,我们认为如果要求用户在使用 StringDType 时重构其 DType 处理代码,将会损害其采用。

我们还希望将来能够添加一个新版本的固定宽度文本 StringDType,它可以使用带长度或编码修饰符的 "T" 字符代码。这将允许迁移到更灵活的文本 dtype,用于结构化数组和其他固定宽度字符串比可变宽度字符串更适合的用例。

缺失数据支持#

缺失数据可以使用哨兵表示

>>> dt = StringDType(na_object=np.nan)
>>> arr = np.array(["hello", nan, "world"], dtype=dt)
>>> arr
array(['hello', nan, 'world'], dtype=StringDType(na_object=nan))
>>> arr[1]
nan
>>> np.isnan(arr[1])
True
>>> np.isnan(arr)
array([False,  True, False])
>>> np.empty(3, dtype=dt)
array(['', '', ''])

我们只提议支持用户提供的哨兵。默认情况下,空数组将填充空字符串。

>>> np.empty(3, dtype=StringDType())
array(['', '', ''], dtype=StringDType())

通过仅支持用户提供的缺失数据哨兵,我们避免了精确解决 NumPy 本身应如何支持缺失数据以及缺失数据对象的正确语义的问题,将决定权留给用户。但是,我们确实检测用户是否提供了一个类似 NaN 的缺失数据值、一个字符串缺失数据值,或者两者都不是。我们将在下面解释如何处理这些情况。

谨慎的读者可能会担心需要处理三种不同类别的缺失数据哨兵的复杂性。这里的复杂性反映了 object 数组的灵活性以及我们发现的下游使用模式。有些用户希望与哨兵的比较产生错误,因此他们使用 None。另一些用户希望比较成功并具有某种有意义的排序,因此他们使用任意的、希望是唯一的字符串。还有一些用户希望使用在比较和算术中表现得像 NaN 或字面上就是 NaN 的东西,这样专门查找精确 NaN 的 NumPy 操作就能正常工作,并且无需在 NumPy 之外重写缺失数据处理。我们相信支持所有这些是可能的,但这需要一些有望可管理的复杂性。

类似 NaN 的哨兵#

类似 NaN 的哨兵在算术运算中返回自身作为结果。这包括 Python 的 nan 浮点数和 Pandas 的缺失数据哨兵 pd.NA。我们选择让类似 NaN 的哨兵在操作中继承这些行为,因此加法的结果就是哨兵。

>>> dt = StringDType(na_object=np.nan)
>>> arr = np.array(["hello", np.nan, "world"], dtype=dt)
>>> arr + arr
array(['hellohello', nan, 'worldworld'], dtype=StringDType(na_object=nan))

我们还选择让类似 NaN 的哨兵排序到数组的末尾,遵循包含 nan 的数组的排序行为。

>>> np.sort(arr)
array(['hello', 'world', nan], dtype=StringDType(na_object=nan))

字符串哨兵#

字符串缺失数据值是 str 的实例或 str 的子类型。

操作将直接使用哨兵值处理缺失条目。这是我们在下游代码中发现此模式的主要用法,例如将像 "__nan__" 这样的缺失数据哨兵传递给低级排序或分区算法。

其他哨兵#

任何其他 Python 对象将在操作或比较中引发错误,就像当前 None 作为 object 数组的缺失数据哨兵时那样。

>>> dt = StringDType(na_object=None)
>>> np.sort(np.array(["hello", None, "world"], dtype=dt))
ValueError: Cannot compare null that is not a string or NaN-like value

由于比较需要引发错误,并且 NumPy 比较 API 在不持有 GIL 的情况下无法在排序期间发出基于值的错误信号,因此使用任意缺失数据哨兵的排序数组将持有 GIL。我们也可以尝试通过重构 NumPy 的比较和排序实现来放宽此限制,以允许在排序操作期间传播基于值的错误。

对 DType 推断的影响#

如果将来我们决定打破向后兼容性,将 StringDType 设置为 str 数据的默认 DType,那么支持任意对象作为缺失数据哨兵可能似乎对实现 DType 推断造成问题。然而,鉴于该 DType 的初始支持将需要直接使用该 DType,并且无法依赖 NumPy 来推断 DType,我们认为这对于缺失数据功能的下游用户来说不会是主要问题。要使用 StringDType,他们需要更新代码以在创建数组时显式指定 DType,因此如果 NumPy 将来更改 DType 推断,他们的代码行为将不会改变,并且缺失数据哨兵将永远不需要参与 DType 推断。

强制转换非字符串#

默认情况下,非字符串数据会被强制转换为字符串

>>> np.array([1, object(), 3.4], dtype=StringDType())
array(['1', '<object object at 0x7faa2497dde0>', '3.4'], dtype=StringDType())

如果不希望出现此行为,可以创建一个禁用字符串强制转换的 DType 实例。

>>> np.array([1, object(), 3.4], dtype=StringDType(coerce=False))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: StringDType only allows string data when string coercion
is disabled

这允许在 NumPy 用于创建数组的相同数据遍历中进行严格数据验证,而下游库无需在单独的、昂贵的输入类数组遍历中实现自己的字符串验证。我们选择不将此作为默认行为,以遵循 NumPy 固定宽度字符串(它会强制转换非字符串)。

类型转换、ufunc 支持和字符串操作函数#

将提供一套完整的往返内置 NumPy DType 的类型转换。此外,我们将添加比较运算符的实现,以及一个接受两个字符串数组的 add 循环,接受字符串和整数数组的 multiply 循环,一个 isnan 循环,以及 NumPy 2.0 中新增的 str_lenisalphaisdecimalisdigitisnumericisspacefindrfindcountstriplstriprstripreplace 字符串 ufuncs 的实现。

isnan ufunc 将对类似 NaN 的哨兵条目返回 True,否则返回 False。比较将按照 unicode 码点顺序对数据进行排序,这与当前固定宽度 unicode DType 的实现方式相同。将来 NumPy 或下游库可能会为 NumPy unicode 字符串数组添加区域设置感知排序、大小写折叠和规范化,但我们目前不提议添加这些功能。

两个 StringDType 实例在用相同的 na_objectcoerce 参数创建时被认为是相等的。对于接受多个字符串参数的 ufuncs,我们还引入了“兼容” StringDType 实例的概念。我们允许在 ufunc 操作中同时使用不同的 DType 实例,如果它们具有相同的 na_object,或者如果只有一个 DType 显式设置了 na_object。我们不考虑字符串强制转换来确定实例是否兼容,尽管如果操作结果是字符串,结果将继承原始操作数中更严格的字符串强制转换设置。

这种“兼容”实例的概念将在二进制 ufuncs 的 resolve_descriptors 函数中强制执行。这一选择使得使用非默认 StringDType 实例变得更容易,因为 Python 字符串被强制转换为默认的 StringDType 实例,因此以下惯用表达式是被允许的。

>>> arr = np.array(["hello", "world"], dtype=StringDType(na_object=None))
>>> arr + "!"
array(['hello!', 'world!'], dtype=StringDType(na_object=None))

如果我们只考虑 StringDType 实例的相等性,这将是一个错误,导致糟糕的用户体验。如果操作数具有不同的 na_object 设置,NumPy 将引发错误,因为结果 DType 的选择不明确。

>>> arr + np.array("!", dtype=StringDType(na_object=""))
TypeError: Cannot find common instance for incompatible dtype instances

np.strings 命名空间#

字符串操作将在 np.strings 命名空间中提供,该命名空间将填充字符串 ufuncs。

>>> np.strings.upper((np.array(["hello", "world"], dtype=StringDType())
array(['HELLO', 'WORLD'], dtype=StringDType())
>>> isinstance(np.strings.upper, np.ufunc)
True

我们认为 np.stringsnp.char 更直观,并且最终将取代 np.char,一旦下游库根据 SPEC-0 支持的最低 NumPy 版本足够新,它们就可以安全地切换到 np.strings,而无需任何基于 NumPy 版本的条件逻辑。

序列化#

由于字符串数据存储在数组缓冲区之外,因此序列化为 npy 格式需要格式修订以支持存储可变宽度伴随数据。作为这项工作的一部分,我们不打算支持序列化为 npynpz 格式,除非指定 allow_pickle=True

这延续了当前 object 字符串数组的情况,它们只能使用 allow_pickle=True 选项保存到 npy 文件。

将来我们可能会决定添加对此的支持,但应注意不要破坏 NumPy 外部可能未维护的解析器。

StringDType 的 C API#

C API 的目标是向用户隐藏字符串数据在堆上的存储细节,并提供一个线程安全的接口,用于读取和写入存储在 StringDType 数组中的字符串。为了实现这一点,我们决定将字符串分为两种不同的表示形式:打包 (packed)解包 (unpacked)。打包字符串直接存在于数组缓冲区中,可能包含足够短的字符串数据,或包含堆分配的元数据,其中存储了字符串的字符。解包字符串公开字符串的字节大小和一个指向字符串数据的 char * 指针。

要访问存储在 NumPy 数组中的字符串的解包字符串数据,用户必须调用一个函数将打包字符串加载到解包字符串中,或调用另一个函数将解包字符串打包到数组中。这些操作需要一个指向数组条目的指针和一个分配器结构体的引用。分配器管理在堆上存储字符串数据所需的簿记。将这种簿记集中在分配器中意味着我们可以自由更改底层分配策略。我们还通过使用互斥锁 (mutex) 保护对分配器的访问来确保线程安全。

下面我们将更详细地描述此设计,列举我们希望添加到 C API 中的类型和函数。在下一节中,我们将描述计划使用此 API 实现的内存布局和堆分配策略。

PyArray_StringDTypePyArray_StringDTypeObject 结构体#

我们将公开 StringDType 元类的结构体以及 StringDType 实例类型的结构体。前者 PyArray_StringDType 将在 C API 中提供,与其他 PyArray_DTypeMeta 实例用于编写 ufunc 和类型转换循环的方式相同。此外,我们将公开以下结构体:

struct PyArray_StringDTypeObject {
    PyArray_Descr base;
    // The object representing a null value
    PyObject *na_object;
    // Flag indicating whether or not to coerce arbitrary objects to strings
    char coerce;
    // Flag indicating the na object is NaN-like
    char has_nan_na;
    // Flag indicating the na object is a string
    char has_string_na;
    // If nonzero, indicates that this instance is owned by an array already
    char array_owned;
    // The string data to use when a default string is needed
    npy_static_string default_string;
    // The name of the missing data object, if any
    npy_static_string na_name;
    // the allocator should only be directly accessed after
    // acquiring the allocator_lock and the lock should
    // be released immediately after the allocator is
    // no longer needed
    npy_string_allocator *allocator;
}

公开此定义有助于未来与其他 dtypes 的集成。

字符串和分配器类型#

解包字符串在 C API 中用 npy_static_string 类型表示,其定义将公开如下:

struct npy_static_string {
    size_t size;
    const char *buf;
};

其中 size 是字符串的字节大小,buf 是一个指向 UTF-8 编码字节流开头(包含字符串数据)的常量指针。这是一个只读的字符串视图,我们不会暴露用于修改这些字符串的公共接口。我们不会在字节流末尾追加一个尾随空字符,因此试图将 buf 字段传递给期望 C 字符串的 API 的用户必须创建一个带尾随空的副本。将来,如果复制到以 null 结尾的缓冲区对于 C API 的下游用户来说成本过高,我们可能会决定始终写入尾随空字节。

此外,我们将公开两个不透明结构体:npy_packed_static_stringnpy_string_allocatorStringDType NumPy 数组中的每个条目都将存储一个 npy_packed_static_string 的内容;这是一个字符串的打包表示。字符串数据要么直接存储在打包字符串中,要么存储在堆上,由附加到与数组关联的独立 npy_string_allocator 结构体管理。打包字符串的精确布局和用于在堆上分配数据的策略将不会公开,用户不应依赖这些细节。

新的 C API 函数#

我们计划公开的 C API 函数分为两类:获取和释放分配器锁的函数,以及加载和打包字符串的函数。

获取和释放分配器#

获取和释放分配器的主要接口是以下一对静态内联函数

static inline npy_string_allocator *
NpyString_acquire_allocator(PyArray_StringDTypeObject *descr)

static inline void
NpyString_release_allocator(npy_string_allocator *allocator)

第一个函数获取附加到描述符实例的分配器锁,并返回指向与该描述符关联的分配器的指针。然后该线程可以使用分配器加载现有打包字符串或将新字符串打包到数组中。一旦需要分配器的操作完成,必须释放分配器锁。在调用 NpyString_release_allocator 之后使用分配器可能导致数据竞争或内存损坏。

还有一些情况,同时处理多个分配器会很方便。例如,add ufunc 接受两个字符串数组并生成第三个字符串数组。这意味着 ufunc 循环需要三个分配器才能加载每个操作数的字符串并将结果打包到输出数组中。更复杂的是,输入和输出操作数不一定是不同的对象,并且由于是同一个数组,操作数可以共享分配器。原则上,我们可以要求用户在 ufunc 循环内部获取和释放锁,但这会比在循环设置中获取所有三个分配器并在循环结束时同时释放它们增加大量的性能开销。

为了处理这些情况,我们还将公开这两个函数的变体,它们接受任意数量的描述符和分配器(NpyString_acquire_allocatorsNpyString_release_allocators)。公开这些函数使得编写同时处理多个分配器的代码变得简单。简单地多次调用 NpyString_acquire_allocatorNpyString_release_allocator 的朴素方法,当 ufunc 操作数共享描述符时,将尝试在同一线程中多次获取相同的锁,从而导致未定义行为。多描述符变体在尝试获取锁之前检查描述符是否相同,从而避免了未定义行为。为了正确操作,用户只需选择接受与所需工作描述符数量相同的分配器获取或释放变体即可。

打包和加载字符串#

访问字符串由以下函数中介

int NpyString_load(
    npy_string_allocator *allocator,
    const npy_packed_static_string *packed_string,
    npy_static_string *unpacked_string)

此函数在错误时返回 -1,这可能发生在存在线程错误或损坏阻止访问堆分配的情况下。成功时它可能返回 1 或 0。如果返回 1,则表示打包字符串的内容是空字符串 (null string),在这种情况下可以执行处理空字符串的特殊逻辑。如果函数返回 0,则表示可以从 unpacked_string 中读取 packed_string 的内容。

字符串打包可以通过以下函数之一进行

int NpyString_pack(
    npy_string_allocator *allocator,
    npy_packed_static_string *packed_string,
    const char *buf, size_t size)

int NpyString_pack_null(
    npy_string_allocator *allocator,
    npy_packed_static_string *packed_string)

第一个函数将 buf 的前 size 个元素内容打包到 packed_string 中。第二个函数将空字符串打包到 packed_string 中。这两个函数都会使与打包字符串关联的任何先前堆分配失效,并且在打包字符串后,仍在作用域中的旧解包表示也将失效。这两个函数在成功时返回 0,失败时返回 -1,例如当 malloc 失败时。

C API 使用示例#

加载字符串#

假设我们正在为 StringDType 编写 ufunc 实现。如果给定 const char *buf 指向 StringDType 数组条目的开头,以及 PyArray_Descr * 指向数组描述符,则可以如下访问底层字符串数据:

npy_string_allocator *allocator = NpyString_acquire_allocator(
        (PyArray_StringDTypeObject *)descr);

npy_static_string sdata = {0, NULL};
npy_packed_static_string *packed_string = (npy_packed_static_string *)buf;
int is_null = 0;

is_null = NpyString_load(allocator, packed_string, &sdata);

if (is_null == -1) {
    // failed to load string, set error
    return -1;
}
else if (is_null) {
    // handle missing string
    // sdata->buf is NULL
    // sdata->size is 0
}
else {
    // sdata->buf is a pointer to the beginning of a string
    // sdata->size is the size of the string
}
NpyString_release_allocator(allocator);
打包字符串#

此示例展示了如何将新字符串打包到数组中

char *str = "Hello world";
size_t size = 11;
npy_packed_static_string *packed_string = (npy_packed_static_string *)buf;

npy_string_allocator *allocator = NpyString_acquire_allocator(
        (PyArray_StringDTypeObject *)descr);

// copy contents of str into packed_string
if (NpyString_pack(allocator, packed_string, str, size) == -1) {
    // string packing failed, set error
    return -1;
}

// packed_string contains a copy of "Hello world"

NpyString_release_allocator(allocator);

Cython 支持和缓冲区协议#

StringDType 不可能支持 Python 缓冲区协议,因此除非 Cython 将来添加特殊支持,否则 Cython 不会支持 StringDType 数组的惯用类型化 memoryview 语法。我们有一些初步的想法,可以更新缓冲区协议 [9] 或利用 Arrow C 数据接口 [10] 来公开对于缓冲区协议无意义的 DType 的 NumPy 数组,但这些努力可能无法及时在 NumPy 2.0 之前实现。这意味着将使用固定宽度字符串数组的传统 Cython 代码调整为与 StringDType 协同工作将是相当困难的。调整与 object 字符串数组协同工作的代码应该会很简单,因为 object 数组也不受缓冲区协议支持,并且在 Cython 中可能没有类型或具有 object 类型。

我们将为作为这项工作一部分而添加的公共 C API 函数添加 Cython nogil 包装器,以方便与下游 Cython 代码集成。

内存布局和堆分配管理#

下面我们详细描述了我们选择的内存布局,但在深入探讨之前,我们想指出,上述 C API 没有公开这些细节中的任何一个。以下所有内容都可能在未来进行修订、改进和更改,因为字符串数据的精确内存布局并未公开。

内存布局和短字符串优化#

每个数组元素都表示为一个联合体,在小端架构上的定义如下:

typedef struct _npy_static_vstring_t {
   size_t offset;
   size_t size_and_flags;
} _npy_static_string_t;

typedef struct _short_string_buffer {
   char buf[sizeof(_npy_static_string_t) - 1];
   unsigned char size_and_flags;
} _short_string_buffer;

typedef union _npy_static_string_u {
 _npy_static_string_t vstring;
 _short_string_buffer direct_buffer;
} _npy_static_string_u;

_npy_static_vstring_t 表示最适用于直接在堆上或在竞技场 (arena) 分配中表示字符串,其中 offset 字段直接包含地址的 size_t 表示,或者是一个竞技场分配中的整数偏移量。_short_string_buffer 表示最适用于短字符串优化,其中字符串数据存储在 direct_buffer 字段中,大小存储在 size_and_flags 字段中。在这两种情况下,size_and_flags 字段都存储字符串的 size 以及位标志。短字符串将大小存储在缓冲区的最后四位中,将 size_and_flags 的前四位保留给标志。堆字符串或竞技场分配中的字符串使用最高有效字节作为标志,将前导字节保留给字符串大小。值得指出的是,这个选择限制了数组中允许存储的最大字符串大小,特别是在 32 位系统上,每个字符串的限制是 16 兆字节——小到足以令人担心影响实际工作流程。

在大端系统上,布局是反向的,size_and_flags 字段首先出现在结构体中。这允许实现总是使用 size_and_flags 字段的最高有效位作为标志。这些结构体的依赖于字节序的布局是实现细节,不会在 API 中公开。

字符串是直接存储在竞技场缓冲区还是在堆中,通过设置字符串数据上的 NPY_OUTSIDE_ARENANPY_STRING_LONG 标志来指示。由于堆分配字符串的最大大小受限于最大的 7 字节无大小整数的大小,因此这些标志永远不会为有效的堆字符串设置。

有关每种内存布局中字符串的一些直观示例,请参阅内存布局示例

竞技场分配器#

在 64 位系统上超过 15 字节,在 32 位系统上超过 7 字节的字符串存储在数组缓冲区之外的堆上。分配的簿记由附加到与数组关联的 StringDType 实例的竞技场分配器管理。该分配器将作为不透明的 npy_string_allocator 结构体公开。其内部布局如下:

struct npy_string_allocator {
    npy_string_malloc_func malloc;
    npy_string_free_func free;
    npy_string_realloc_func realloc;
    npy_string_arena arena;
    PyThread_type_lock *allocator_lock;
};

这允许我们将内存分配函数组合在一起,并根据需要选择不同的运行时分配函数。分配器的使用受到互斥锁 (mutex) 的保护,有关线程安全的更多讨论请参阅下文。

内存分配由 npy_string_arena 结构体成员处理,其布局如下:

struct npy_string_arena {
    size_t cursor;
    size_t size;
    char *buffer;
};

其中 buffer 是指向堆分配竞技场开头的指针,size 是该分配的大小,cursor 是竞技场中上一次竞技场分配结束的位置。竞技场使用指数级扩展缓冲区填充,扩展因子为 1.25。

竞技场中的每个字符串条目都以大小为前缀,大小存储在 charsize_t 中,具体取决于字符串的长度。长度在 16 或 8(取决于架构)到 255 之间的字符串以 char 大小存储。我们内部称这些为“中等”字符串。这种选择通过每条中等长度字符串减少 7 字节的开销,从而降低了在堆上存储较小字符串的开销。竞技场中长度超过 255 字节的字符串会设置 NPY_STRING_LONG 标志。

如果打包字符串的内容被释放,然后分配给一个与最初存储在打包字符串中的字符串大小相同或更小的新字符串,则现有的短字符串或竞技场分配将被重复使用。然而,有一个例外,当竞技场中的字符串被短字符串覆盖时,竞技场元数据会丢失,竞技场分配无法重复使用。

如果字符串被放大,竞技场缓冲区中的现有空间将无法使用,因此我们转而通过 malloc 直接在堆上分配空间,并设置 NPY_STRING_OUTSIDE_ARENANPY_STRING_LONG 标志。请注意,在这种情况下,即使字符串长度小于 255 字节,也可以设置 NPY_STRING_LONG。由于堆地址会覆盖竞技场偏移量,并且未来的字符串替换将存储在堆上或直接作为短字符串存储在数组缓冲区中。

无论存储在哪里,一旦字符串被初始化,它都会被标记上 NPY_STRING_INITIALIZED 标志。这使我们能够清楚地区分未初始化的空字符串和已变为空字符串的字符串。

分配的大小存储在竞技场中,以便在字符串变异时可以重用竞技场分配。原则上,我们可以禁止重用竞技场缓冲区,并且不将大小存储在竞技场中。这可能节省内存或提高性能,也可能不会,具体取决于确切的使用模式。目前,我们倾向于在字符串变异时避免不必要的堆分配,但原则上我们可以通过选择始终将变异的竞技场字符串存储为堆字符串并忽略竞技场分配来简化实现。有关在多线程上下文中如何处理 NumPy 数组可变性的更多详细信息,请参阅下文。

使用每数组竞技场分配器可确保相邻数组元素的字符串缓冲区通常在堆上彼此靠近。我们不保证相邻数组元素在堆上是连续的,以支持短字符串优化、缺失数据和允许数组条目变异。有关这些主题如何影响内存布局的更多讨论,请参阅下文。

变异和线程安全#

当多个线程访问和修改数组时,变异会引入数据竞争和使用后释放错误的可能。此外,如果我们在竞技场缓冲区中分配变异的字符串,并要求旧字符串被新字符串替换时必须是连续存储,那么修改单个字符串可能会触发对整个数组的竞技场缓冲区重新分配。与 object 字符串数组或固定宽度字符串相比,这是一种病态的性能下降。

一个解决方案是禁用变异,但不可避免地会有下游使用 object 字符串数组来变异数组元素的情况,而我们希望支持这些情况。

相反,我们选择将附加到 PyArray_StringDType 实例的 npy_string_allocator 实例与 PyThread_type_lock 互斥锁配对。静态字符串 C API 中任何允许操作堆分配数据的函数都接受一个 allocator 参数。要正确使用 C API,线程必须在任何使用 allocator 之前获取分配器互斥锁。

PyThread_type_lock 互斥锁相对较重,并且不提供更复杂的允许同时多个读者的锁原语。作为 GIL 移除项目的一部分,CPython 正在向 C API 添加新的同步原语,供 NumPy 等项目使用。当这种情况发生时,我们可以更新锁定策略以允许同时多个读线程,以及 NumPy 中一旦 GIL 移除后所需的其他线程错误修复。

释放字符串#

在丢弃或重用打包字符串之前,必须释放现有字符串。API 的构建要求所有字符串都如此,即使是没有堆分配的短字符串也如此。在所有情况下,打包字符串中的所有数据都将被清零,除了标志(它们会被保留)。

内存布局示例#

我们创建了三种可能的字符串内存布局的示意图。所有图表都假设为 64 位小端架构。

_images/nep-0055-short-string-memory-layout.svg

短字符串将字符串数据直接存储在数组缓冲区中。在小端架构上,字符串数据首先出现,接着是一个字节,该字节为四个标志留出空间,并在最后 4 位中以无符号整数形式存储字符串大小。在此示例中,字符串内容为“Hello world”,大小为 11。标志表示此字符串存储在竞技场外部且已初始化。

_images/nep-0055-arena-string-memory-layout.svg

竞技场字符串将字符串数据存储在由附加到数组的 StringDType 实例管理的堆分配竞技场缓冲区中。在此示例中,字符串内容为“Numpy is a very cool library”,存储在竞技场分配的偏移量 0x94C 处。请注意,size 存储了两次,一次在 size_and_flags 字段中,另一次在竞技场分配中。这有助于在字符串被修改时重用竞技场分配。另请注意,由于字符串长度足够小,可以放入 unsigned char 中,因此这是一个“中等”长度字符串,大小在竞技场分配中只需要一个字节。大于 255 字节的竞技场字符串将需要在竞技场中使用 8 字节来存储 size_t 大小。唯一设置的标志表示此字符串已初始化。

_images/nep-0055-heap-string-memory-layout.svg

堆字符串将字符串数据存储在由 PyMem_RawMalloc 返回的缓冲区中,并且不存储竞技场缓冲区中的偏移量,而是直接存储 malloc 返回的堆地址。在此示例中,字符串内容为“Numpy is a very cool library”,并存储在堆地址 0x4d3d3d3。该字符串设置了三个标志,表示它是一个“长”字符串(例如,不是短字符串),存储在竞技场外部,并且已初始化。请注意,如果此字符串存储在竞技场内部,它将不会设置长字符串标志,因为它存储所需的字节数少于 256。

空字符串和缺失数据#

我们选择的布局有一个好处,即由 calloc 返回的新创建的数组缓冲区在构造时将是一个填充了空字符串的数组,因为没有设置任何标志的字符串是一个未初始化的零长度竞技场字符串。这并不是空字符串唯一有效的表示形式,因为可以设置其他标志来指示空字符串是否与预先存在的短字符串或竞技场字符串相关联。

缺失字符串将具有相同的表示形式,只是它们的 flags 字段中总是会设置一个标志,即 NPY_STRING_MISSING。用户在访问解包字符串缓冲区之前需要检查字符串是否为 null,并且我们已将 C API 设置为在解包字符串时强制执行 null 检查。缺失字符串和空字符串都可以根据打包字符串表示中的数据进行检测,并且不需要在竞技场分配中或额外的堆分配中留出相应的空间。

实现#

我们有一个开放的拉取请求[12],已准备好合并到NumPy中,用于添加StringDType。

我们已经创建了一个支持使用StringDType创建Pandas数据结构的Pandas开发分支[13]。这展示了在大量使用对象字符串数组的下游库中支持StringDType所需的重构工作。

如果获得批准,本NEP剩余的大部分工作将是更新文档和完善NumPy 2.0版本。我们已经完成了以下工作:

  • 创建一个np.strings命名空间,并直接在该命名空间中公开字符串ufuncs。

  • StringDType实现从外部扩展模块移至NumPy内部,并在适当的地方重构NumPy。这个新的DType将通过一个包含文档更新的大型拉取请求添加。在可能的情况下,我们将在发布主拉取请求之前,将与StringDType无关的修复和重构提取到较小的拉取请求中。

我们将继续进行以下工作:

  • 处理NumPy中与新DType相关的剩余问题。特别是,我们已经意识到,NumPycopyswap的剩余用法应迁移到使用类型转换(cast)或尚未添加的单元素复制DType API槽。我们还需要确保DType类可以在Python API中所有有意义的地方与DType实例互换使用,并在所有其他可以传入DType实例但DType类没有意义的地方添加有用的错误信息。

替代方案#

主要的替代方案是维持现状,并提供对象数组作为变长字符串数组的解决方案。虽然这可行,但这意味着内存使用和性能的即时改进以及未来的性能改进将不会很快实现,NumPy将失去与对文本数据数组支持更好的其他生态系统的相关性。

我们认为提议的DType与改进的固定宽度二进制字符串DType(可以表示任意二进制数据或任何编码的文本)并非互斥,并且在添加StringDType之后,一旦NumPy中对字符串数据的整体支持得到改善,未来添加此类DType将更加容易。

讨论#

参考文献和脚注#