我们如何让使用 chDB 查询 Pandas DataFrames 的速度提升 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 DataFrame 查询的初始版本实现如下

1. 将内存中的 DataFrame 序列化为 Parquet，写入临时文件或 memfd。我们最初考虑将数据序列化为 Arrow Buffer，但我们的测试表明 Parquet 更快。
2. 修改 SQL 语句中的数据源表，将其替换为 ClickHouse 的 File 表引擎，并传入临时文件的文件描述符。
3. 运行 ClickHouse 引擎，将输出格式设置为 Parquet。
4. 读取 Parquet Buffer 并将其转换为 DataFrame。

此实现导致大部分时间都花在序列化、反序列化和内存复制上。即使使用 memfd，性能仍然不尽如人意。

如下面的图表所示，ClickBench 基准测试中的几乎每个查询都花费了 chDB 超过 30 秒。

2024 年 6 月，chDB v2 引入了 DataFrame 上的 SQL，允许您轻松地在 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 Dict 变量也可以类似的方式作为表进行查询

import chdb
import pandas as pd
import 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 成为世界上最快的 DataFrame SQL 引擎，一次一行代码。

我发现向 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 DataFrames 时，不可避免地要调用 CPython 的一部分（Python 的 C 实现，也是主流的 Python 实现）。这带来了两个重大挑战：GIL 和对象引用计数。

由于 Python GIL（全局解释器锁）的存在，任何 CPython 函数调用都需要先获取 GIL。如果 GIL 的粒度太大，将直接导致 ClickHouse 的多线程引擎在 GIL 的约束下退化为串行执行；如果 GIL 粒度太小，线程之间频繁的锁竞争将会发生，这甚至可能使执行速度比单线程程序更慢。

在 CPython 中，全局解释器锁，或 GIL，是一个互斥锁，用于保护对 Python 对象的访问，防止多个线程同时执行 Python 字节码。GIL 防止了竞争条件并确保了线程安全。关于 Python GIL 如何在这些领域提供帮助的详细解释可以在这里找到。简而言之，这种互斥锁是必要的，主要是因为 CPython 的内存管理不是线程安全的。

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

Python 具有自动垃圾回收机制，这使得用 Python 编写代码非常容易。但是，如果您不小心引用了内存中现有的 Python 对象或创建了一个新对象，则必须控制引用计数；否则，可能会导致双重释放错误或内存泄漏。

因此，对于 chDB 来说，唯一可行的方法是尽可能避免调用 CPython API 函数来读取 CPython 对象。这听起来是不是很疯狂？！

让我简要解释一下，尽管 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. 在 C++ 中重写了 Python 字符串编码和解码逻辑。

最后一点可能有点难以理解，所以让我详细说明一下

由于历史原因，Python 的 str 数据结构极其复杂。它可以存储为 UTF-8、UTF-16、UTF-32 甚至更晦涩的编码。当用户进一步操作 str 时，Python 运行时可能会将它们全部转换为 UTF-8 编码，并可能将它们缓存在 str 结构中。

这导致需要处理许多情况。如前所述，任何对 CPython 函数的调用都需要首先获取 GIL。如果您使用 Python 的内部实现转换为 UTF-8，则一次只能有一个线程工作。因此，出于必要，我们在 C++ 中重新实现了 str 编码转换逻辑。

我们的努力直接导致了 Q23 的另一次重大性能飞跃，将时间从 8.6 秒减少到 0.56 秒，提高了 15 倍！????????????

既然我们已经解决了这些问题，我们认为应该看看 chDB 与 DuckDB（一种流行的进程内分析数据库）相比表现如何。

下图显示了当查询包含 1000 万行 ClickBench 数据的 DataFrames 时，这两个数据库的性能

注：1. 以上基准测试数据是在配备 EPYC 9654 + 128G + HP FX900 4TB NVMe 的硬件上，使用 3000 万行 ClickBench 数据进行测试的。相关代码：pd_zerocopy.ipynb。

总结图表

当然，性能不是数据库中的一切。chDB v2 的 87 倍性能提升非常酷，但这种方法并未涵盖多样化且复杂的 Python 数据查询需求。

因此，我开始考虑创建一种机制，允许用户定义自己的表返回逻辑。这样，chDB 用户可以将 Python 的灵活性与 ClickHouse 的高性能相结合，以查询他们喜欢的任何数据集。

经过几周的开发，我们现在有了 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 data
reader = MyReader({
    "a": [1, 2, 3, 4, 5, 6],
    "b": ["tom", "jerry", "auxten", "tom", "jerry", "auxten"],
})

# Execute a query on the Python reader and display results
chdb.query("SELECT b, sum(a) FROM Python('reader') GROUP BY b").show()

借助 chdb.PyReader，您可以使用 Python 定义自己的数据返回逻辑。所以我做了一个 使用 SQL 查询 Google 日历的演示。通过运行

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

您可以检索所有您接受的会议邀请

通过 chDB 提供的 API，您可以轻松地将许多返回 JSON 数组的 API 视为 ClickHouse 表来运行 SQL 查询。您无需存储额外数据或手动定义表结构。

以上所有功能都在 chDB v2.0.2 及更高版本中提供，您可以通过运行以下命令安装

pip install "chdb>=2.0.2"

如果您有兴趣使用 chDB 构建自己的应用程序，我们欢迎您加入我们的 Discord。不要忘记给 chDB 一个 GitHub 星标。并查看 chDB 文档。