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


已发布 2022 年 11 月 07 日
作者 tustvold 和 alamb

介绍

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

行格式

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

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

          Input Arrays                          Row Format
           (Columns)

该编码经过精心设计,无需转义:区分字节是 sentinel(例如 null)的一部分还是值,永远不会含糊不清。

无符号整数

要编码一个非空无符号整数,首先写入字节 0x01,然后是整数的字节,从最高有效位开始,即大端序。null 被编码为 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 值之前或之后)。行格式通过选择性地将 null 编码为 0xFF 而不是 0x00 来支持此选项,每个列都进行编码。

结论

希望这两篇文章能让你了解可比行格式的可能性以及它的工作原理。随时查看文档以获取入门说明,并在我们的 bugtracker 上报告任何问题。

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

我们已经使用它将 DataFusion 项目中保留排序的合并性能提高了一倍以上。我们预计将其应用于排序、分组、连接和窗口函数运算符时,性能提升幅度会相似或更大。

与往常一样,Arrow 社区 非常期待看到您用它构建的成果!