如何使用 chDB 使 Pandas DataFrame 查询提速 87 倍

我从开始开发 chDB（一个在进程中运行的 ClickHouse 嵌入式版本）已经接近两年，并且在六个月前 chDB 正式加入了 ClickHouse。在这篇博客中，我将分享过去几个月中进行的性能优化工作。

构建 chDB 时，我面临的一个早期挑战是在不损失性能的前提下，基于 ClickHouse Local 支持从多种数据源进行查询。值得注意的是，ClickHouse Local 的所有数据输入和输出都是通过文件描述符完成的，如下图所示。

虽然对 ClickHouse Local 来说这并不是问题，但对于 chDB 这样需要直接处理 Pandas、Numpy 或 PyArrow 等库生成数据的进程内引擎，这确实带来了一些挑战。

为了在这些数据上执行高效的 SQL 查询，chDB 需要满足以下几点要求：

1. 零复制：通过利用 Python 的 memoryview，实现 ClickHouse 与 Python 进程之间的直接内存映射。

2. 向量化读取：充分利用现有统计数据，并且在读取过程中考虑 CPU 和内存的硬件特性。

chDB 的初始版本在设计时追求简化。以处理内存数据为例，针对 Pandas DataFrame 的查询，chDB 初始版本的实现步骤如下：

1. 将内存中的 DataFrame 序列化为 Parquet 格式，并写入临时文件或 memfd。我们最初考虑过将数据序列化为 Arrow 缓冲区，但测试结果显示 Parquet 更快。

2. 修改 SQL 语句中的数据源表，将其替换为 ClickHouse 的文件表引擎，并传入临时文件的文件描述符。

3. 运行 ClickHouse 引擎，并将输出格式设置为 Parquet。

4. 读取 Parquet 缓冲区，并将其转换为 DataFrame。

这一实现导致大部分的时间都花费在序列化、反序列化以及内存复制的过程中。即使我们使用了 memfd，性能仍然不尽如人意。

如下面的图表所示，在 ClickBench 基准测试中，几乎每个查询都需要 chDB 超过 30 秒的时间。

2024 年 6 月，chDB v2 引入了 SQL on DataFrame 功能，让您可以轻松地将 DataFrame 变量当作表来运行 SQL，如下所示：

import chdb
df = pd.DataFrame({"a": [1, 2, 3], "b": ["one", "two", "three"]})chdb.query("SELECT * FROM Python(df)").show()

同样，您也可以对 Numpy 数组、PyArrow 表以及 Python 字典变量进行类似的查询。

import chdbimport pandas as pdimport pyarrow as pa
data = {    "a": [1, 2, 3, 4, 5, 6],    "b": ["tom", "jerry", "auxten", "tom", "jerry", "auxten"],}chdb.query("""SELECT b, sum(a) FROM Python(data) GROUP BY b"").show()
arrow_table = pa.table(data)chdb.query("""SELECT b, sum(a) FROM Python(arrow_table) GROUP BY b""").show()

在接下来的部分，我将分享我们是如何解决这些效率挑战的，并使 chDB 成为全球最快的 SQL on DataFrame 引擎的过程。

我发现为 ClickHouse 添加一个表函数的过程相对简单。这个过程分为三个步骤，且不需要涉及复杂的 C++ 逻辑：

1. 声明并注册一个名为 TableFunctionPython 的表函数。

2. 定义 StoragePython 的逻辑，主要解决如何读取数据和获取表的架构。

3. 定义 PythonSource 的逻辑，主要用于在并发管道中返回数据块。

首先，我们需要声明并注册一个名为 TableFunctionPython 的表函数。这样，ClickHouse 的解析器在将 SQL 转换为 AST 时，就知道存在一个名为 Python 的表引擎：

注册过程也非常简单，主要是提供必要的描述和示例，并声明 “Python” 是不区分大小写的：

StoragePython 类的主要作用是作为 ClickHouse 数据处理管道的一部分，而数据处理的大部分工作实际上是在 PythonSource 类中完成的。在早期版本的 ClickHouse 中，IStorage.read 函数负责数据处理。但现在，它已经成为物理执行计划的一部分。

PythonSource 类继承自 ISource，而 ISource.generate 函数负责在管道开始运行后生成数据。

虽然 chDB 的 Python 表引擎整体架构相对简单，但在实现过程中我们遇到了许多意想不到的问题。大部分问题源自 C++ 和 Python 交互时的性能瓶颈。

例如，当在 chDB 中读取内存中的 Pandas DataFrame 时，不可避免地要调用部分 CPython（Python 的 C 实现，也是最常用的实现）。这带来了两个主要挑战：GIL（全局解释器锁）和对象引用计数。

由于 Python 存在 GIL（全局解释器锁），任何 CPython 函数调用都需要先获取 GIL。如果 GIL 的锁定范围过大，ClickHouse 的多线程引擎将会因受限于 GIL 而退化为串行执行；而如果 GIL 的锁定范围过小，线程之间频繁争夺锁会导致性能下降，甚至比单线程程序还要慢。

在 CPython 中，全局解释器锁（GIL）是一个用于保护 Python 对象访问的互斥锁，防止多个线程同时执行 Python 字节码。GIL 的存在主要是因为 CPython 的内存管理并非线程安全。有关 GIL 如何确保线程安全的详细解释可以参考此链接。

--- https://wiki.python.org/moin/GlobalInterpreterLock

Python 拥有自动垃圾回收功能，编写代码变得很简单。然而，如果在内存中引用了现有的 Python 对象或创建了新对象，就需要管理引用计数，否则会导致内存泄漏或双重释放错误。

因此，在 chDB 中，我们必须尽量避免调用 CPython API 函数来操作 Python 对象。这听起来可能有些极端，但确实是我们提高性能的唯一选择。

接下来，我将简要介绍我们是如何克服这些困难，尽管 Python 本身存在限制，仍然使 chDB 成为了 Pandas DataFrame 上最快的 SQL 引擎之一。

在这篇博客的开头，我们提到 chDB v1 由于需要多达四次额外的序列化和反序列化步骤，导致在 ClickBench 中的每个查询执行时间都超过了 30 秒。

第一个优化目标是减少这些额外的步骤，并直接读取 Python 对象。这一优化显著缩短了大多数 ClickBench 查询的执行时间。对于最耗时的 Q23 查询，执行时间减少了近四倍，降至 8.6 秒。

然而，我们仍然希望进一步提升性能，以便在 Python 的 GIL 和垃圾回收机制（GC）限制下，保持 ClickHouse 的高效性能。我们通过以下方法实现了这一点：

1. 尽量减少 CPython API 函数调用次数。当调用不可避免时，我们集中处理这些调用，避免在数据管道启动后再调用 CPython 函数。

2. 批量进行数据复制，充分利用 ClickHouse 的 SIMD 优化 memcpy。

3. 将 Python 字符串的编码与解码逻辑重写为 C++ 实现。

最后一点可能较为复杂，下面我来详细解释：

由于历史原因，Python 的字符串（str）数据结构相当复杂。它可以以 UTF-8、UTF-16、UTF-32 或更罕见的编码格式存储。当用户对字符串进行操作时，Python 的运行时可能会将其转换为 UTF-8 并缓存到字符串结构中。

这导致了多种需要处理的情况。如前所述，调用任何 CPython 函数前都需要获取 GIL。如果使用 Python 内部机制进行 UTF-8 转换，执行时只能单线程处理。为了解决这一问题，我们选择在 C++ 中重新实现字符串编码转换逻辑。

我们的努力使得 Q23 查询的性能又有了显著提升，执行时间从 8.6 秒缩短到 0.56 秒，提升了 15 倍！

解决了这些问题后，我们决定对比 chDB 与当前流行的内嵌分析数据库 DuckDB 的性能表现。

下图展示了在查询包含 1000 万行 ClickBench 数据的 DataFrame 时，chDB 和 DuckDB 的性能对比：

注意：以上基准测试数据是在配置为 EPYC 9654 + 128G + HP FX900 4TB NVMe 的硬件上测试的，使用了 3000 万行 ClickBench 数据。相关代码可参考 pd_zerocopy.ipynb【https://github.com/chdb-io/chdb/blob/readPyObj/tests/pd_zerocopy.ipynb】。

总结图表：

尽管 chDB v2 在性能上实现了 87 倍的提升，这无疑是一个令人惊叹的成就，但仅靠这种方式，仍难以满足日益多样化和复杂的 Python 数据查询需求。

因此，我开始着手开发一种机制，让用户能够自行定义表的返回逻辑。这样一来，chDB 用户便可以在享受 ClickHouse 高性能的同时，结合 Python 的灵活性来查询任意数据集。

经过数周的开发，我们推出了 chdb.PyReader。通过继承此类并实现其中的 read 函数，用户可以用 Python 自定义 ClickHouse 表函数返回的数据。示例如下：

import chdb
class MyReader(chdb.PyReader):    def __init__(self, data):        # some basic init
    def read(self, col_names, count):        # return col_names*count block

然后，我们就可以像这样调用新的 reader：

# Initialize reader with sample datareader = MyReader({    "a": [1, 2, 3, 4, 5, 6],    "b": ["tom", "jerry", "auxten", "tom", "jerry", "auxten"],})
# Execute a query on the Python reader and display resultschdb.query("SELECT b, sum(a) FROM Python('reader') GROUP BY b").show()

借助 chdb.PyReader，您可以使用 Python 自定义数据返回逻辑。为了展示这一点，我制作了一个使用 SQL 查询 Google Calendar 的示例。只需运行以下命令：

python google_cal.py \  "SELECT summary, organizer_email, parseDateTimeBestEffortOrNull(start_dateTime) WHERE status = 'confirmed';"

您便可以轻松获取所有已接受的会议邀请：

通过 chDB 提供的 API，您可以将返回 JSON 数组的众多 API 转换为 ClickHouse 表，从而直接运行 SQL 查询，而无需存储额外数据或手动定义表结构。这极大地简化了处理 API 数据的过程。

以上功能已在 chDB v2.0.2 及更高版本中实现，您可以通过以下命令进行安装：

pip install "chdb>=2.0.2"

如果您有兴趣利用 chDB 构建自己的应用程序，欢迎加入我们的 Discord 社区！别忘了在 GitHub 上为 chDB 点个星，并查阅 chDB 文档以了解更多信息【https://clickhouse.com/docs/en/chdb】。