Apache Arrow Rust 中快速且内存高效的多列排序,第 2 部分


发布 2022 年 11 月 7 日
作者: tustvold 和 alamb

简介

在这篇文章的第 1 部分中,我们描述了多列排序问题以及高效实现它的挑战。这第二篇文章解释了 Apache Arrow Rust 实现中新的行格式如何工作以及它是如何构建的。

行格式

行格式是通过连接每个列的编码形式创建的可变长度字节序列。每个列的编码取决于其数据类型(和排序选项)。

   ┌─────┐   ┌─────┐   ┌─────┐
   │     │   │     │   │     │
   ├─────┤ ┌ ┼─────┼ ─ ┼─────┼ ┐              ┏━━━━━━━━━━━━━┓
   │     │   │     │   │     │  ─────────────▶┃             ┃
   ├─────┤ └ ┼─────┼ ─ ┼─────┼ ┘              ┗━━━━━━━━━━━━━┛
   │     │   │     │   │     │
   └─────┘   └─────┘   └─────┘
               ...
   ┌─────┐ ┌ ┬─────┬ ─ ┬─────┬ ┐              ┏━━━━━━━━┓
   │     │   │     │   │     │  ─────────────▶┃        ┃
   └─────┘ └ ┴─────┴ ─ ┴─────┴ ┘              ┗━━━━━━━━┛
   Customer    State    Orders
    UInt64      Utf8     F64

          Input Arrays                          Row Format
           (Columns)

编码经过精心设计,无需转义:字节是哨兵(例如 null)的一部分还是值的一部分,永远不会模棱两可。

无符号整数

要编码非空无符号整数,首先写入字节0x01,然后是整数的字节,从最高有效位开始,即大端序。空值编码为0x00字节,后跟整数零值的编码字节。

              ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐
   3          │03│00│00│00│      │01│00│00│00│03│
              └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘
              ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐
  258         │02│01│00│00│      │01│00│00│01│02│
              └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘
              ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐
 23423        │7F│5B│00│00│      │01│00│00│5B│7F│
              └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘
              ┌──┬──┬──┬──┐      ┌──┬──┬──┬──┬──┐
 NULL         │??│??│??│??│      │00│00│00│00│00│
              └──┴──┴──┴──┘      └──┴──┴──┴──┴──┘

             32-bit (4 bytes)        Row Format
 Value        Little Endian

有符号整数

在 Rust 和大多数现代计算机体系结构中,有符号整数使用二进制补码进行编码,其中一个数字通过翻转所有位并加 1 来取反。因此,翻转最高位并将结果视为无符号整数可以保留顺序。然后可以使用上一节中描述的无符号整数的相同编码来编码此无符号整数。例如:

       ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┬──┐
    5  │05│00│00│00│       │05│00│00│80│       │01│80│00│00│05│
       └──┴──┴──┴──┘       └──┴──┴──┴──┘       └──┴──┴──┴──┴──┘
       ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┐       ┌──┬──┬──┬──┬──┐
   -5  │FB│FF│FF│FF│       │FB│FF│FF│7F│       │01│7F│FF│FF│FB│
       └──┴──┴──┴──┘       └──┴──┴──┴──┘       └──┴──┴──┴──┴──┘

 Value  32-bit (4 bytes)    High bit flipped      Row Format
         Little Endian

浮点数

浮点值可以根据IEEE 754 totalOrder 谓词(在 Rust 中由f32::total_cmp实现)进行排序。此排序将浮点值的字节解释为相应大小的有符号小端整数,在负数情况下翻转除符号位以外的所有位。

因此,浮点值通过将其转换为适当大小的有符号整数表示,然后使用上一节中描述的相同的有符号整数编码来编码为行格式。

字节数组(包括字符串)

与上面的原始类型不同,字节数组是可变长度的。对于短字符串,例如我们上面示例中的state,可以将所有值填充到最长值的长度,并使用一些固定值(例如0x00)来生成固定长度的行。这是 DuckDB 博客中描述的用于编码c_birth_country的方法。

然而,字符串列中的值长度差异很大,或者执行开始时不知道最大长度,这使得将字符串填充到固定长度是不可取和/或不切实际的。因此,Rust Arrow 行格式使用可变长度编码。

我们需要一种能够明确终止字节数组末尾的编码。这不仅允许从行格式中恢复原始值,而且确保在与包含较短字节数组的行进行比较时,较长字节数组的字节不会与不同列的字节进行比较。

空字节数组编码为单个0x00字节。类似地,空字节数组编码为单个0x01字节。

要编码非空非空数组,首先写入单个0x02字节。然后以 32 字节块的形式写入数组,每个完整块后跟一个0xFF字节作为继续标记。最后一个块用0x00填充到 32 字节,然后是此最后一个块的未填充长度作为单个字节,代替继续标记。

请注意,以下示例编码为了简洁起见使用了 4 字节的块大小,而不是 32 字节。

                      ┌───┬───┬───┬───┬───┬───┐
 "MEEP"               │02 │'M'│'E'│'E'│'P'│04 │
                      └───┴───┴───┴───┴───┴───┘

                      ┌───┐
 ""                   │01 |
                      └───┘

 NULL                 ┌───┐
                      │00 │
                      └───┘

"Defenestration"      ┌───┬───┬───┬───┬───┬───┐
                      │02 │'D'│'e'│'f'│'e'│FF │
                      └───┼───┼───┼───┼───┼───┤
                          │'n'│'e'│'s'│'t'│FF │
                          ├───┼───┼───┼───┼───┤
                          │'r'│'a'│'t'│'i'│FF │
                          ├───┼───┼───┼───┼───┤
                          │'o'│'n'│00 │00 │02 │
                          └───┴───┴───┴───┴───┘

这种方法松散地受到COBS 编码的启发,并选择它而不是更传统的字节填充,因为它更适合向量化,特别是带有 AVX-256 的硬件可以在一条指令中复制一个 32 字节块。

字典数组

字典编码数据(在 pandas 中称为分类)变得越来越重要,因为它们可以非常有效地存储和处理低基数数据。

编码字典数组的一个简单方法是直接使用前面描述的原始值的编码来编码逻辑值。然而,这会失去字典编码在减少内存和 CPU 消耗方面的优势。

使事情进一步复杂化的是,Arrow 字典编码的实现非常通用,我们无法对字典的内容做出任何假设。特别是,我们不能假设字典值已排序,也不能假设同一字典用于列中的所有数组。

以下示例显示了如何使用两个不同的字典在两个数组中编码字符串列。第一个批次中的字典键012对应于与第二个字典中相同键不同的值。

┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
  ┌───────────┐ ┌─────┐    │
│ │"Fabulous" │ │  0  │
  ├───────────┤ ├─────┤    │
│ │   "Bar"   │ │  2  │
  ├───────────┤ ├─────┤    │       ┌───────────┐
│ │  "Soup"   │ │  2  │            │"Fabulous" │
  └───────────┘ ├─────┤    │       ├───────────┤
│               │  0  │            │  "Soup"   │
                ├─────┤    │       ├───────────┤
│               │  1  │            │  "Soup"   │
                └─────┘    │       ├───────────┤
│                                  │"Fabulous" │
                 Values    │       ├───────────┤
│ Dictionary   (indexes in         │   "Bar"   │
               dictionary) │       ├───────────┤
│                                  │   "ZZ"    │
 ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘       ├───────────┤
┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─        │   "Bar"   │
                           │       ├───────────┤
│ ┌───────────┐ ┌─────┐            │   "ZZ"    │
  │"Fabulous" │ │  1  │    │       ├───────────┤
│ ├───────────┤ ├─────┤            │"Fabulous" │
  │   "ZZ"    │ │  2  │    │       └───────────┘
│ ├───────────┤ ├─────┤
  │   "Bar"   │ │  1  │    │
│ └───────────┘ ├─────┤
                │  0  │    │      Logical column
│               └─────┘               values
                Values     │
│  Dictionary (indexes in
              dictionary)  │
 ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘

允许我们高效地为此类数据创建行格式的关键观察是,给定一个字节数组,总是可以通过添加一个额外的字节来创建在排序顺序中位于其之前或之后的新字节数组。

因此,我们可以逐步构建从字典值到可变长度字节数组的顺序保留映射,而无需事先知道所有可能的字典值,而是在遇到新字典值时引入映射。

┌──────────┐                 ┌─────┐
│  "Bar"   │ ───────────────▶│ 01  │
└──────────┘                 └─────┘
┌──────────┐                 ┌─────┬─────┐
│"Fabulous"│ ───────────────▶│ 01  │ 02  │
└──────────┘                 └─────┴─────┘
┌──────────┐                 ┌─────┐
│  "Soup"  │ ───────────────▶│ 05  │
└──────────┘                 └─────┘
┌──────────┐                 ┌─────┐
│   "ZZ"   │ ───────────────▶│ 07  │
└──────────┘                 └─────┘

    Example Order Preserving Mapping

用于生成此映射的数据结构的详细信息超出了本博客文章的范围,但可能是未来文章的主题。您可以在此处找到代码

数据结构还确保没有值包含0x00,因此我们可以直接使用0x00作为结束分隔符来编码数组。

空值编码为单个0x00字节,非空值编码为单个0x01字节,后跟由顺序保留映射确定的以0x00终止的字节数组。

                          ┌─────┬─────┬─────┬─────┐
   "Fabulous"             │ 01  │ 01  │ 02  │ 00  │
                          └─────┴─────┴─────┴─────┘

                          ┌─────┬─────┬─────┐
   "ZZ"                   │ 01  │ 07  │ 00  │
                          └─────┴─────┴─────┘

                          ┌─────┐
    NULL                  │ 00  │
                          └─────┘

     Input                  Row Format

排序选项

我们迄今为止忽略的一个细节是如何支持升序和降序排序(例如 SQL 中的ASCDESC)。Arrow Rust 行格式通过简单地反转编码表示的字节(除了用于空值编码的初始字节)来支持这些选项,按列进行。

类似地,支持 SQL 兼容排序还需要一种可以指定NULL值顺序的格式(在所有非NULL值之前或之后)。行格式通过可选地将空值编码为0xFF而不是0x00来支持此选项,按列进行。

结论

希望这两篇文章能让您了解可比较行格式的可能之处以及它的工作原理。请随时查看文档以获取入门说明,并在我们的错误跟踪器上报告任何问题。

使用此格式进行词典排序比基于比较器的方法快3 倍以上,对于字符串、字典和大量列的排序,其优势尤其明显。

我们还已经使用它将 DataFusion 项目中保留排序的合并性能翻倍以上,并期望在将其应用于排序、分组、连接和窗口函数运算符时也能获得类似或更大的性能提升。

一如既往,Arrow 社区非常期待看到您用它构建什么!