devbox
IT小白
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

Apache Arrow kernel设计与实现

作者：IT小白 | 2024-02-18 03:30:31

赞

踩

apache arrowinternal

Arrow kernel设计与实现

目录

  • 1.什么是kernel

  • 2.Kernel盘点

  • 3.执行计算流

  • 4.如何实现Kernel执行三步曲？

1.什么是kernel

最近在实现几个自定义的Arrow Kernel，需要非常了解Arrow源码，里面的代码设计是什么样子，怎么快速上手，如何高度自定制开发，里面有哪些比较有意思的设计等等问题，将在本文一一展开。

Arrow是一个列存格式，同时也有一个streaming execution engine，整个执行器的Plan节点，可以拆分为Source、Sink、Agg、Proj、Filter等等。

以Agg为例，对于PostgreSQL来说Agg可以分为两阶段，PartialAgg与FinalAgg，而在Arrow这边只有一个mean运算，也就是一次聚集。

两阶段聚集的逻辑为：以avg为例

  • 第一阶段

    • 在各个节点进行计算，先计算出各个节点的sum、count

  • 第二阶段

    • 在master节点上收到各个子节点的sum、count，然后每个累加之后，做除法，得到avg。

对于Arrow来说，使用Agg便是往AggOption里面设置一个mean函数，这个函数直接得到的是一个avg结果，跟我们的预期不符，因此需要拆分为：

  • avg_trans（第一阶段）

    • 输入为任意类型，输出为struct<sum, count>

  • avg_final（第二阶段）

    • 输入为struct<sum, count>，输出为avg结果

由于Arrow在每个阶段支持consume->merge->finalize，所以我们可以在第二阶段做merge操作，finalize阶段直接计算avg结果，这样变得非常容易了。

当然，里面还有很多实现细节，诸如：struct如何生成、如何识别struct输入，struct如何进行计算，里面的类型转换是什么样子呢？

在写这篇文章之前，发现自己理解的比较肤浅，写完之后，一梳理全都顺了，好东西都是要分享的，期待大家一起交流，共同进步~

2.Kernel盘点

61575d1c00afc9ee93c0f3a7b48095c8.png

1.值描述符(ValueDescr)

包含数据类别(Shape)与数据类型(DataType)。

  • 数据类别：分为数组、标量值、任意类型(前面两者之一)。

  • 数据类型抽象：DataType为所有类型的基类，例如：定长、嵌套类型等。

对于一个数据，总是有类型的，那么在Arrow中就有一个对应的ValueDesc，用来描述输入数据的类型是什么，由于在向量化中传入的数据可以为嵌套类型、高精度类型等等，而不仅仅包含传统的C类型，所以需要自定义一个ValueDescr用来描述这个结构。

2.输入数据类型(InputType)

允许我们根据输入的数据类型去决定如何提取输入数据的行为。通常比较直接的就是你传递的是什么，输入就是什么，那么这被称为ANY_TYPE，对应ValueDescr的Any；当需要自定义输入时，我们可以选用EXACT_TYPE，例如：我现在输入的是int类型，但是我想将其转为uint，那么便可以选用这种；最后一种情况就是高度自定义，可以自己定义一个类型匹配器，决定输入的类型是否匹配，这种称为USE_TYPE_MATCHER。

3.输出类型(OutputType)

就更加有意思了，它划分为固定输出类型与可被计算的输出类型。

  • 固定输出类型

    • 你给我传递什么，我就用什么作为输出类型。

  • 可被计算输出类型

    • 你给我传递一个解析器，根据你传递的类型(一般为输入类型)去生成一个输出类型。

不得不说这个解析器设计简直牛逼！

举个场景：基于PG的逻辑我们需要做AGG两阶段聚集，根据输入的int类型数据，第一阶段产生的是struct<sum, cout>，普通的输入与输出是一致的，但是对于这种场景完全不一致，所以需要自己编写一个解析器，可以完美的解决这个问题。

4.内核签名(KernelSignature)

内核签名表示内核(kernel)当中计算的数据签名，设计思想为：你给我一堆输入类型，一个输出类型。

⚠️输入类型可以是多个，输出是一个，为何这样设计呢？

举个例子：我现在要计算count(bool) 有一列是bool，我要计算count(bool)，得出的结果是什么？

毋庸置疑是一个int/uint，那么本身是int的输出就还是int，这样需求不就来了，存在多对一关系。

实际实现中，内核签名还需要对输入类型做hash保证唯一性。

5.内核(Kernel)

内核是其他所有计算内核的基类，定义了公共拥有的内容。在Kernel当中定义了内核签名与初始化。

内核签名就不说了，一个内核肯定需要知道当前输入数据与输出数据张什么样子。

内核初始化，需要重点讲一下，每个计算内核都需要做自己的事情，那么如何保证这一点呢？

这便是KernelInit的强大之处!简单来说，Kernel提供了统一的初始化接口，内部InitAll会初始化所有的Kernel，每次初始化会调用对应Kernel的KernelInit函数，签名见后面，每个内核要做自己的事情可以在KernelInit中去做，怎么保证这一点？

便是通过KernelState来实现的。

  1. using KernelInit = std::function<Result<std::unique_ptr<KernelState>>(
  2.     KernelContext*, const KernelInitArgs&)>;

KernelState定义了一个接口，是一个opaque类型。每个计算内核只需要继承KernelState便可以自己定义的KernelState。

  1. struct ARROW_EXPORT KernelState {
  2.   virtual ~KernelState() = default;
  3. };

例如：ScalarAggregator相关：

struct ScalarAggregator : public KernelState { };

这里之所以称为ScalarAggregator，是因为返回值是标量，而不是输入是标量，输入可以是数组/标量的。另外一种普遍的Agg就是group by，这里被称为HashAgg。

在每个Init时，用户可以根据传入的参数(KernelInitArgs)与上下文(KernelContext)来生成自己的KernelState。接下来我们分别讲解这里提到的一些内容。

6.内核参数( KernelInitArgs)

内核参数非常简单，kernel与inputs前面都介绍过了，这里直接说这两个功能：

  • kernel

一般用来提取输出类型，在注册阶段会把输入与输出类型注册到KernelSignal中，我们知道OuputType可以是解析器，那么怎么提取呢？

这里可以通过kernel来获取，例如：

  1. ARROW_ASSIGN_OR_RAISE(auto out_type,
  2.                         args.kernel->signature->out_type().Resolve(ctx, args.inputs));

  • inputs

一堆ValueDesc，表示输入数据类型签名。

  • options

FunctionOptions这个就比较复杂了，后面专门讲解，可以简单理解为存储计算函数，例如现在要计算sum(i)，那么这里便是sum函数相关的一些信息。

  1. struct KernelInitArgs {
  2.   const Kernel* kernel;
  3.   const std::vector<ValueDescr>& inputs;
  4.   const FunctionOptions* options;
  5. };

7.内核上下文( kernelContext)

内核上下文用来执行特定KernelState。

里面包含了执行上下文(ExecContext)与内核状态(KernelState)。

执行上下文跟执行器相关了，后面再详细阐述。

接下来，我们用实际的例子来深入内核。

3.执行计算流

在这里我们将回答几个问题，我们在一个Plan里面添加了节点经常需要设置Option，为何要设置呢？如果不设置会怎么样呢？

答案是如果不设置那就找不到你的计算函数了，例如：avg会写入到option里面，在构建AggNode时，将会从option里面找到Agg函数，option里面的agg可能是多个，所以要一个个遍历，然后拿到每个agg的输入描述结构列表，从Kernel签名会对其进行Match，随后构建KernelContext，构建KernelInitArgs，此时Kernel准备工作就绪，调用我们的Init函数，此时会初始化得到我们想要的State，把这个State丢进context里面，下次执行从context进行获取即可，最后需要通过Kernel签名中的OutputType拿到对应的Resolver，提取出Out类型，构造出Agg节点放入plan中。

以上便是一个完整的添加节点到plan当中的流程。

ce6e65316e8bc46410180d4c9011271b.png

添加完节点，那肯定要执行，所以当执行器运行起来后，会通过前面的KernelContext执行state的三步曲。依次调用consume、merge、finalize。

对应如下图所示:

6bb58cabeb76e2d3b9f46f2cd5773ef7.png

4.如何实现Kernel执行三步曲？

以Sum为例：

  1. Consume阶段，我们需要消费数据，数据可以是一个数组也可以是一个标量，所以在代码中都做了处理，下面呆木进行了本部分删减，只展示核心逻辑。

  1. Status Consume(KernelContext*, const ExecBatch& batch) override {
  2.   if (batch[0].is_array()) {
  3.     const auto& data = batch[0].array();
  4.     this->count += data->length - data->GetNullCount();
  5.     this->sum += SumArray<CType, SumCType, SimdLevel>(*data);
  6.   } else {
  7.     const auto& data = *batch[0].scalar();
  8.     this->count += data.is_valid * batch.length;
  9.     this->sum += internal::UnboxScalar<ArrowType>::Unbox(data) * batch.length;
  10.     }
  11.   }
  12.   return Status::OK();
  13. }

2.Merge

两两累加。

  1. Status MergeFrom(KernelContext*, KernelState&& src) override {
  2.   const auto& other = checked_cast<const ThisType&>(src);
  3.   this->count += other.count;
  4.   this->sum += other.sum;
  5.   this->nulls_observed = this->nulls_observed || other.nulls_observed;
  6.   return Status::OK();
  7. }

3.Finalze

由于Sum是不需要三阶段的，所以上述Merge完就是结果了，这里只是赋值，不做处理，但是对于mean之类的，需要做计算。

  1. Status Finalize(KernelContext*, Datum* out) override {
  2.   if ((!options.skip_nulls && this->nulls_observed) ||
  3.       (this->count < options.min_count)) {
  4.     out->value = std::make_shared<OutputType>(out_type);
  5.   } else {
  6.     out->value = std::make_shared<OutputType>(this->sum, out_type);
  7.   }
  8.   return Status::OK();
  9. }

本节完~

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/104277
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号