email-tracker/external/duckdb/tools/juliapkg/src/table_function.jl

#=
//===--------------------------------------------------------------------===//
// Table Function Bind
//===--------------------------------------------------------------------===//
=#
struct BindInfo
    handle::duckdb_bind_info
    main_function::Any

    function BindInfo(handle::duckdb_bind_info, main_function)
        result = new(handle, main_function)
        return result
    end
end

mutable struct InfoWrapper
    main_function::Any
    info::Any

    function InfoWrapper(main_function, info)
        return new(main_function, info)
    end
end

function parameter_count(bind_info::BindInfo)
    return duckdb_bind_get_parameter_count(bind_info.handle)
end

function get_parameter(bind_info::BindInfo, index::Int64)
    return Value(duckdb_bind_get_parameter(bind_info.handle, index))
end

function set_stats_cardinality(bind_info::BindInfo, cardinality::UInt64, is_exact::Bool)
    duckdb_bind_set_cardinality(bind_info.handle, cardinality, is_exact)
    return
end

function add_result_column(bind_info::BindInfo, name::AbstractString, type::DataType)
    return add_result_column(bind_info, name, create_logical_type(type))
end

function add_result_column(bind_info::BindInfo, name::AbstractString, type::LogicalType)
    return duckdb_bind_add_result_column(bind_info.handle, name, type.handle)
end

function get_extra_data(bind_info::BindInfo)
    return bind_info.main_function.extra_data
end

function _add_global_object(main_function, object)
    begin
        lock(main_function.global_lock)
        try
            push!(main_function.global_objects, object)
        finally
            unlock(main_function.global_lock)
        end
    end
    return
end

function _remove_global_object(main_function, object)
    begin
        lock(main_function.global_lock)
        try
            delete!(main_function.global_objects, object)
        finally
            unlock(main_function.global_lock)
        end
    end
    return
end

function _table_bind_cleanup(data::Ptr{Cvoid})
    info::InfoWrapper = unsafe_pointer_to_objref(data)
    _remove_global_object(info.main_function, info)
    return
end

function get_exception_info()
    error = ""
    for (exc, bt) in current_exceptions()
        error = string(error, sprint(showerror, exc, bt))
    end
    return error
end

function _table_bind_function(info::duckdb_bind_info)
    try
        main_function = unsafe_pointer_to_objref(duckdb_bind_get_extra_info(info))
        binfo = BindInfo(info, main_function)
        bind_data = InfoWrapper(main_function, main_function.bind_func(binfo))
        bind_data_pointer = pointer_from_objref(bind_data)
        _add_global_object(main_function, bind_data)
        duckdb_bind_set_bind_data(info, bind_data_pointer, @cfunction(_table_bind_cleanup, Cvoid, (Ptr{Cvoid},)))
    catch
        duckdb_bind_set_error(info, get_exception_info())
        return
    end
    return
end

#=
//===--------------------------------------------------------------------===//
// Table Function Init
//===--------------------------------------------------------------------===//
=#
struct InitInfo
    handle::duckdb_init_info
    main_function::Any

    function InitInfo(handle::duckdb_init_info, main_function)
        result = new(handle, main_function)
        return result
    end
end

function _table_init_function_generic(info::duckdb_init_info, init_fun::Function)
    try
        main_function = unsafe_pointer_to_objref(duckdb_init_get_extra_info(info))
        binfo = InitInfo(info, main_function)
        init_data = InfoWrapper(main_function, init_fun(binfo))
        init_data_pointer = pointer_from_objref(init_data)
        _add_global_object(main_function, init_data)
        duckdb_init_set_init_data(info, init_data_pointer, @cfunction(_table_bind_cleanup, Cvoid, (Ptr{Cvoid},)))
    catch
        duckdb_init_set_error(info, get_exception_info())
        return
    end
    return
end

function _table_init_function(info::duckdb_init_info)
    main_function = unsafe_pointer_to_objref(duckdb_init_get_extra_info(info))
    return _table_init_function_generic(info, main_function.init_func)
end

function _table_local_init_function(info::duckdb_init_info)
    main_function = unsafe_pointer_to_objref(duckdb_init_get_extra_info(info))
    return _table_init_function_generic(info, main_function.init_local_func)
end

function get_bind_info(info::InitInfo, ::Type{T})::T where {T}
    return unsafe_pointer_to_objref(duckdb_init_get_bind_data(info.handle)).info
end

function get_extra_data(info::InitInfo)
    return info.main_function.extra_data
end

function set_max_threads(info::InitInfo, max_threads)
    return duckdb_init_set_max_threads(info.handle, max_threads)
end

function get_projected_columns(info::InitInfo)::Vector{Int64}
    result::Vector{Int64} = Vector()
    column_count = duckdb_init_get_column_count(info.handle)
    for i in 1:column_count
        push!(result, duckdb_init_get_column_index(info.handle, i))
    end
    return result
end

function _empty_init_info(info::DuckDB.InitInfo)
    return missing
end

#=
//===--------------------------------------------------------------------===//
// Main Table Function
//===--------------------------------------------------------------------===//
=#
struct FunctionInfo
    handle::duckdb_function_info
    main_function::Any

    function FunctionInfo(handle::duckdb_function_info, main_function)
        result = new(handle, main_function)
        return result
    end
end

function get_bind_info(info::FunctionInfo, ::Type{T})::T where {T}
    return unsafe_pointer_to_objref(duckdb_function_get_bind_data(info.handle)).info
end

function get_init_info(info::FunctionInfo, ::Type{T})::T where {T}
    return unsafe_pointer_to_objref(duckdb_function_get_init_data(info.handle)).info
end

function get_local_info(info::FunctionInfo, ::Type{T})::T where {T}
    return unsafe_pointer_to_objref(duckdb_function_get_local_init_data(info.handle)).info
end

function _table_main_function(info::duckdb_function_info, chunk::duckdb_data_chunk)
    main_function::TableFunction = unsafe_pointer_to_objref(duckdb_function_get_extra_info(info))
    binfo::FunctionInfo = FunctionInfo(info, main_function)
    try
        main_function.main_func(binfo, DataChunk(chunk, false))
    catch
        duckdb_function_set_error(info, get_exception_info())
    end
    return
end

#=
//===--------------------------------------------------------------------===//
// Table Function
//===--------------------------------------------------------------------===//
=#
"""
DuckDB table function
"""
mutable struct TableFunction
    handle::duckdb_table_function
    bind_func::Function
    init_func::Function
    init_local_func::Function
    main_func::Function
    extra_data::Any
    global_objects::Set{Any}
    global_lock::ReentrantLock

    function TableFunction(
        name::AbstractString,
        parameters::Vector{LogicalType},
        bind_func::Function,
        init_func::Function,
        init_local_func::Function,
        main_func::Function,
        extra_data::Any,
        projection_pushdown::Bool
    )
        handle = duckdb_create_table_function()
        duckdb_table_function_set_name(handle, name)
        for param in parameters
            duckdb_table_function_add_parameter(handle, param.handle)
        end
        result = new(handle, bind_func, init_func, init_local_func, main_func, extra_data, Set(), ReentrantLock())
        finalizer(_destroy_table_function, result)

        duckdb_table_function_set_extra_info(handle, pointer_from_objref(result), C_NULL)
        duckdb_table_function_set_bind(handle, @cfunction(_table_bind_function, Cvoid, (duckdb_bind_info,)))
        duckdb_table_function_set_init(handle, @cfunction(_table_init_function, Cvoid, (duckdb_init_info,)))
        duckdb_table_function_set_local_init(handle, @cfunction(_table_local_init_function, Cvoid, (duckdb_init_info,)))
        duckdb_table_function_set_function(
            handle,
            @cfunction(_table_main_function, Cvoid, (duckdb_function_info, duckdb_data_chunk))
        )
        duckdb_table_function_supports_projection_pushdown(handle, projection_pushdown)

        return result
    end
end

function _destroy_table_function(func::TableFunction)
    # disconnect from DB
    if func.handle != C_NULL
        duckdb_destroy_table_function(func.handle)
    end
    return func.handle = C_NULL
end

function create_table_function(
    con::Connection,
    name::AbstractString,
    parameters::Vector{LogicalType},
    bind_func::Function,
    init_func::Function,
    main_func::Function,
    extra_data::Any = missing,
    projection_pushdown::Bool = false,
    init_local_func::Union{Missing, Function} = missing
)
    if init_local_func === missing
        init_local_func = _empty_init_info
    end
    fun = TableFunction(
        name,
        parameters,
        bind_func,
        init_func,
        init_local_func,
        main_func,
        extra_data,
        projection_pushdown
    )
    if duckdb_register_table_function(con.handle, fun.handle) != DuckDBSuccess
        throw(QueryException(string("Failed to register table function \"", name, "\"")))
    end
    push!(con.db.functions, fun)
    return
end

function create_table_function(
    con::Connection,
    name::AbstractString,
    parameters::Vector{DataType},
    bind_func::Function,
    init_func::Function,
    main_func::Function,
    extra_data::Any = missing,
    projection_pushdown::Bool = false,
    init_local_func::Union{Missing, Function} = missing
)
    parameter_types::Vector{LogicalType} = Vector()
    for parameter_type in parameters
        push!(parameter_types, create_logical_type(parameter_type))
    end
    return create_table_function(
        con,
        name,
        parameter_types,
        bind_func,
        init_func,
        main_func,
        extra_data,
        projection_pushdown,
        init_local_func
    )
end

function create_table_function(
    db::DB,
    name::AbstractString,
    parameters::Vector{LogicalType},
    bind_func::Function,
    init_func::Function,
    main_func::Function,
    extra_data::Any = missing,
    projection_pushdown::Bool = false,
    init_local_func::Union{Missing, Function} = missing
)
    return create_table_function(
        db.main_connection,
        name,
        parameters,
        bind_func,
        init_func,
        main_func,
        extra_data,
        projection_pushdown,
        init_local_func
    )
end

function create_table_function(
    db::DB,
    name::AbstractString,
    parameters::Vector{DataType},
    bind_func::Function,
    init_func::Function,
    main_func::Function,
    extra_data::Any = missing,
    projection_pushdown::Bool = false,
    init_local_func::Union{Missing, Function} = missing
)
    return create_table_function(
        db.main_connection,
        name,
        parameters,
        bind_func,
        init_func,
        main_func,
        extra_data,
        projection_pushdown,
        init_local_func
    )
end