사용자 정의 함수 작성하기

사용자 정의 함수 작성하기 (Write User Defined Functions)

Rust 클라이언트는 Rust로 작성한 사용자 정의 함수를 등록할 수 있어요. vscalar 기능으로 스칼라 함수, vtab 기능으로 테이블 함수를 만듭니다. 둘 다 DuckDB의 컬럼형 데이터 청크를 다루므로 단일 호출이 행 배치를 처리해요.

출처: 문서

본문

개요 (Overview)

Rust 클라이언트는 Rust로 작성한 사용자 정의 함수를 등록할 수 있어요: vscalar 기능으로 스칼라 함수, vtab 기능으로 테이블 함수를 등록해요. 둘 다 DuckDB의 컬럼형 데이터 청크(chunk)에서 동작하므로 단일 호출이 행 배치를 처리해요. 아래 섹션은 각각을 빌드하고 등록한 다음, 로드 가능한 익스텐션으로 패키징하는 템플릿을 안내해요.

스칼라 함수 (Scalar Functions)

스칼라 함수는 VScalar 트레이트를 구현하는 타입으로, register_scalar_function()으로 연결에 등록해요. DuckDB는 데이터 청크마다 함수의 invoke()를 한 번 호출해서 입력 행 배치를 전달하고 해당 출력 값을 써넣길 기대해요. signatures()는 인자와 반환 타입을 선언해요.

다음은 크레이트의 vscalar 예시에서 가져온 add_one(BIGINT) -> BIGINT를 정의해요. 단일 BIGINT 인자를 i64 슬라이스로 읽고 각 결과를 출력 벡터에 씁니다:

use duckdb::{
    Connection, Result,
    core::{DataChunkHandle, LogicalTypeHandle, LogicalTypeId},
    vscalar::{ScalarFunctionSignature, VScalar},
    vtab::arrow::WritableVector,
};

struct AddOne;

impl VScalar for AddOne {
    // 함수별 상태는 필요 없어요.
    type State = ();

    fn invoke(
        _: &Self::State,
        input: &mut DataChunkHandle,
        output: &mut dyn WritableVector,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let len = input.len();
        let src = input.flat_vector(0);
        let src = unsafe { src.as_slice_with_len::<i64>(len) };

        let mut out = output.flat_vector();
        let dst = unsafe { out.as_mut_slice_with_len::<i64>(len) };

        for (d, s) in dst.iter_mut().zip(src) {
            *d = s + 1;
        }
        Ok(())
    }

    fn signatures() -> Vec<ScalarFunctionSignature> {
        vec![ScalarFunctionSignature::exact(
            vec![LogicalTypeHandle::from(LogicalTypeId::Bigint)],
            LogicalTypeHandle::from(LogicalTypeId::Bigint),
        )]
    }
}

fn main() -> Result<()> {
    let conn = Connection::open_in_memory()?;
    conn.register_scalar_function::<AddOne>("add_one")?;

    let mut stmt = conn.prepare("SELECT add_one(i) FROM range(5) t(i)")?;
    let values: Vec<i64> = stmt.query_map([], |row| row.get(0))?.collect::<Result<_>>()?;
    assert_eq!(values, vec![1, 2, 3, 4, 5]);
    Ok(())
}

등록되면 함수는 register_scalar_function()에 넘긴 이름으로 SQL에서 호출돼요:

SELECT add_one(41);

associated type State는 호출 간 공유되는 함수별 상태를 담고, 필요 없으면 ()을 사용해요. 예시는 간결함을 위해 입력을 타입 슬라이스로 읽고 non-null 입력을 가정해요. 프로덕션 함수는 각 값의 유효성도 확인해야 해요. 이 기능은 features = ["vscalar"]가 필요해요.

테이블 함수 (Table Functions)

테이블 함수는 VTab 트레이트를 구현하는 타입으로, register_table_function()으로 등록하고 다른 테이블처럼 쿼리해요. VTab은 작업을 세 콜백으로 나눠요:

  • bind는 문이 준비될 때 한 번 실행돼요. add_result_column()으로 결과 컬럼을 선언하고 get_parameter()로 호출의 인자를 읽어, 읽기 전용 바인드 데이터를 반환해요.
  • init은 실행 전에 한 번 실행되고 커서 같은 스캔별 상태를 반환해요.
  • func는 호출마다 행 청크 하나를 만들어요. 출력 벡터를 채우고 청크 길이를 설정해요. 길이 0은 스캔을 끝내요. DuckDB는 스캔이 끝날 때까지 반복 호출하므로 큰 결과도 여러 호출로 전달돼요.
  • parameters는 함수의 파라미터 타입을 선언해요.

크레이트의 vtab 예시는 0..count의 각 정수마다 한 행을 반환하는 numbers(BIGINT) 함수를 정의해요. 증가하는 복잡도의 컬럼들로 구성돼요. 바인드 데이터, init 상태, 청크별 실행 모델을 보여주는 골격:

use duckdb::{
    Connection, Result,
    core::{DataChunkHandle, LogicalTypeHandle, LogicalTypeId},
    vtab::{BindInfo, InitInfo, TableFunctionInfo, VTab},
};
use std::sync::atomic::{AtomicUsize, Ordering};

struct Numbers;

// 호출 인자에서 한 번 계산; 실행 중에는 읽기 전용.
struct NumbersBind {
    count: usize,
}

// 스캔 진행을 추적: func는 출력 청크마다 한 번 호출.
struct NumbersInit {
    cursor: AtomicUsize,
}

impl VTab for Numbers {
    type BindData = NumbersBind;
    type InitData = NumbersInit;

    fn bind(bind: &BindInfo) -> Result<Self::BindData, Box<dyn std::error::Error>> {
        bind.add_result_column("n", LogicalTypeId::Bigint.into());
        let count = bind.get_parameter(0).to_int64().max(0) as usize;
        Ok(NumbersBind { count })
    }

    fn init(_: &InitInfo) -> Result<Self::InitData, Box<dyn std::error::Error>> {
        Ok(NumbersInit { cursor: AtomicUsize::new(0) })
    }

    fn func(
        func: &TableFunctionInfo<Self>,
        output: &mut DataChunkHandle,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let count = func.get_bind_data().count;
        let cursor = &func.get_init_data().cursor;

        // 청크는 벡터 하나 분량의 행만 담을 수 있어요.
        let capacity = output.flat_vector(0).capacity();
        let start = cursor.fetch_add(capacity, Ordering::Relaxed);
        if start >= count {
            output.set_len(0); // 더 이상 행 없음: 스캔 종료
            return Ok(());
        }
        let rows = capacity.min(count - start);

        let mut v = output.flat_vector(0);
        let slice = unsafe { v.as_mut_slice_with_len::<i64>(rows) };
        for (i, slot) in slice.iter_mut().enumerate() {
            *slot = (start + i) as i64;
        }
        output.set_len(rows);
        Ok(())
    }

    fn parameters() -> Option<Vec<LogicalTypeHandle>> {
        Some(vec![LogicalTypeId::Bigint.into()])
    }
}

fn main() -> Result<()> {
    let conn = Connection::open_in_memory()?;
    conn.register_table_function::<Numbers>("numbers")?;

    let total: i64 = conn.query_row("SELECT count(*) FROM numbers(5000)", [], |row| row.get(0))?;
    assert_eq!(total, 5000);
    Ok(())
}

SELECT * FROM numbers(6);

func는 원자적 fetch_add로 결과의 슬라이스를 차지하므로, DuckDB가 여러 스레드에서 실행해도 스캔이 올바르게 유지돼요. 전체 예시는 DOUBLE, VARCHAR(set_null으로 NULL 처리), LIST<BIGINT>, STRUCT 컬럼도 채우며 각 벡터 종류를 쓰는 방법을 보여줘요. Rust 테이블 함수 인터페이스는 DuckDB의 C API 테이블 함수를 최대한 그대로 따르고 있어요. 이 기능은 features = ["vtab"]가 필요하며, vtab-arrow 기능은 Handle Results에서 다루는 내장 ArrowVTab을 추가해요.

로드 가능한 익스텐션 빌드하기 (Building a Loadable Extension)

위 함수들은 클라이언트 애플리케이션 안에서 Connection에 등록돼요. 같은 VScalar와 VTab 인터페이스를 로드 가능한 DuckDB 익스텐션으로 패키징할 수도 있는데, 어떤 DuckDB 클라이언트든 LOAD할 수 있어요. 여기서는 loadable-extension 기능과 duckdb_entrypoint_c_api 매크로로 익스텐션 진입점을 선언해요. 크레이트의 hello-ext 예시가 hello(name) 테이블 함수를 이렇게 등록해요.

경고 (Warning) 익스텐션을 빌드할 때만 loadable-extension을 활성화하세요. 클라이언트 애플리케이션에서 이 기능은 Connection::open_in_memory() 같은 호출을 익스텐션 컨텍스트 밖에서 실행하기 때문에 panic을 일으켜요.

단순한 cargo build는 로드 가능한 파일을 만들지 않아요: DuckDB는 유효한 메타데이터 footer가 있는 .duckdb_extension으로 끝나는 파일을 요구해요. 공식 extension-template-rs 템플릿에서 시작하세요. 그 make debug가 duckdb -unsigned로 로드되는 .duckdb_extension을 만들어요.

더 알아보기 (Further Reading)

  • Handle Results — 내장 ArrowVTab 테이블 함수와 사용자 정의 함수와 공유하는 Arrow 벡터 API.
  • C API — Rust VTab API가 따르는 테이블 함수 인터페이스.
  • Extensions — 로드 가능한 Rust 익스텐션이 연결되는 DuckDB의 익스텐션 시스템.
  • Run Queries — SQL에서 등록된 함수 호출하기.