Apache Arrow 构建简单的UDF,用于DataSet Scan

环境配置方面请参考下面的博客。
https://www.jianshu.com/p/6f6b5c88acc4
本例主要演示在Apache Arrow环境中构建自定义的UDF,进行自定义运算的过程。
本例定义一个三元的求和函数,求三列之和。
代码如下,
conanfile.txt

[requires]
boost/1.72.0
arrow/15.0.0

[generators]
cmake

CMakeLists.txt

cmake_minimum_required(VERSION 3.3)

project(9_udf_in_dataset)

set(ENV{PKG_CONFIG_PATH} "$ENV{PKG_CONFIG_PATH}:/usr/local/lib/pkgconfig/")

set ( CMAKE_CXX_FLAGS "-pthread")
set(CMAKE_CXX_STANDARD 17)
add_definitions(-g)

include(${CMAKE_BINARY_DIR}/conanbuildinfo.cmake)
conan_basic_setup()

include_directories(${INCLUDE_DIRS})
LINK_DIRECTORIES(${LINK_DIRS})

file( GLOB main_file_list ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp) 
file( GLOB sources ${CMAKE_CURRENT_SOURCE_DIR}/*.cc)

foreach( main_file ${main_file_list} )
    file(RELATIVE_PATH filename ${CMAKE_CURRENT_SOURCE_DIR} ${main_file})
    string(REPLACE ".cpp" "" file ${filename})
    add_executable(${file}  ${main_file} ${sources})
    target_link_libraries(${file}  ${CONAN_LIBS_ARROW} ${CONAN_LIBS} pthread)
endforeach( main_file ${main_file_list})

udf_in_dataset.cpp

#include <arrow/api.h>
#include <arrow/compute/api.h>
#include <arrow/csv/api.h>
#include "arrow/acero/exec_plan.h"
#include "arrow/compute/expression.h"

#include <arrow/dataset/dataset.h>
#include <arrow/dataset/plan.h>
#include <arrow/dataset/scanner.h>

#include <arrow/io/interfaces.h>
#include <arrow/io/memory.h>
#include <arrow/io/stdio.h>

#include <arrow/filesystem/filesystem.h>

#include <arrow/result.h>
#include <arrow/status.h>

#include <arrow/util/vector.h>

#include <iostream>
#include <vector>


namespace ds = arrow::dataset;
namespace cp = arrow::compute;


char ch_csv_data[] = R"csv(k1,k2,k3
1,4,7
2,5,8
11,20,21
3,6,9)csv";

cp::FunctionDoc func_doc{
    "User-defined-function usage to demonstrate registering an out-of-tree function",
    "returns x + y + z",
    {"x", "y", "z"},
    "UDFOptions"
};

arrow::Result<std::shared_ptr<ds::Dataset>> CreateDatasetFromCsvData() {
    arrow::io::IOContext const& io_context = arrow::io::default_io_context();
    std::shared_ptr<arrow::io::InputStream> input;
    std::string csv_data = ch_csv_data;
    std::cout << csv_data << std::endl;

    std::string_view sv = csv_data;
    input = arrow::io::BufferReader::FromString(std::string(sv));
    auto read_options = arrow::csv::ReadOptions::Defaults();
    auto parse_options = arrow::csv::ParseOptions::Defaults();
    auto convert_options = arrow::csv::ConvertOptions::Defaults();

    ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::csv::TableReader> table_reader,
        arrow::csv::TableReader::Make(io_context, input,
            read_options, parse_options, convert_options)); 
    ARROW_ASSIGN_OR_RAISE(auto maybe_table, table_reader->Read());
    auto ds_ = std::make_shared<ds::InMemoryDataset>(maybe_table);
    arrow::Result<std::shared_ptr<ds::InMemoryDataset>> result(std::move(ds_));
    return result;
}

arrow::Status SampleFunction(cp::KernelContext* ctx, cp::ExecSpan const& batch,
    cp::ExecResult* out) {
    
    // return x + y + z
    const int64_t* x = batch[0].array.GetValues<int64_t>(1);
    const int64_t* y = batch[1].array.GetValues<int64_t>(1);
    const int64_t* z = batch[2].array.GetValues<int64_t>(1);
    int64_t* out_values = out->array_span_mutable()->GetValues<int64_t>(1);
    for (int64_t i = 0; i < batch.length; ++i) {
        *out_values++ = *x++ + *y++ + *z++;
    }
    return arrow::Status::OK();
}

arrow::Status UDFDatasetScan() {
    ARROW_ASSIGN_OR_RAISE(auto data_set, CreateDatasetFromCsvData());
    ARROW_ASSIGN_OR_RAISE(auto scan_builder, data_set->NewScan());

    // Customized add three function
    std::string const name = "add_three";
    auto func = std::make_shared<cp::ScalarFunction>(name, cp::Arity::Ternary(), func_doc);
    cp::ScalarKernel kernel({arrow::int64(), arrow::int64(), arrow::int64()},
        arrow::int64(), SampleFunction);

    kernel.mem_allocation = cp::MemAllocation::PREALLOCATE;
    kernel.null_handling = cp::NullHandling::INTERSECTION;
    ARROW_RETURN_NOT_OK(func->AddKernel(std::move(kernel)));

    auto registry = cp::GetFunctionRegistry();
    ARROW_RETURN_NOT_OK(registry->AddFunction(std::move(func)));

    // Start to project the result dataset
    ARROW_RETURN_NOT_OK(scan_builder->Project({
        cp::field_ref("k1"),
        cp::field_ref("k2"),
        cp::field_ref("k3"),
        cp::call("add_three", {cp::field_ref("k1"), cp::field_ref("k2"), cp::field_ref("k3")})
    }, {"k1", "k2", "k3", "sum"}));

    ARROW_ASSIGN_OR_RAISE(auto scanner, scan_builder->Finish());
    ARROW_ASSIGN_OR_RAISE(auto table, scanner->ToTable());
    std::cout << "Result " << table->num_rows() << " rows" << std::endl;
    std::cout << table->ToString() << std::endl;
    return arrow::Status::OK();
} 


int main(int argc, char* argv[]) {
    arrow::Status st = UDFDatasetScan();
    if(!st.ok()) {
        std::cerr << "Error occurred: " << st.message() << std::endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}

程序输出如下,


image.png

image.png
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 159,835评论 4 364
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 67,598评论 1 295
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 109,569评论 0 244
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 44,159评论 0 213
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 52,533评论 3 287
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,710评论 1 222
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,923评论 2 313
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,674评论 0 203
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 34,421评论 1 246
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,622评论 2 245
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 32,115评论 1 260
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,428评论 2 254
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 33,114评论 3 238
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 26,097评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,875评论 0 197
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,753评论 2 276
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,649评论 2 271

推荐阅读更多精彩内容