如何在Apache Arrow中从多个标量函数返回StructArray

实现方案 #

基于Apache Arrow现有API可以很方便实现你的需求，不需要从头开发时间计算逻辑，以下是两种常用实现方式：

方案1：直接构造StructArray（适合临时使用场景）

无需注册新函数，直接调用内置时间组件函数后拼接为StructArray即可，代码示例如下：

#include <arrow/compute/api.h>
#include <arrow/array/array_struct.h>

arrow::Result<std::shared_ptr<arrow::StructArray>> ExtractDateComponents(const std::shared_ptr<arrow::Array>& date_array) {
  // 调用内置计算函数提取各日期组件
  ARROW_ASSIGN_OR_RAISE(auto year, arrow::compute::CallFunction("year", {date_array}));
  ARROW_ASSIGN_OR_RAISE(auto month, arrow::compute::CallFunction("month", {date_array}));
  ARROW_ASSIGN_OR_RAISE(auto day, arrow::compute::CallFunction("day", {date_array}));

  // 可按需求扩展新增 hour、minute、second、微秒等时间组件
  return arrow::StructArray::Make(
    {year.array(), month.array(), day.array()},
    {"year", "month", "day"}
  );
}

该方案实现简单，无需修改全局函数注册表，按需调整返回的组件即可。

方案2：注册自定义Compute函数（适合全局高频复用场景）

如果需要在多个地方调用该能力，可以将逻辑注册为全局Compute函数，之后和调用内置year/month函数的方式完全一致：

1. 首先定义Kernel执行逻辑和注册逻辑：

#include <arrow/compute/function.h>
#include <arrow/compute/registry.h>

arrow::Status ExtractDateComponentsKernel(arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* out) {
  const auto& input = batch.values[0].array();
  // 提取日期组件
  ARROW_ASSIGN_OR_RAISE(auto year, arrow::compute::CallFunction("year", {input}));
  ARROW_ASSIGN_OR_RAISE(auto month, arrow::compute::CallFunction("month", {input}));
  ARROW_ASSIGN_OR_RAISE(auto day, arrow::compute::CallFunction("day", {input}));
  // 构造返回的StructArray
  ARROW_ASSIGN_OR_RAISE(auto struct_arr, arrow::StructArray::Make(
    {year.array(), month.array(), day.array()},
    {"year", "month", "day"}
  ));
  *out = struct_arr;
  return arrow::Status::OK();
}

// 全局初始化阶段执行一次注册即可
arrow::Status RegisterDateComponentsFunction() {
  auto func = std::make_shared<arrow::compute::ScalarFunction>(
    "extract_date_components",
    arrow::compute::Arity::Unary(),
    nullptr
  );
  // 配置Kernel签名，可按需求扩展支持Date64、Timestamp等更多输入类型
  arrow::compute::ScalarKernel kernel;
  kernel.signature = arrow::compute::KernelSignature::Make(
    {arrow::date32()},
    arrow::struct_({
      arrow::field("year", arrow::int64()),
      arrow::field("month", arrow::int64()),
      arrow::field("day", arrow::int64())
    })
  );
  kernel.exec = ExtractDateComponentsKernel;
  ARROW_RETURN_NOT_OK(func->AddKernel(std::move(kernel)));
  // 注册到全局函数注册表
  return arrow::compute::GetFunctionRegistry()->AddFunction(std::move(func));
}

2. 注册完成后的调用方式：

// 调用方式和内置Compute函数完全一致
ARROW_ASSIGN_OR_RAISE(auto date_components, arrow::compute::CallFunction("extract_date_components", {array}));
// 直接拿到StructArray，无需单独管理多个Datum
auto struct_arr = date_components.struct_array();

如果后续需要新增更细粒度的时间组件，只需要在Kernel逻辑中新增对应内置函数的调用，同时更新Struct的字段定义即可，调用侧代码无需修改。

内容的提问来源于stack exchange，提问作者Will Ayd