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 进行字符串操作。

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

动机与范围#

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

NumPy 中字符串支持的历史#

NumPy 对文本数据的支持是根据早期用户需求以及 Python 生态系统的变化有机演变的。

NumPy 中增加了对字符串的支持,以支持 `NumArray` 的 `chararray` 类型的用户。这方面的遗留影响至今仍在 NumPy API 中可见:与字符串相关的函数位于 `np.char` 中,以支持 `np.char.chararray` 类。这个类并非正式弃用,但自 NumPy 1.4 起,其模块文档字符串中已有一条注释建议改用字符串 dtypes。

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

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

随着 NumPy 对 Python 3 支持的到来,由于向后兼容性的考虑,字符串 DTypes 大多被保留下来,尽管 `unicode` DType 成为了 `str` 数据的默认 DType,而旧的 `string` DType 被重命名为 `bytes_` DType。这个改变使得 NumPy 处于一个欠佳的境地:它将一个最初用于以 null 结尾的字节串的数据类型作为所有 Python `bytes` 数据的类型,并且默认的字符串类型在内存中的表示比表示单字节 ASCII 或 Latin-1 编码所需的数据消耗四倍的内存。

固定宽度字符串的问题#

现有的两种字符串 DTypes 都代表固定宽度序列,允许在数组缓冲区中存储字符串数据。这避免了向 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 字符串之间透明转换的方式处理“NumPy 字符串数组”,包括支持缺失字符串。`object` DType 部分满足了这一需求,尽管性能较慢且没有类型检查。

作为本次讨论的结果,改进字符串数据的支持已被添加到 NumPy 项目路线图[6] 中,明确要求添加一种更适合内存映射具有任何或无编码的字节的数据类型,以及一种支持缺失数据的可变宽度字符串 DType,以替换 `object` 字符串数组的用法。

拟议工作#

本 NEP 提议为 NumPy 添加 `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 优化。

  • 更新 `npy` 和 `npz` 文件格式,以允许存储任意长度的附加数据。

尽管我们明确排除将这些项目作为此工作的一部分来实现,但添加新的字符串 DType 有助于为未来实现其中一些项目的工作奠定基础。

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

虽然我们不提议向 NumPy 添加缺失数据标记,但我们提议支持一个可选的、用户提供的缺失数据标记,因此这在一定程度上使 NumPy 更接近官方支持缺失数据。我们正试图避免解决 NEP 26 中描述的分歧NEP 26,此提案不要求也不排除将来向 `ndarray` 添加缺失数据标记或基于位标志的缺失数据支持。

用法与影响#

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

希望将字符串数据加载到 NumPy 中并利用 NumPy 特性(如花式高级索引)的用户将有一个自然的选项,与固定宽度 Unicode 字符串相比,它提供了显著的内存节省,并且与 `object` 字符串数组相比,它具有更好的验证保证和与 NumPy 的整体集成。转向一流的字符串 DType 还消除了在字符串操作期间获取 GIL 的需要,从而为使用 `object` 字符串数组不可能实现的未来优化开辟了道路。

性能#

在此,我们简要描述了我们使用实验性 DType API 在 NumPy 之外实现的 `StringDType` 原型版本的初步性能测量。本节中的所有基准测试均在运行 Ubuntu 22.04 和 Python 3.11.3(使用 pyenv 编译)的 Dell XPS 13 9380 上进行。NumPy、Pandas 和 `StringDType` 原型都使用 meson release 版本编译。

目前,`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=np.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` DTypes 的速度明显更快,因为 `data` 列表中的对象可以直接在数组中进行内插,而 `StrDType` 和 `StringDType` 需要复制字符串数据,并且 `StringDType` 需要将数据转换为 UTF-8 并执行数组缓冲区之外的其他堆分配。未来,如果 Python 的字符串内部表示迁移到 UTF-8,`StringDType` 的字符串加载性能应该会有所提高。

字符串操作性能相似。

In [7]: %timeit np.array([s.capitalize() for s in data], dtype=np.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。接下来,我们将描述缺失数据处理支持和数组元素严格字符串类型检查支持。然后,我们将讨论我们将定义的转换和 ufunc 实现,并讨论我们为新的 `np.strings` 命名空间制定的计划,以直接在 Python API 中公开字符串 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` 不以数组元素的aught 长度为参数,任意长或短的字符串可以在同一个数组中存在,而无需为短字符串的填充字节保留存储空间。

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

我们提议将 Python 内置的 `str` 关联为 DType 的标量类型。

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

虽然这会产生一个 API 上的不便之处,即内置 DType 类到 NumPy 中标量的映射将不再是一对一的(`unicode` DType 的标量类型是 `str`),但这避免了需要为该目的定义、优化或维护 `str` 子类,或者其他用于维持这种一对一映射的技巧。为保持向后兼容性,为 Python 字符串列表检测到的 DType 将保持为固定宽度的 Unicode 字符串。

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

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

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

我们还希望在未来能够添加一种新的固定宽度文本版本的 `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 本身,以便 NumPy 专门查找 NaN 的操作能够正常工作,而无需在 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 DTypes 的完整的往返转换。此外,我们将为比较运算符添加实现,以及一个接受两个字符串数组的 `add` 循环,接受字符串和整数数组的 `multiply` 循环,一个 `isnan` 循环,以及 `str_len`、`isalpha`、`isdecimal`、`isdigit`、`isnumeric`、`isspace`、`find`、`rfind`、`count`、`strip`、`lstrip`、`rstrip` 和 `replace` 字符串 ufuncs 的实现,这些将在 NumPy 2.0 中新可用。

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

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

这种“兼容”实例的概念将在二进制 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,一旦下游库支持的最低 NumPy 版本(根据 SPEC-0)足够高,可以安全地切换到 np.strings 而无需进行任何依赖于 NumPy 版本的逻辑条件判断。

序列化#

由于字符串数据存储在数组缓冲区之外,因此序列化到 npy 格式需要格式修订才能支持存储可变长度的辅助数据。我们不打算在此过程中完成此操作,因此在未指定 allow_pickle=True 的情况下,我们不打算支持序列化到 npynpz 格式。

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

未来我们可能会决定添加对此的支持,但应谨慎行事,以免破坏可能未维护的 NumPy 外部解析器。

C API for StringDType#

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

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

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

PyArray_StringDTypePyArray_StringDTypeObject 结构#

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

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 编码的字节流开头的 const 指针,其中包含字符串数据。这是对字符串的只读视图,我们不会公开修改这些字符串的接口。我们不会在字节流后附加空字符,因此尝试将 buf 字段传递给期望 C 字符串的 API 的用户必须创建一个带有空字符结尾的副本。未来,如果事实证明复制到 null 终止缓冲区对 C API 的下游用户来说成本过高,我们可能会决定始终写入 null 字节。

此外,我们将公开两个不透明结构: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 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,则表示打包字符串的内容是空字符串,在这种情况下可以执行处理空字符串的特殊逻辑。如果函数返回 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 实现。如果我们获得指向 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 不支持 StringDType 数组的惯用类型内存视图语法,除非将来在 Cython 中添加特殊支持。我们有一些初步的想法,可以通过更新缓冲区协议[9]或利用 Arrow C 数据接口[10]来公开不适合缓冲区协议的 DType 的 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 表示最适合直接表示位于堆上的字符串或位于内存区域分配中的字符串,其中 offset 字段包含地址的 size_t 表示,或者内存区域分配中的整数偏移量。_short_string_buffer 表示最适合短字符串优化,其中字符串数据存储在 direct_buffer 字段中,大小存储在 size_and_flags 字段中。在这两种情况下,size_and_flags 字段都存储字符串的 size 以及位标志。短字符串将大小存储在缓冲区的最后四位中,为标志保留 size_and_flags 的前四位。堆字符串或内存区域分配中的字符串使用最高有效字节作为标志,将前几个字节留给字符串大小。值得指出的是,这个选择限制了允许存储在数组中的字符串的最大大小,尤其是在 32 位系统上,每个字符串的限制为 16MB——这足以影响实际工作流程。

在大端系统上,布局是相反的,size_and_flags 字段出现在结构体的开头。这使得实现能够始终使用 size_and_flags 字段的最高有效位作为标志。这些结构体的字节序相关的布局是一个实现细节,并且不在 API 中公开。

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

有关这些内存布局的直观示例,请参见 内存布局示例

内存区域分配器#

长度超过 15 字节(64 位系统)和 7 字节(32 位系统)的字符串存储在数组缓冲区之外的堆上。分配的记账由附加到数组关联的 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;
};

这允许我们将内存分配函数分组,并在运行时选择不同的分配函数。分配器的使用由互斥锁保护,有关线程安全的更多讨论请参见下文。

内存分配由 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 数组可变性的更多详细信息,请参见下文。

使用每个数组的内存区域分配器可确保附近数组条目的字符串缓冲区通常位于堆上的附近。我们不对相邻数组条目在堆上的连续性做保证,以支持短字符串优化、缺失数据以及允许修改数组条目。有关这些主题如何影响内存布局的更多讨论,请参见下文。

修改和线程安全#

修改会引入数据竞争和使用后释放错误的可能性,当数组被多个线程访问和修改时。此外,如果我们在内存区域缓冲区中分配修改后的字符串,并强制连续存储,将旧字符串替换为新字符串,那么修改单个字符串可能会触发对整个数组的内存区域缓冲区的重新分配。与对象字符串数组或固定宽度字符串相比,这是一种非常糟糕的性能退化。

一种解决方案是禁用修改,但不可避免地会有一些对象字符串数组的下游用途会修改数组元素,而我们希望支持这些用途。

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

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

释放字符串#

在丢弃或重用打包字符串之前,必须先释放现有字符串。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 返回的新创建的数组缓冲区将默认填充空字符串,因为一个没有设置标志的字符串是一个未初始化的零长度内存区域字符串。这不是空字符串的唯一有效表示,因为可以设置其他标志来指示空字符串与预先存在的短字符串或内存区域字符串关联。

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

实现#

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

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

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

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

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

我们将继续做以下工作:

  • 处理 NumPy 中与新 DType 相关的剩余问题。特别是,我们已经知道 NumPy 中 copyswap 的剩余用法应迁移到使用 cast 或一个尚未添加的单元素复制 DType API 插槽。我们还需要确保 DType 类可以在 Python API 中与 DType 实例互换使用(在所有合适的地方),并在所有其他 DType 实例可以传递但 DType 类不适用的地方添加有用的错误。

替代方案#

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

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

讨论#

参考文献和脚注#