Ragged

介绍

Ragged 是一个库，允许像处理 NumPy 或 CuPy 数组一样操作乱序数组，遵循 Array API 规范。

例如，这是一个 乱序/乱序数组

>>> import ragged
>>> a = ragged.array([[[1.1, 2.2, 3.3], []], [[4.4]], [], [[5.5, 6.6, 7.7, 8.8], [9.9]]])
>>> a
ragged.array([
    [[1.1, 2.2, 3.3], []],
    [[4.4]],
    [],
    [[5.5, 6.6, 7.7, 8.8], [9.9]]
])

值都是浮点数，因此 a.dtype 是 float64，

>>> a.dtype
dtype('float64')

但 a.shape 具有非整数维度，以考虑到其列表长度可能不均匀

>>> a.shape
(4, None, None)

通常，ragged.array 可以具有任何混合的规则和不规则维度，尽管 shape[0]（长度）始终是整数。这个约定遵循 Array API 对 array.shape 的规范，它必须是 int 或 None 的元组

array.shape: Tuple[Optional[int], ...]

(我们使用 None 来表示没有单一值大小的维度，这与 Array API 指定 未知 大小维度的意图不同，但它遵循技术规范。Array API 消费者库可以尝试使用 Ragged 来找出它们是否准备好处理乱序数组。)

所有正常的元素级和减少函数都适用，以及切片

>>> ragged.sqrt(a)
ragged.array([
    [[1.05, 1.48, 1.82], []],
    [[2.1]],
    [],
    [[2.35, 2.57, 2.77, 2.97], [3.15]]
])

>>> ragged.sum(a, axis=0)
ragged.array([
    [11, 8.8, 11, 8.8],
    [9.9]
])

>>> ragged.sum(a, axis=-1)
ragged.array([
    [6.6, 0],
    [4.4],
    [],
    [28.6, 9.9]
])

>>> a[-1, 0, 2]
ragged.array(7.7)

>>> a[a * 10 % 2 == 0]
ragged.array([
    [[2.2], []],
    [[4.4]],
    [],
    [[6.6, 8.8], []]
])

所有在 Array API 中的方法、属性和函数都将被实现于 Ragged 中，包括不必要遵循 Array API 的便利性。查看标记为 "todo" 的 开放问题 以了解仍需编写的 Array API 函数（总共 120 个）。

Ragged 有两个 device 值，分别为 "cpu"（由 NumPy 支持）和 "cuda"（由 CuPy 支持）。最终，所有操作对于 CPU 和 GPU 将是相同的。

实现

Ragged 是使用 Awkward Array （代码，文档）实现的，它是一个适用于任意树状（类似 JSON）数据的数组库。由于其通用性，Awkward Array 无法遵循 Array API——实际上，它的数组对象不能有独立的 dtype 和 shape 属性（数组的 type 无法分解）。因此，Ragged 是

• 特别化 Awkward Array 以处理固定长度和可变长度的列表中的数值数据，
• 并且是 形式化 以符合 Array API 和其完全类型化的协议。

在 为什么这个库存在？ 中了解更多详情，请访问 讨论 标签。

Ragged 是 Awkward Array 的一个薄包装，限制其只处理稀疏数组，并将函数参数和返回值转换为符合规范。

Awkward Array 则在时间和内存效率上都很高，适合大数据集。以下是一个例子

import gc      # control for garbage collection
import psutil  # measure process memory
import time    # measure time

import math
import ragged

this_process = psutil.Process()

def measure_memory(task):
    gc.collect()
    start_memory = this_process.memory_full_info().uss
    out = task()
    gc.collect()
    stop_memory = this_process.memory_full_info().uss
    print(f"memory: {(stop_memory - start_memory) * 1e-9:.3f} GB")
    return out

def measure_time(task):
    gc.disable()
    start_time = time.perf_counter()
    out = task()
    stop_time = time.perf_counter()
    gc.enable()
    print(f"time: {stop_time - start_time:.3f} sec")
    return out

def make_big_python_object():
    out = []
    for i in range(10000000):
        out.append([j * 1.1 for j in range(i % 10)])
    return out

def make_ragged_array():
    return ragged.array(pyobj)

def compute_on_python_object():
    out = []
    for row in pyobj:
        out.append([math.sqrt(x) for x in row])
    return out

def compute_on_ragged_array():
    return ragged.sqrt(arr)

ragged.array 的大小减少了 3 倍

>>> pyobj = measure_memory(make_big_python_object)
memory: 2.687 GB

>>> arr = measure_memory(make_ragged_array)
memory: 0.877 GB

并且对其进行的样本计算（每个值的平方根）速度快了 50 倍

>>> result = measure_time(compute_on_python_object)
time: 4.180 sec

>>> result = measure_time(compute_on_ragged_array)
time: 0.082 sec

Awkward Array 和 Ragged 通常比它们的 Python 等价物更小、更快，原因与 NumPy 比起 Python 列表更小、更快的原因相同。有关更多信息，请参阅 Awkward Array 的 论文和演示文稿。

安装

Ragged 在 PyPI

pip install ragged

上，并将在将来出现在 conda-forge 上。

ragged 是一个纯 Python 库，它只依赖于 awkward（该库反过来只依赖于 numpy 和一个编译扩展）。原则上（即最终），ragged 可以加载到 Pyodide 和 JupyterLite。

致谢

此工作的支持由 NSF 奖助金 OAC-2103945 和 Awkward Array 贡献者 的慷慨帮助提供。