开发者文档#

pylibcudf 是 libcudf 的一个轻量级 Cython 封装器。它旨在提供一个接近零开销的接口，用于在 Python 中访问 libcudf。使用调用 pylibcudf 的 Cython 化代码，应该能够实现接近原生 C++ 的性能，同时还允许从 Python 进行相当高效的使用。除了这些要求之外，pylibcudf 还必须与其他的 Python 库自然地集成。换句话说，它应该与标准的 Python 容器、社区协议（如 __cuda_array_interface__）以及常见的词汇类型（如 CuPy 数组）进行相当透明的互操作。

通用设计原则#

为了实现 pylibcudf 的目标，我们遵循以下设计原则集合

  • 每个公共函数或方法都应声明为 cpdef。这使得它们可以在 Cython 和 Python 代码中使用。与 cdef 函数相比，这会产生一些轻微的开销，但我们认为这是可以接受的，因为 1) 绝大多数用户将使用纯 Python 而不是 Cython，并且 2) cpdef 函数相对于 cdef 函数的开销在纳秒级别，而 CUDA 内核启动开销在微秒级别，因此这些函数开销在 pylibcudf 的典型使用中可以忽略不计。

  • 使用的每个变量都应强类型化，并且是原始类型（int, float 等）或 cdef 类。C++ 中的任何枚举都应使用 cpdef enum 来镜像，这将在 Cython 中创建 C 风格的枚举，并在 Python 中创建 PEP 435 风格的 Python 枚举，这些枚举将在 Python 中自动使用。

  • 代码中的所有类型标注都应使用 Cython 语法，而不是 PEP 484 Python 类型标注语法。这不仅确保了与 Cython < 3 的兼容性，而且即使在 Cython 3 中，PEP 484 的支持在撰写本文时仍然不完整。

  • 所有 cudf 代码应仅与 pylibcudf 交互，绝不直接与 libcudf 交互。目前并非如此，但这是库正在发展的方向。

  • 理想情况下，pylibcudf 除了 rmm 之外不应依赖任何 RAPIDS 组件，并且通常应具有最少的运行时依赖。

  • 类型存根是手动提供和生成的。添加新功能时，请确保相应的类型存根得到适当更新。

与 libcudf 的关系#

总的来说，pylibcudf 和 libcudf 之间的关系可以从两个方面来理解：数据结构和算法。

数据结构#

通常，libcudf 中的每个类型都应该有一个对应的 Cython cdef 类，该类具有一个属性 self.c_obj: unique_ptr[${underlying_type}]，用于拥有底层 libcudf 类型的一个实例。每个类型还应该实现一个相应的方法 cdef ${cython_type} from_libcudf(${underlying_type} dt)，以便能够从底层的 libcudf 实例构造 Cython 对象。根据类型的性质，该函数可能需要接受一个 unique_ptr 并获得所有权，例如 cdef ${cython_type} from_libcudf(unique_ptr[${underlying_type}] obj)。对于拥有 GPU 数据的类型来说，通常会是这种情况，可能需要进一步规范化。

例如，libcudf::data_type 映射到 pylibcudf.DataType，其结构如下（实现省略）

cdef class DataType:
    cdef data_type c_obj

    cpdef TypeId id(self)
    cpdef int32_t scale(self)

    @staticmethod
    cdef DataType from_libcudf(data_type dt)

这使得 pylibcudf 函数可以接受一个类型化的 DataType 参数，然后通过访问参数的 c_obj 来轻松调用底层的 libcudf 算法。

pylibcudf 的表和列#

上述规则集的主要例外是 libcudf 的核心数据拥有类型，cudf::table 和 cudf::column。libcudf 使用基于智能指针的现代 C++ 惯用法来避免资源泄露并使代码异常安全。为了避免传递原始指针，并确保所有权语义清晰，libcudf 具有与数据拥有类型相对应的独立 view 类型。例如，cudf::column 拥有数据，而 cudf::column_view 表示对数据列的视图，cudf::mutable_column_view 表示可变视图。column_view 不一定必须引用由 cudf::column 拥有的数据；任何内存缓冲区都可以。这种分离允许 libcudf 算法清晰地传达所有权预期，并允许多个视图同时存在于同一数据中。

虽然 libcudf 算法接受视图作为输入，但任何分配数据的算法必须返回 cudf::column 和 cudf::table 对象。libcudf 的所有权模型对 pylibcudf 来说是个问题，pylibcudf 必须能够与 PyTorch 或 Numba 等其他 Python 库提供的数据无缝互操作。因此，pylibcudf 采用以下策略

  • pylibcudf 定义了 gpumemoryview 类型，它（类似于 Python memoryview 类型）表示一个视图，指向由另一个对象拥有的内存，并通过 Python 标准的引用计数机制使其保持活跃。gpumemoryview 可以由任何实现 CUDA 数组接口协议 的对象构造。

    • 这种类型最终将被通用化，以便在 pylibcudf 外部重用。

  • pylibcudf 定义了自己的 Table 和 Column 类。

    • Table 会维护对其包含的 Column 的 Python 引用，因此多个 Table 可以共享同一个 Column。

    • Column 由其数据缓冲区（对于嵌套类型可能包含子项）和其空值掩码的 gpumemoryview 组成。

  • pylibcudf.Table 和 pylibcudf.Column 提供了对查看相同列/内存的 cudf::table_view 和 cudf::column_view 对象的便捷访问。在基于底层的 libcudf 算法实现任何 pylibcudf 算法时，可以使用这些视图对象。具体来说，这些类中的每一个都拥有 libcudf 视图类型的一个实例，并提供一个 view 方法，该方法可用于访问该对象的指针，以便传递给 libcudf。

算法#

pylibcudf 算法应该看起来几乎与 libcudf 算法完全相同。任何 libcudf 函数都应在 pylibcudf 中以完全相同的签名进行镜像，并将 libcudf 类型映射到相应的 pylibcudf 类型。所有对 libcudf 算法的调用都应尽早执行任何必需的 Python 预处理，然后在调用 libcudf 之前释放 GIL。例如，以下是 gather 的实现

cpdef Table gather(
    Table source_table,
    Column gather_map,
    OutOfBoundsPolicy bounds_policy
):
    cdef unique_ptr[table] c_result
    with nogil:
        c_result = move(
            cpp_copying.gather(
                source_table.view(),
                gather_map.view(),
                bounds_policy
            )
        )
    return Table.from_libcudf(move(c_result))

上面的代码片段有几个值得注意的地方

  • 从 libcudf 返回的对象立即转换为 pylibcudf 类型。

  • cudf::gather 接受一个 cudf::out_of_bounds_policy 枚举参数。OutOfBoundsPolicy 是 pylibcudf 中此类型的一个别名，与我们的 Python 命名约定（使用 CapsCase 而不是 snake_case）相匹配。

测试#

在编写 pylibcudf 测试时，重要的是要记住所有 API 都应该已经在 libcudf 的 C++ 层中进行了测试。pylibcudf 测试的主要目的是确保 绑定 的正确性；底层实现的正确性通常应在 libcudf 中进行验证。如果 pylibcudf 测试发现了 libcudf 的错误，则应添加相应的 libcudf 测试来覆盖此情况，而不是仅仅依赖于 pylibcudf 测试。

pylibcudf 的 conftest.py 包含一些标准的参数化 dtype fixture 列表，这些列表又可以用于参数化其他 fixture。分配数据的 fixture 应尽可能利用这些 dtype 列表，以简化对重要类型矩阵的测试。在适当的情况下，可以添加新的 fixture 列表。

为了尽可能高效地运行测试，测试套件应充分利用 fixture。遵循的最简单通用结构是让 pyarrow array/table/scalar fixture 由 dtype 列表之一进行参数化。然后，可以使用简单的 from_arrow 调用创建相应的 pylibcudf fixture。这种方法确保了各种测试在类型上的全局一致性覆盖。

一般来说，pylibcudf 测试应优先对照相应的 pyarrow 实现进行验证，而不是硬编码数据。如果没有 pyarrow 实现，另一种替代方法是编写一个纯 Python 实现，该实现循环遍历 Table/Column 的值，前提是存在 pylibcudf 实现的标量 Python 等效项（这对于字符串方法尤其重要）。

这种方法对输入数据的变化更具弹性，特别是在考虑上述 fixture 策略的情况下。utils 模块提供了用于比较 pylibcudf 和 pyarrow 类型之间的标准工具。

以下是演示上述要点的示例

import pyarrow as pa
import pyarrow.compute as pc
import pytest
from cudf._lib import pylibcudf as plc
from utils import assert_column_eq

# The pa_dtype fixture is defined in conftest.py.
@pytest.fixture(scope="module")
def pa_column(pa_dtype):
    pa.array([1, 2, 3])


@pytest.fixture(scope="module")
def column(pa_column):
    return plc.interop.from_arrow(pa_column)


def test_foo(pa_column, column):
    index = 1
    result = plc.foo(column)
    expected = pa.foo(pa_column)

    assert_column_eq(result, expected)

关于应测试内容的指导原则

  • 测试应该全面覆盖 API，包括确保良好测试覆盖率所需的所有可能的参数组合。

  • pylibcudf 不应尝试对大型数据量进行压力测试，而应转交给 libcudf 测试。

    • 例外情况：在 C++ 中难以构建合适的大型测试（例如为 I/O 测试创建合适的输入数据）的特殊情况下，可以改为在 pylibcudf 中添加测试。

  • 始终应测试可为空的数据。

  • 应该测试预期的异常。编写测试时应从用户的角度出发，如果 API 当前没有抛出适当的异常，则应该更新它。

    • 重要说明：如果异常应该由 libcudf 生成，则应更新底层的 libcudf API 以在 C++ 中抛出所需的异常。在非平凡的情况下，此类更改可能需要与 libcudf 开发人员协商。此问题 提供了一个概述并指出了应涵盖大多数用例的可接受异常类型。在极少数情况下，可能需要在 error.hpp 中引入新的 C++ 异常。如果是这样，此异常还需要在 exception_handler.pxd 中映射到适当的 Python 异常。

关于如何最好地使用 pytest 的一些指导原则。

  • 默认情况下，生成设备数据容器的 fixture 应该是模块范围的，并被测试视为不可变。在 GPU 上分配数据开销很大且会减慢测试速度。几乎所有的 pylibcudf 操作都是非就地操作，因此模块范围的 fixture 通常不会带来问题。会话范围的 fixture 也可以工作，但它们更难理解，因为它们位于不同的模块中，并且如果它们因任何原因需要更改，可能会影响任意数量的测试。模块范围是一个很好的平衡。

  • 在必要的情况下，可变 fixture 应按此命名（例如 mutable_col）并具有函数范围。如果可能，可以通过简单地复制相应的模块范围不可变 fixture 来实现它们，以避免重复生成逻辑。

测试应按照 pylibcudf 模块进行组织，即每个 pylibcudf 模块对应一个测试模块。

cuDF Python 测试指南的以下章节通常也适用于 pylibcudf，除非被上述任何陈述取代

其他说明#

Cython 作用域枚举#

Cython 3 引入了对作用域枚举的支持。然而，这种支持存在一些错误以及一些容易掉进去的陷阱。我们对枚举的使用旨在最小化代码的复杂性，同时规避 Cython 的限制。

警告

随着 Cython 的更新和我们对最佳实践理解的演进，本节的指导可能会经常更改。

  • 所有声明 C++ 枚举的 pxd 文件都应该使用 cpdef enum class 声明。

    • 原因：此声明使 C++ 枚举可在 Cython 代码中使用，同时透明地创建 Python 枚举。

  • 任何仅包含 C++ 声明的 pxd 文件，如果其中任何声明是作用域枚举，则仍必须具有相应的 pyx 文件。

    • 原因：Python 枚举的创建需要 Cython 实际生成必要的 Python C API 代码，如果仅存在 pxd 文件，则不会发生这种情况。

  • 如果某个 C++ 枚举是 pylibcudf 模块公共 API 的一部分，则应将其直接导入（而非 cimport）到 pyx 文件中，并使用符合我们 Python 类命名约定 (CapsCase) 而非 C++ 命名约定 (snake_case) 的名称进行别名化。

  • 原因：我们希望将枚举暴露给模块的 Python 和 Cython 使用者。作为副作用，这种别名化避免了 此 Cython 错误。

  • 注意：一旦上述 Cython 错误得到解决，当 cimport 枚举时，也应该在 pylibcudf pxd 文件中对其进行别名化，以便 Python 和 Cython 的用法一致。

以下是枚举的适当使用示例。

# pylibcudf/libcudf/copying.pxd
cdef extern from "cudf/copying.hpp" namespace "cudf" nogil:
    # cpdef here so that we export both a cdef enum class and a Python enum.Enum.
    cpdef enum class out_of_bounds_policy(bool):
        NULLIFY
        DONT_CHECK


# pylibcudf/libcudf/copying.pyx
# This file is empty, but is required to compile the Python enum in pylibcudf/libcudf/copying.pxd
# Ensure this file is included in pylibcudf/libcudf/CMakeLists.txt


# pylibcudf/copying.pxd

# cimport the enum using the exact name
# Once https://github.com/cython/cython/issues/5609 is resolved,
# this import should instead be
# from pylibcudf.libcudf.copying cimport out_of_bounds_policy as OutOfBoundsPolicy
from pylibcudf.libcudf.copying cimport out_of_bounds_policy


# pylibcudf/copying.pyx
# Access cpp.copying members that aren't part of this module's public API via
# this module alias
from pylibcudf.libcudf cimport copying as cpp_copying
from pylibcudf.libcudf.copying cimport out_of_bounds_policy

# This import exposes the enum in the public API of this module.
# It requires a no-cython-lint tag because it will be unused: all typing of
# parameters etc will need to use the Cython name `out_of_bounds_policy` until
# the Cython bug is resolved.
from pylibcudf.libcudf.copying import \
    out_of_bounds_policy as OutOfBoundsPolicy  # no-cython-lint

处理 libcudf 中的重载函数#

作为 C++ 库，libcudf 大量使用了函数重载。例如，libcudf 中存在以下两个函数

std::unique_ptr<table> empty_like(table_view const& input_table);
std::unique_ptr<column> empty_like(column_view const& input);

然而，Cython 不直接支持这种方式的重载，而是遵循 Python 的语义，即每个函数名必须唯一标识一个函数。因此，在实现上述重载函数的 pylibcudf 封装器时，应使用 Cython 的 融合类型 (fused types)。融合类型是 Cython 的泛型编程版本，在这种情况下相当于编写模板函数，这些函数会编译成与不同 C++ 重载相对应的独立副本。对于上述函数，等效的 Cython 函数是

ctypedef fused ColumnOrTable:
    Table
    Column

cpdef ColumnOrTable empty_like(ColumnOrTable input)

Cython 支持根据参数类型对融合类型函数的内容进行专门化处理，因此可以使用适当的条件来编码任何特定于类型的逻辑。请参阅 pylibcudf 源代码以获取实现此类函数的示例。

如果 libcudf 为同一函数提供了多个重载且参数数量不同，请在 Cython 定义中指定最大参数数量，并将重载之间不共享的参数设置为 None。如果用户尝试为特定重载类型传递不受支持的参数，则应引发 ValueError。

最后，如果您认为这种不一致可以在 libcudf 侧解决，请考虑提交 libcudf issue。

类型存根#

由于像 mypy 和 pyright 这样的静态类型检查器无法解析 Cython 代码，我们为 pylibcudf 包提供了类型存根。这些目前是手动维护的，与相应的 pylibcudf 文件并行。

每个 pyx 文件都应该有一个匹配的 pyi 文件来提供类型存根。大多数函数都可以直接暴露。一些指导原则

  • 对于 libcudf 中带类型的整数参数，使用 int 作为类型标注。

  • 对于在 Cython 中被标注为 list，但函数体进行了更详细检查的函数，请尝试在类型中编码详细信息。

  • 对于 Cython 融合类型，有两种选择

    1. 如果融合类型在函数签名中只出现一次，使用 Union 类型；

    2. 如果融合类型出现多次（或同时作为输入和输出类型），使用 TypeVar，并将融合类型中的变体作为约束提供。

举例来说，pylibcudf.copying.split 在 Cython 中的类型是

ctypedef fused ColumnOrTable:
    Table
    Column

cpdef list split(ColumnOrTable input, list splits): ...

在这里，我们只使用了融合类型一次，并且 list 参数没有指定它们的值。在这里，如果输入提供 Column，则输出接收 list[Column]，如果输入提供 Table，则输出接收 list[Table]。

在类型存根中，我们可以使用 TypeVar 来编码这一点，我们还可以为 splits 参数提供类型信息，表明分割值必须是整数

ColumnOrTable = TypeVar("ColumnOrTable", Column, Table)

def split(input: ColumnOrTable, splits: list[int]) -> list[ColumnOrTable]: ...

相反，pylibcudf.copying.scatter 在其输入中只使用了融合类型一次

ctypedef fused TableOrListOfScalars:
    Table
    list

cpdef Table scatter(
    TableOrListOfScalars source, Column scatter_map, Table target
)

在这种情况下，在类型存根中我们可以使用正常的联合类型

def scatter(
    source: Table | list[Scalar], scatter_map: Column, target: Table
) -> Table: ...