给DuckDB c++插件模板添加标量函数模板实现大数四则运算-优快云博客

在前文编写的聚合函数的实现文件mpz.cpp末尾增加统一的GMP运算模板函数以及相关字符串工具函数和注册函数。

    static duckdb::string_t StoreString(const std::string& input) {
        duckdb::string_t result;
        auto* raw = reinterpret_cast<duckdb_string_t*>(&result);
        
        if (input.size() <= 12) {
            raw->value.inlined.length = input.size();
            memcpy(raw->value.inlined.inlined, input.data(), input.size());
        } else {
            raw->value.pointer.length = input.size();
            memcpy(raw->value.pointer.prefix, input.data(), 4);
            raw->value.pointer.ptr = (char*)malloc(input.size());
            memcpy(raw->value.pointer.ptr, input.data(), input.size());
        }
        return result;
    }
    // 统一的GMP运算模板
    template<int Operation>
    static duckdb::string_t gmp_operation_impl(duckdb::string_t a, duckdb::string_t b) {
        
        std::string a_str, b_str;
        try {
            a_str = GetNumericString(a);
            b_str = GetNumericString(b);
        } catch (const std::exception& e) {
            throw std::runtime_error(std::string("Invalid input: ") + e.what());
        }

        mpz_t num1, num2, result;
        mpz_init(num1);
        mpz_init(num2);
        mpz_init(result);
        
        if (mpz_set_str(num1, a_str.c_str(), 10) != 0 ||
            mpz_set_str(num2, b_str.c_str(), 10) != 0) {
            mpz_clear(num1);
            mpz_clear(num2);
            mpz_clear(result);
            throw std::runtime_error("Failed to parse number");
        }
        
        // 根据操作类型执行不同的GMP运算
        switch(Operation) {
            case 0: mpz_add(result, num1, num2); break;
            case 1: mpz_sub(result, num1, num2); break;
            case 2: mpz_mul(result, num1, num2); break;
            case 3: mpz_tdiv_q(result, num1, num2); break; // 整数除法
            case 4: { // 开n次方
                unsigned long n = mpz_get_ui(num2);
                if (n == 0) {
                    mpz_clear(num1);
                    mpz_clear(num2);
                    mpz_clear(result);
                    throw std::runtime_error("Root degree cannot be zero");
                }
                mpz_root(result, num1, n);
                break;
            }
            default:
                throw std::runtime_error("Unknown operation");
        }
        
        char* res_str = mpz_get_str(nullptr, 10, result);
        duckdb::string_t ret = StoreString(std::string(res_str));
        
        mpz_clear(num1);
        mpz_clear(num2);
        mpz_clear(result);
        free(res_str);

        return ret;
    }

    void registerGMPFunctions(duckdb::DatabaseInstance &db_) {

        //auto db_ = std::make_unique<duckdb::DuckDB>(db_path);
        auto conn_ = std::make_unique<duckdb::Connection>(db_);

        
        // 注册GMP运算函数
        conn_->CreateScalarFunction("mpz_add", &gmp_operation_impl<0>);
        conn_->CreateScalarFunction("mpz_sub", &gmp_operation_impl<1>);
        conn_->CreateScalarFunction("mpz_mul", &gmp_operation_impl<2>);
        conn_->CreateScalarFunction("mpz_div", &gmp_operation_impl<3>);
        conn_->CreateScalarFunction("mpz_root", &gmp_operation_impl<4>);
    }

将上面的注册函数调用添加到quack_extension.cpp的static void LoadInternal(DatabaseInstance &instance) 函数最后。

    registerGMPFunctions(instance);

正常编译、添加元数据即可。注意链接gmp动态库。注意将libduckdb.so动态库目录加入LD_LIBRARY_PATH。

export LD_LIBRARY_PATH=/par

g++ -fPIC -shared -o libtest2.so quack_extension.cpp -I /par/duck/src/include -lssl -lcrypto -I include -lduckdb -L /par/duck/build/src -I /par/include  -lgmp
root@6ae32a5ffcde:/par/agg# python3 /par/appendmetadata.py -l libtest2.so -n quack -dv v1.3.0  --duckdb-platform linux_amd64 --extension-version 0.1 --abi-type ""
Creating extension binary:
 - Input file: libtest2.so
 - Output file: quack.duckdb_extension
 - Metadata:
   - FIELD8 (unused)            = EMPTY
   - FIELD7 (unused)            = EMPTY
   - FIELD6 (unused)            = EMPTY
   - FIELD5 (abi_type)          =
   - FIELD4 (extension_version) = 0.1
   - FIELD3 (duckdb_version)    = v1.3.0
   - FIELD2 (duckdb_platform)   = linux_amd64
   - FIELD1 (header signature)  = 4 (special value to identify a duckdb extension)
root@6ae32a5ffcde:/par/agg# /par/duckdb130 -unsigned
DuckDB v1.3.0 (Ossivalis) 71c5c07cdd
Enter ".help" for usage hints.
D load '/par/agg/quack.duckdb_extension';
D select function_name from duckdb_functions() where function_name like 'mpz%';
┌───────────────┐
│ function_name │
│    varchar    │
├───────────────┤
│ mpz_root      │
│ mpz_add       │
│ mpz_div       │
│ mpz_mul       │
│ mpz_sub       │
│ mpz_sum       │
└───────────────┘
D select mpz_add('1234567898901234567890','898901234898901234');
┌─────────────────────────────────────────────────────────┐
│ mpz_add('1234567898901234567890', '898901234898901234') │
│                         varchar                         │
├─────────────────────────────────────────────────────────┤
│ 1235466800136133469124                                  │
└─────────────────────────────────────────────────────────┘

D select mpz_mul('1234567898901234567890','898901234898901234');
┌─────────────────────────────────────────────────────────┐
│ mpz_mul('1234567898901234567890', '898901234898901234') │
│                         varchar                         │
├─────────────────────────────────────────────────────────┤
│ 1109754608888861604978884773459777776260                │
└─────────────────────────────────────────────────────────┘
D select mpz_sub('1234567898901234567890','898901234898901234');
┌─────────────────────────────────────────────────────────┐
│ mpz_sub('1234567898901234567890', '898901234898901234') │
│                         varchar                         │
├─────────────────────────────────────────────────────────┤
│ 1233668997666335666656                                  │
└─────────────────────────────────────────────────────────┘
D select mpz_div('1234567898901234567890','898901234898901234');
┌─────────────────────────────────────────────────────────┐
│ mpz_div('1234567898901234567890', '898901234898901234') │
│                         varchar                         │
├─────────────────────────────────────────────────────────┤
│ 1373                                                    │
└─────────────────────────────────────────────────────────┘

D select mpz_root('1234567898901234567890898901234898901234','2');
┌───────────────────────────────────────────────────────────┐
│ mpz_root('1234567898901234567890898901234898901234', '2') │
│                          varchar                          │
├───────────────────────────────────────────────────────────┤
│ 35136418413111410463                                      │
└───────────────────────────────────────────────────────────┘