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/[email protected]/thread/Y5CIKBZKMIOWSRYLJ64WV6DKM37QR76B/

摘要#

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

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

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

  • 一个更直观的、面向用户的 API,用于处理 Python 字符串数组,无需考虑内存中的数组表示。

动机和范围#

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

NumPy 中字符串支持的历史#

NumPy 中对文本数据的支持是根据早期用户的需求,然后是 Python 生态系统的变化而逐渐发展的。

对字符串的支持已添加到 NumPy 中,以支持 NumArray chararray 类型的用户。这些遗留痕迹在 NumPy API 中仍然可见:字符串相关功能位于 np.char 中,以支持 np.char.chararray 类。该类没有正式弃用,但在模块文档字符串中有一条注释,建议自 NumPy 1.4 起使用字符串 DType。

NumPy 的 bytes_ DType 最初用于在 Python 3 支持添加到 NumPy 之前表示 Python 2 str 类型。当 bytes_ DType 用于表示 Python 2 字符串或其他以空值结尾的字节序列时,它最有意义。但是,忽略尾随空值意味着 bytes_ DType 仅适用于不包含尾随空值的固定宽度字节流,因此它可能不适合用于通用的字节流,在这些字节流中需要将尾随空值往返于 NumPy 字符串。

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

随着 Python 3 支持在 NumPy 中的到来,由于向后兼容性的考虑,字符串 DType 大体上保持不变,尽管 unicode DType 成为 str 数据的默认 DType,并且旧的 string DType 被重命名为 bytes_ DType。这一变化使 NumPy 处于一个次优状态,即发布了一个最初用于以空值结尾的字节字符串的数据类型,作为所有 python bytes 数据的数据类型,以及一个默认字符串类型,其内存表示所消耗的内存是使用单字节 ASCII 或 Latin-1 编码可以很好地表示的数据所需的内存的四倍。

固定长度字符串的问题#

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

在 NumPy 数组中下游使用字符串数据证明了对可变长度字符串数据类型的需求。实际上,许多下游库为了避免使用固定长度字符串而带来的可用性问题,而是使用 object 数组来存储字符串。特别是,Pandas 已明确弃用对 NumPy 固定长度字符串的支持,它将 NumPy 固定长度字符串数组强制转换为 object 字符串数组或 PyArrow 支持的字符串数组,并且在将来将只支持通过 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 字符串数组的下游使用。这项工作将大量利用 NumPy 中的最新改进,以改进对用户定义 DType 的支持,因此我们也必然会研究 NumPy 中的数据类型内部。具体来说,我们建议

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

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

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

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

以下内容不在本次工作范围内

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

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

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

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

  • 更新 npynpz 文件格式以允许存储任意长度的侧边车数据。

虽然我们明确排除了在本次工作中实现这些项目,但添加新的字符串 DType 有助于为将来实现这些项目中的部分内容奠定基础。

如果实施本 NEP,将使将来更容易添加新的固定长度文本 DType,方法是将字符串操作移动到长期支持的命名空间,并改进 NumPy 中处理字符串的内部基础设施。我们还建议使用一种内存布局,该布局应该在某些情况下适合 SIMD 优化,从而增加了将来将字符串操作编写为 SIMD 优化的 ufunc 的回报。

虽然我们没有建议向 NumPy 添加缺失数据哨兵,但我们建议添加对可选的、用户提供的缺失数据哨兵的支持,因此这确实使 NumPy 离正式支持缺失数据更近了一步。我们试图避免解决 NEP 26 中描述的分歧,并且本提案不要求或排除非将来向 ndarray 添加缺失数据哨兵或基于位标志的缺失数据支持。

使用和影响#

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

希望将字符串数据加载到 NumPy 并利用 NumPy 特性(如高级索引)的用户,可以选择一种方法,它在内存节省方面明显优于固定宽度 Unicode 字符串,并且比对象字符串数组提供了更好的验证保证和与 NumPy 的整体集成。迁移到一流的字符串 DType 还消除了在字符串操作期间获取 GIL 的需求,从而解锁了对象字符串数组无法实现的未来优化。

性能#

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

目前,StringDType 原型的性能与对象数组和固定宽度字符串数组相当。唯一例外是使用 Python 字符串创建数组,性能比对象数组稍慢,与固定宽度 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)

在这个例子中,对象 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 中基于迭代器的操作实现缓慢。当我们完成将这些操作重写为 ufunc 时,将能够实现大幅的性能提升。以 add ufunc 为例,我们已经为 StringDType 原型实现了该 ufunc。

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 使用对象字符串数组,这是直到最近才使用的默认设置,我们将看到在 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 支持的字符串数组,这正是为了避免这种情况以及与对象字符串数组相关的其他性能成本。

向后兼容性#

我们没有建议更改 Python 字符串的 DType 推断,并且预计不会对 NumPy 的现有用法产生任何影响。

详细描述#

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

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

Python API for StringDType#

新的 DType 可以通过 np.dtypes 命名空间访问。

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

此外,我们建议将字符 "T"(代表文本)保留用于 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 不受数组元素的最大长度的限制,任意长或短字符串都可以存在于同一个数组中,而无需为短字符串中的填充字节预留存储空间。

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

我们建议将 Python str 内置函数作为 DType 的标量类型。

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

虽然这确实在 NumPy 中创建了一个 API 缺陷,即从内置 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 的缺失数据值、字符串缺失数据值,还是两者都没有。我们将在下面解释如何处理这些情况。

谨慎的读者可能会担心需要处理三种不同类别的缺失数据哨兵的复杂性。这里的复杂性反映了对象数组的灵活性以及我们发现的下游使用模式。一些用户希望与哨兵的比较出错,因此他们使用 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))

字符串哨兵#

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

操作将直接使用哨兵值来表示缺失项。这是我们在下游代码中发现的这种模式的主要用法,其中类似 "__nan__" 的缺失数据哨兵传递给低级排序或分区算法。

其他哨兵#

任何其他 Python 对象在操作或比较中都会引发错误,就像 None 目前在对象数组中作为缺失数据哨兵一样。

>>> 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 循环,以及对 str_lenisalphaisdecimalisdigitisnumericisspacefindrfindcountstriplstriprstripreplace 字符串 ufunc 的实现,这些 ufunc 将在 NumPy 2.0 中新增。

isnan ufunc 将对类似 NaN 的哨兵项返回 True,对其他项返回 False。比较将根据 Unicode 代码点对数据进行排序,这与目前为固定宽度 Unicode DType 实现的方式相同。将来,NumPy 或下游库可能会为 NumPy Unicode 字符串数组添加区域感知排序、大小写折叠和规范化,但我们目前没有建议添加这些功能。

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

这种“兼容”实例的概念将在二元通用函数的 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 命名空间中可用,该命名空间将填充字符串通用函数。

>>> 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

这是目前对象字符串数组状况的延续,对象字符串数组只能使用 allow_pickle=True 选项保存到 npy 文件中。

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

用于 StringDType 的 C API#

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

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

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

The PyArray_StringDTypePyArray_StringDTypeObject 结构#

我们将公开 StringDType 元类的结构和 StringDType 实例类型的结构。前者 PyArray_StringDType 将与其他 PyArray_DTypeMeta 实例一样在 C API 中可用,用于编写通用函数和转换循环。此外,我们将公开以下结构

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;
}

公开此定义便于将来与其他数据类型集成。

字符串和分配器类型#

解包字符串在 C API 中使用 npy_static_string 类型表示,该类型将使用以下定义公开

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

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

此外,我们将公开两个不透明结构,npy_packed_static_stringnpy_string_allocator。每个 StringDType 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 通用函数接受两个字符串数组,并生成第三个字符串数组。这意味着通用函数循环需要三个分配器才能加载每个操作数的字符串并将结果打包到输出数组中。由于输入和输出操作数不必是不同的对象,并且操作数可以通过是同一个数组而共享分配器,这使得问题变得更加复杂。原则上,我们可以要求用户在通用函数循环内部获取和释放锁,但这与在循环设置中获取所有三个分配器并在循环结束之后同时释放它们相比,将增加大量的性能开销。

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

打包和加载字符串#

访问字符串由以下函数控制

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

此函数在发生错误时返回 -1,这可能发生在存在线程错误或损坏导致无法访问堆分配时。成功时,它可以返回 1 或 0。如果它返回 1,则表示打包字符串的内容为空字符串,在这种情况下,可以针对空字符串执行特殊逻辑。如果函数返回 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 的通用函数实现。如果我们被赋予指向 StringDType 数组条目开头的 const char *buf 指针,以及指向数组描述符的 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 数组的惯用类型化内存视图语法。我们有一些初步的想法,可以更新缓冲协议[9]或利用 Arrow C 数据接口[10]来公开不适用于缓冲协议的数据类型的 NumPy 数组,但这些工作可能不会及时在 NumPy 2.0 中完成。这意味着将使用固定宽度字符串数组的遗留 Cython 代码调整为使用 StringDType 将是困难的。将使用对象字符串数组的代码调整应该很简单,因为对象数组也不受缓冲协议的支持,并且在 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 表示,或指向 arena 分配的整数偏移量。 _short_string_buffer 表示对于小字符串优化最为有用,字符串数据存储在 direct_buffer 字段中,大小存储在 size_and_flags 字段中。 在这两种情况下,size_and_flags 字段都存储字符串的 size 以及位标志。 小字符串将大小存储在缓冲区的最后 4 位,保留 size_and_flags 的前 4 位用于标志。 堆字符串或 arena 分配中的字符串使用最高有效字节作为标志,保留前导字节用于字符串大小。 值得指出的是,这种选择限制了允许存储在数组中的最大字符串大小,特别是在 32 位系统上,每个字符串的限制为 16 兆字节 - 足够小,不必担心会影响实际工作流程。

在大端系统上,布局相反,size_and_flags 字段首先出现在结构体中。 这使实现始终可以使用 size_and_flags 字段的最高有效位作为标志。 这些结构体的依赖于字节序的布局是实现细节,在 API 中没有公开暴露。

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

有关这些内存布局中字符串的一些可视化示例,请参见 内存布局示例

Arena 分配器#

长度超过 15 字节的字符串(64 位系统)和长度超过 7 字节的字符串(32 位系统)存储在数组缓冲区外部的堆上。 与数组关联的 StringDType 实例附带的 arena 分配器管理分配的簿记。 分配器将作为不透明的 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;
};

这使我们可以将内存分配函数组合在一起,并在需要时在运行时选择不同的分配函数。 分配器的使用受互斥锁保护,有关线程安全性的更多讨论,请参见下文。

npy_string_arena 结构体成员处理内存分配,它具有以下布局

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

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

arena 中的每个字符串条目之前都带有大小,存储在 charsize_t 中,具体取决于字符串的长度。 长度在 16 或 8(取决于体系结构)到 255 之间的字符串存储在一个 char 大小中。 在内部,我们将这些字符串称为“中等”字符串。 此选择通过减少每条中等长度字符串 7 字节的开销来减少在堆上存储较小字符串的开销。 arena 中长度超过 255 字节的字符串具有设置的 NPY_STRING_LONG 标志。

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

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

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

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

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

变异和线程安全性#

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

一种解决方案是禁用变异,但不可避免地,对象字符串数组的下游使用将会变异我们希望支持的数组元素。

相反,我们选择将附加到 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。 标志指示该字符串存储在 arena 外部,并且已初始化。

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

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

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

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

空字符串和丢失数据#

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

丢失的字符串将具有相同的表示形式,只是它们始终有一个标志,即 NPY_STRING_MISSING 设置在标志字段中。 用户需要在访问未打包的字符串缓冲区之前检查字符串是否为空,并且我们已经以这样一种方式设置了 C API,即在每次解包字符串时强制执行空检查。 丢失的字符串和空字符串都可以根据已打包字符串表示形式中的数据进行检测,并且不需要在 arena 分配中或额外的堆分配中占用相应的空间。

实现#

我们有一个待合并到 NumPy 中的公开拉取请求 [12],它将添加 StringDType。

我们创建了 Pandas 的一个开发分支,该分支支持使用 StringDType 创建 Pandas 数据结构 [13]。 这说明了在下游库中支持 StringDType 所需的重构,这些库大量使用了对象字符串数组。

如果被接受,此 NEP 的大部分剩余工作将是更新文档和完善 NumPy 2.0 版本。 我们已经完成了以下工作

  • 创建一个 np.strings 命名空间,并将字符串 ufunc 直接暴露在该命名空间中。

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

我们将继续执行以下操作

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

替代方案#

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

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

讨论#

参考资料和脚注#