我们如何将使用 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 的 DataFrame 查询的初始版本实现如下

1. 将内存中的 DataFrame 序列化为 Parquet，写入临时文件或memfd。我们最初考虑过将数据序列化到 Arrow Buffer，但我们的测试表明 Parquet 速度更快。
2. 修改 SQL 语句中的数据源表，将其替换为 ClickHouse 的 File Table 引擎，并将临时文件的描述符传入。
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 DataFrame 时，不可避免地要调用 CPython 的部分内容（Python 的 C 实现，也是主流的 Python 实现）。这带来了两个重大挑战：GIL 和对象引用计数。

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

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

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

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

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

让我简要解释一下我们为了让 chDB 成为 Pandas DataFrame 上最快的 SQL 引擎之一而做出的努力，尽管 Python 拖了后腿。

在博文开头，我们提到 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 的表现如何，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

然后，我们可以像这样使用我们的新读取器：

# 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 点赞。还可以查看 chDB 文档。