KnowForge/Apache Gluten 简体中文
开发者概览

Bolt 函数开发

qianmoQqianmoQ· 更新于 2026-10-02· 阅读 17 分钟· 0 次阅读
登录后可跨设备保存划线和私人笔记登录

Bolt 中实现 Spark SQL 函数的开发者指南

本指南介绍如何在 Bolt 原生引擎中实现并注册 Spark SQL 的内置函数，供 Gluten 使用。理解 Bolt 的函数架构及其与 Gluten 的交互方式，是扩展其能力的关键。

高层支持与命名约定

Bolt 为函数维护两个主要命名空间，位于 bolt/functions/：

  1. prestosql：包含一个丰富的函数库，最初以 Presto 兼容语义实现。这些函数默认注册。
  2. sparksql：包含专门为匹配 Spark SQL 语义而实现或覆盖的函数。

当 Gluten 初始化 Bolt 后端时，它首先注册 prestosql 命名空间中的所有函数，然后注册 sparksql 命名空间中的函数。这种分层机制使得 sparksql 中具有相同名称和签名的函数能够透明地覆盖其 prestosql 中的对应函数。这对于确保与 Spark 的语义兼容性至关重要。

如果所需函数在两个命名空间中都缺失，或者 prestosql 中某个函数的行为与 Spark 不一致（即使签名匹配），则必须向 sparksql 文件夹中添加新的实现。

Gluten 到 Bolt 的函数映射： Gluten 将 Spark 的 Expression 节点翻译为 Substrait 计划。Bolt 随后消费该计划，将 Substrait 函数调用映射到其已注册的原生函数。Spark SQL 中使用的函数名会被转换为小写字符串（例如，BitwiseAnd 转换为 bitwise_and），并以此在 Bolt 的函数注册表中查找相应的原生实现。

具体示例

  • 标量函数：Spark SQL 函数 hash(col1, col2, ...) 映射到 Bolt 中的 hash 函数，其对应的 C++ 实现位于 bolt/functions/sparksql/Hash.cpp。
  • 聚合函数：Spark 中的 SUM() 函数映射到 Bolt 中的 sum 聚合函数。其实现 SumAggregate 可在 bolt/functions/sparksql/aggregates/SumAggregate.cpp 中找到。该实现处理了多种数值类型，并展示了部分聚合与最终聚合的管理方式。

Bolt 函数的结构剖析

从本质上讲，Bolt 函数是一个 C++ 结构体，它实现了一组特定的方法并遵循某些约定。

标量函数（简单函数）

Bolt 中的简单函数每次处理一行数据。它们定义在结构体内部，通常以模板形式编写以处理不同的数据类型。

关键组成部分：

  • BOLT_DEFINE_FUNCTION_TYPES(T)：函数结构体内必需的宏，用于定义必要的类型别名，如 out_type<Varchar> 和 arg_type<Varchar>。
  • call() 方法：非空输入的主执行逻辑。第一个参数始终是输出的非 const 引用（result），随后是输入的 const 引用。void 返回类型表示当输入非空时，函数绝不会产生空输出。
  • callNullable() 方法：如果函数需要显式处理 NULL 输入（即输入为空并不意味着输出为空），则应实现该方法。输入参数以指针形式传递，当输入为空时这些指针将是 nullptr。bool 从 call() 或 callNullable() 返回的值用于控制结果的可空性（true 表示非空，false 表示为空）。
  • is_deterministic 标志：一个静态的 constexpr bool，当函数是非确定性的（例如 rand()）时，可将其设置为 false。默认情况下，函数被视为确定性的。

示例：一个简单的 plus 函数

// From bolt/examples/SimpleFunctions.cpp
template <typename TExecParams>
struct MyPlusFunction {
  FOLLY_ALWAYS_INLINE void
  call(int64_t& out, const int64_t& a, const int64_t& b) {
    out = a + b;
  }
};

聚合函数

聚合函数更为复杂，通常继承自 bolt::exec::Aggregate。它们需要一些方法来管理跨行的累加状态。

  • 累加器（char*）：用于保存中间聚合状态的原始内存缓冲区，其布局由函数自身定义。
  • initializeNewGroups()：为新的一组分组初始化累加器。
  • addRawInput() / addIntermediateResults() / addSingleGroupRawInput()：这些方法定义了从原始输入行或部分聚合结果中累加数据的核心逻辑。
  • extractValues() / extractRawValues()：从累加器中提取最终结果或中间结果，并将其写入输出向量。
  • destroy()：清理累加器所持有的资源（例如复杂对象）。

窗口函数

窗口函数通过继承 bolt::exec::WindowFunction 并实现以下方法来定义：

  • resetPartition()：在引擎开始处理新分区时调用。
  • apply()：针对分区中的一批行调用，用于计算函数的输出。该方法会接收窗口帧边界和对等组边界。

函数注册

标量函数注册

标量函数使用 registerFunction 或 RegistrationHelpers.h 提供的辅助模板进行注册。注册时需要指定函数名称、其 C++ 实现结构体，以及函数签名（返回类型后跟参数类型）。

// From bolt/examples/SimpleFunctions.cpp

// Register a specific signature
registerFunction<MyPlusFunction, int64_t, int64_t, int64_t>({"my_plus"});

// Use helpers for common patterns (e.g., all numeric types)
functions::registerBinaryNumeric<MyPlusTemplatedFunction>({"my_other_plus"});

所有 Spark SQL 函数都必须通过调用特定类别的注册函数（例如 registerStringFunctions(prefix)）在 bolt/functions/sparksql/registration/Register.cpp 中完成注册。

聚合函数注册

聚合函数通过 bolt/functions/sparksql/aggregates/Register.cpp 中的 registerAggregateFunction 进行注册。该操作会将函数名称与其对应的实现类以及一组受支持的签名关联起来。

// From bolt/functions/sparksql/aggregates/SumAggregate.cpp
exec::registerAggregateFunction(
    name,
    std::move(signatures),
    [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr>& argTypes,
        const TypePtr& resultType,
        const core::QueryConfig& config) -> std::unique_ptr<exec::Aggregate> {
      // Factory lambda that creates an instance of the aggregate function
      // based on input/output types and aggregation step (partial, final, etc.).
      ...
      return std::make_unique<SumAggregate<...>>(...);
    },
    ...
);

窗口函数注册

窗口函数使用 registerWindowFunction 在 bolt/functions/sparksql/window/WindowFunctionsRegistration.cpp 中进行注册。

// From bolt/functions/sparksql/window/WindowFunctionsRegistration.cpp
void registerWindowFunctions(const std::string& prefix) {
  functions::window::registerNthValueInteger(prefix + "nth_value");
  functions::window::registerRowNumberInteger(prefix + "row_number");
  // ... and so on
}

语义与注意事项

空值处理

  • 默认行为：如果简单函数的任意输入为 NULL，框架会自动生成 NULL 输出，而不会调用该函数。
  • 自定义行为：若要覆盖此行为，请实现 callNullable(result, arg1, arg2, ...)。参数以指针形式传递，当值为 NULL 时，这些指针将 nullptr。这对于诸如 coalesce 或 is_null 之类的函数非常有用。

ANSI SQL 模式

Gluten 目前不支持 Spark 的 ANSI 模式。 当 spark.sql.ansi.enabled 为 true 时，Gluten 通常会回退到原生 Spark 执行。因此，Bolt 中的函数实现只需考虑 ANSI 模式为 关闭 时的行为。这一点对于诸如 cast 之类的函数尤为重要，对于无效转换（例如字符串转 int）的错误处理不应抛出异常，而应返回 NULL。

确定性

默认情况下，所有函数都被视为确定性的。如果你的函数对相同输入可能产生不同结果（例如 rand()、now()），则必须将其声明为非确定性函数：

template <typename T>
struct MyNonDeterministicFunction {
  static constexpr bool is_deterministic = false;
  // ... implementation ...
};

错误处理

对于可恢复的错误（例如，函数的输入格式无效），Bolt 函数应使用 BOLT_USER_FAIL 或 BOLT_USER_CHECK。该机制允许 Bolt 将定义良好的错误传播回 Gluten 和 Spark，而不是导致进程崩溃。不可恢复的逻辑错误仍应使用 BOLT_CHECK 或 BOLT_FAIL。

BOLT_USER_CHECK(
    isValid(input), "Invalid input format for my_function: {}", input);

性能与内存

  • 向量化：虽然简单函数在概念上是逐行处理的，但框架确保它们能在列式数据上高效执行。请避免在 call 方法中使用分支或复杂逻辑。

  • 内存池：函数可以通过 memory::MemoryPool 进行内存分配。对于字符串结果，StringWriter（out_type<Varchar>）会自动处理缓冲区管理。

  • 字符串优化：

    • reuse_strings_from_arg：对于返回输入子串的函数（例如 substr、split），设置 static constexpr int32_t reuse_strings_from_arg = <arg_index>; 可以使输出成为对输入缓冲区的零拷贝视图，从而显著减少内存流量。
    • ASCII 快速路径：如果某个字符串函数在处理纯 ASCII 输入时可以更高效地实现，请提供一个 callAscii() 方法。引擎会自动检测仅含 ASCII 的输入向量，并分派到该优化路径。

测试

新函数必须附带测试。

  • 位置：新函数（例如 my_function）的单元测试应放在对应的新文件中，例如 bolt/functions/sparksql/tests/MyFunctionTest.cpp。
  • 框架：测试通常继承自 SparkFunctionBaseTest，并使用 evaluateOnce 或 evaluate 以测试数据调用函数并验证结果。
  • 交叉验证：确保 Bolt 实现与 Spark 的行为一致至关重要。最有效的方法是在 Gluten 仓库中添加集成测试，针对同一条 Spark SQL 查询分别在原生 Spark 和 Gluten-Bolt 后端上运行，然后比较结果。

综合示例：一个完整的例子

下面是一个添加新标量函数 my_scalar_add(a, b) 的最小示例，该函数将两个整数相加。

  1. 在 bolt/functions/sparksql/MyScalarAdd.h（新文件）中实现该函数：

    #pragma once

#include "bolt/functions/Udf.h"

namespace bytedance::bolt::functions::sparksql {

template <typename T>
struct MyScalarAddFunction {
  BOLT_DEFINE_FUNCTION_TYPES(T);

  FOLLY_ALWAYS_INLINE void call(
      int64_t& result,
      const int64_t& a,
      const int64_t& b) {
    result = a + b;
  }
};
} // namespace bytedance::bolt::functions::sparksql

  2. 在 bolt/functions/sparksql/registration/RegisterMisc.cpp（或新的注册文件）中注册该函数：

    #include "bolt/functions/sparksql/MyScalarAdd.h"
#include "bolt/functions/lib/RegistrationHelpers.h"

namespace bytedance::bolt::functions::sparksql {

void registerMyFunctions(const std::string& prefix) {
  registerFunction<MyScalarAddFunction, int64_t, int64_t, int64_t>(
      {prefix + "my_scalar_add"});
}

// In registerMiscFunctions (or a new top-level registration function)
void registerMiscFunctions(const std::string& prefix) {
    // ... other registrations
    registerMyFunctions(prefix);
}
}

  3. 在 bolt/functions/sparksql/tests/MyScalarAddTest.cpp 中添加测试：

    #include "bolt/functions/sparksql/tests/SparkFunctionBaseTest.h"

using namespace bytedance::bolt;
using namespace bytedance::bolt::test;

class MyScalarAddTest : public functions::test::SparkFunctionBaseTest {
protected:
  void testAdd(int64_t a, int64_t b, int64_t expected) {
    auto result = evaluateOnce<int64_t>(
        "my_scalar_add(c0, c1)",
        makeRowVector({
            makeConstant((int64_t)a, 1),
            makeConstant((int64_t)b, 1),
        }));
    ASSERT_EQ(result.value(), expected);
  }
};

TEST_F(MyScalarAddTest, basic) {
  testAdd(10, 20, 30);
  testAdd(-5, 5, 0);
}

    这种结构化的方法确保新函数能够被正确实现、注册和测试，从而保持与 Spark SQL 的兼容性以及 Bolt 引擎内的高性能表现。

评论

登录后参与评论

正在加载评论…