如何添加 lambda 函数？

本文假设您熟悉 Presto 的 lambda 函数 prestodb

介绍

Pollux 支持 lambda 函数，用于实现对数组和映射的计算。lambda 函数是高阶函数，其参数本身也是函数。 例如，:func:filter 函数接受一个数组或映射和一个谓词，并返回符合谓词的数组或映射元素的子集。 :func:transform 函数接受一个数组和一个函数，将该函数应用于数组的每个元素，并返回一个结果数组。

以下是在 Presto SQL 中使用过滤和转换函数的示例：

    > select filter(array[1, 2, 3, 4], x -> x % 2 = 0);
     [2, 4]

    > select transform(array[1, 2, 3, 4], x -> x * 2)
     [2, 4, 6, 8]

x -> x % 2 = 0 是一个 Lambda 表达式。它由签名和函数体组成，两者之间用“箭头”分隔。请注意，在 Presto SQL 中，Lambda 的签名 仅包含参数名称。参数类型根据上下文推断， 例如，根据filter的数组参数类型推断。

Lambda 可以使用捕获操作来访问封闭函数作用域内的任何列。假设我们有一个数据集，其中包含一个数组列a和一个整数列b：

a: array(integer)	b: integer
[1, 2, 3, 4]	3
[3, 1, 5, 6, 7]	4

我们可以过滤数组a，只保留大于或等于b的元素：

    > select filter(a, x -> x >= b)

    [3, 4]
    [5, 6, 7]

这里，lambda 表达式中的b是捕获符。Lambda 表达式可以使用 零个或多个捕获符。

此外，还可以将不同的 lambda 表达式应用于数据集中的不同行。例如，我们可以过滤数组a，如果b为偶数，则保留偶数元素；如果b为奇数，则保留奇数元素。

    > select filter(a, if(b % 2 == 0, x -> x % 2 == 0, x % 2 == 1))

    [1, 3]
    [6]

注意：在撰写本文时，Presto 不支持此语法。

函数向量

在 Pollux 中，lambda 函数必须实现为向量函数。这些函数接收的 lambda 输入是 FUNCTION 类型的向量。例如， filter函数接收两个向量：一个 ARRAY 类型的向量和另一个 FUNCTION 类型的向量。

函数类型是一种嵌套类型，其子级包含 lambda 参数类型， 后跟 lambda 返回类型。上述filter函数的 lambda 参数的确切类型为 FUNCTION(INTEGER(), BOOLEAN())。

函数向量使用 FunctionVector 类实现。这些向量存储可调用对象，这些对象以紧凑的形式表示可执行的 lambda 表达式。在大多数情况下， 所有行的 lambda 表达式都相同，但正如我们上面所见，不同的行可以与不同的 lambda 关联。FunctionVector 存储一个不同 lambda 的列表以及每个 lambda 适用的一组行。每个 lambda 表达式都表示为一个 Callable 类型的对象，该对象允许对一组行执行 lambda 表达式。

    class Callable {
      bool hasCapture() const;

      void apply(
          const SelectivityVector& rows,
          BufferPtr wrapCapture,
          exec::EvalCtx& context,
          const std::vector<VectorPtr>& args,
          VectorPtr& result);
    };

Callable 的“apply”方法与 VectorFunction 的apply方法类似，因为它接受一组待求值的行和一个表示输入数据的向量列表。 例如，filter函数使用 Callable::apply 对输入数组元素执行 lambda 表达式。在本例中，rows表示元素向量的行，args 包含一个元素向量。“result”是一个布尔向量，对于符合谓词的元素，结果为“true”，对于不符合谓词的元素，结果为false。

除了rows和args之外，Callable::apply() 方法还接受一个可选的wrapCapture缓冲区参数。如果 lambda 表达式使 用捕获，例如，如果 Callable::hasCapture() 返回 true，则必须指定此参数。“wrapCapture”缓冲区用于将顶层捕获行与数组元素或映射键或值的嵌套行对齐。

考虑filter(a, x -> x >= b)示例。x >= b表达式需要两个输入：x和b。其中，x是数组中共有 9 行的元素，而b是只有 2 行的顶层列。 为了对齐“x”和“b”，我们需要重复“b”的次数，次数与相应数组中的元素数量相同。

如果有多个捕获，则所有捕获都需要以相同的方式对齐，例如，它们的值需要重复的次数与相应数组或映射中的元素数量 相同。Callable::apply() 中的 "wrapCapture" 参数用于指定一个索引缓冲区，该缓冲区可用于将捕获包装到字典向 量中以实现此对齐。Callable 对象已经包含用于捕获的向量，因此无需将它们包含在 "apply()" 方法的 "args" 参数中。

与其他向量不同，FunctionVector 不允许访问单个行的 Callable 对象。相反，它提供了一个迭代器，该迭代器返回 Callable 对象的唯一实例以及它们应用到的一组行。

例如，filter 函数可以像这样迭代不同的 Callable：

    auto it = args[1]->asUnchecked<FunctionVector>()->iterator(rows);
    while (auto entry = it.next()) {
        ... entry.callable is a pointer to Callable ...
        ... entry.rows is the set of rows this Callable applies to ...
    }

大多数情况下，Callable 只有一个实例，但函数实现需要允许多个实例。

FunctionVector::iterator() 方法接受 SelectivityVector 参数，该参数将返回的迭代器限制 为指定行的子集。这些行通常是 lambda 函数执行时所依据的行，例如 VectorFunction::apply() 方法的 rows 参数。

端到端流程

表达式树中的 lambda 函数调用由一个 CallTypedExpr 节点和一个 LambdaTypedExpr 子节点表示。filter(a, x -> x % 2 = 0)应表示如下：

请注意，LambdaTypedExpr 节点没有任何子节点。表示 Lambda 主体的表达式包含在 LambdaTypedExpr 节点内。

此表达式树被编译为可执行表达式树。LambdaTypedExpr 被编译为特殊形式的 LambdaExpr，其中包含 编译后的主体（可执行表达式的实例，例如 std::shared_ptr<Expr>）以及用于捕获的 FieldReference 实例列表。LambdaExpr 的执行结果是一个 FunctionVector。 LambdaExpr::evalSpecialForm() 创建 Callable 实例并将其存储在 FunctionVector 中。

Lambda 函数签名

要指定 lambda 函数的签名，请使用function(argType1, argType2,.., returnType)语法来指定 lambda 参数的类型。以下是filter函数签名的示例：

    // array(T), function(T, boolean) -> array(T)
    return {exec::FunctionSignatureBuilder()
                .typeVariable("T")
                .returnType("array(T)")
                .argumentType("array(T)")
                .argumentType("function(T,boolean)")
                .build()};

Testing

测试框架完全支持评估 lambda 表达式。只需像在 Presto SQL 中一样编写表达式即可：

  auto result = evaluate("filter(a, x -> (x >= b))", data);

上例中，data 应包含一个数组类型的列“a”，以及一个与数组元素类型匹配的列“b”。例如，“a”可以是 一个数组（整数），“b”可以是整数。

唯一需要注意的是，你需要将 lambda 表达式主体放在括号中。x -> (x >= b) 可以正常工作，但x -> x >= b 则不行。