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 2 str 类型,然后在 NumPy 中添加了 Python 3 支持。bytes DType 在用于表示 Python 2 字符串或其他以 null 结尾的字节序列时最有效。但是,忽略尾随 null 表示 bytes_ DType 只适用于不包含尾随 null 的定宽字节流,因此对于需要通过 NumPy 字符串进行往返的通用字节流来说,它可能是一个存在问题匹配。

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 命名空间中公开字符串 ufunc,用于与字符串支持相关的函数和类型,从而为将来弃用 np.char 提供迁移路径。

以下内容不在此工作范围内

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

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

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

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

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

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

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

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

用法和影响#

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

希望将字符串数据加载到 NumPy 中并利用 NumPy 功能(如高级花式索引)的用户将拥有一个自然的选择,该选择与固定宽度 Unicode 字符串相比具有很大的内存节省,并且比对象字符串数组具有更好的验证保证和整体集成与 NumPy 的集成。迁移到一流的字符串 DType 也消除了在字符串操作期间获取 GIL 的需要,从而释放了对象字符串数组无法实现的未来优化。

性能#

在这里,我们简要描述了我们在 NumPy 之外使用实验性 DType API 实现的 StringDType 原型版本的初步性能测量。本节中的所有基准测试都在运行 Ubuntu 22.04 和使用 pyenv 编译的 Python 3.11.3 的 Dell 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 时,我们将实现大幅度的性能改进。以我们已为 StringDType 原型实现的 add 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,以及我们为初始实现选择的内存布局和堆分配策略的详细信息。

用于 StringDType 的 Python API#

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

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

此外,我们建议为 np.dtype 保留字符 "T"(文本的缩写),因此以上内容与以下内容相同:

>>> 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,现在DType类可以从np.dtypes命名空间导入,其余的NumPy DType API可以朝着这个方向发展,因此即使严格来说没有必要,我们也会在文档中包含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))

字符串哨兵#

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

运算将直接使用哨兵值表示缺失条目。这是我们在下游代码中发现的这种模式的主要用法,其中类似于"__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的实现,这些将在NumPy 2.0中新增。

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

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

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

>>> 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实现的内存布局和堆分配策略。

PyArray_StringDTypePyArray_StringDTypeObject结构体#

我们将公开StringDType元类的结构体,以及StringDType实例类型的结构体。前者PyArray_StringDType将与其他PyArray_DTypeMeta实例一样在C API中可用,用于编写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;
}

公开此定义可以简化将来与其他dtype的集成。

字符串和分配器类型#

解包字符串在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)。公开这些函数使得编写同时使用多个分配器的代码变得简单直接。当ufunc操作数共享描述符时,简单地多次调用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 编写 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 中添加特殊支持,否则 Cython 将不支持 StringDType 数组的惯用类型化内存视图语法。我们有一些初步的想法,可以通过更新缓冲区协议[9]或使用 Arrow C 数据接口[10]来公开 DType 的 NumPy 数组,这些 DType 在缓冲区协议中没有意义,但这些努力可能无法及时完成 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 以及位标志。小字符串将大小存储在缓冲区的最后四位中,为标志保留 size_and_flags 的前四位。堆字符串或 arena 分配中的字符串使用最高有效字节作为标志,为字符串大小保留前导字节。值得指出的是,这种选择限制了允许存储在数组中的最大字符串大小,尤其是在 32 位系统上,每个字符串的限制为 16 MB——足够小到不用担心会影响实际工作流程。

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

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

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

Arena 分配器#

在 64 位系统上超过 15 个字节,在 32 位系统上超过 7 个字节的字符串存储在数组缓冲区外部的堆上。分配的簿记由附加到与数组关联的 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 个字节来减少在堆上存储较小字符串的开销。长度超过 255 个字节的 arena 中的字符串已设置 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 等项目使用。发生这种情况时,我们可以更新锁定策略以允许多个同时读取线程,以及在删除 GIL 后 NumPy 中所需的线程错误的其他修复。

释放字符串#

在丢弃或重新使用打包字符串之前,必须释放现有字符串。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 中添加 StringDType 的开放拉取请求 [12]

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

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

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

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

我们将继续执行以下操作:

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

替代方案#

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

我们认为所提出的 DType 与改进的固定宽度二进制字符串 DType 不相互排斥,后者可以表示任意二进制数据或任何编码的文本,并且在将来添加这种 DType 将更容易,因为在添加 StringDType 后,NumPy 中对字符串数据支持的整体改进。

讨论#

参考文献和脚注#