Skip to content

Commit 0ba29b3

Browse files
0ax10ax1
authored
feat: duckdb vortex extension scaffolding (#2719)
Upfront scaffolding before integrating multi-file-reader scan functionality.
1 parent 52ffd0f commit 0ba29b3

File tree

1 file changed

+102
-34
lines changed

1 file changed

+102
-34
lines changed

‎duckdb-vortex/src/vortex_duckdb_extension.cpp‎

Lines changed: 102 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,52 +1,70 @@
11
#define DUCKDB_EXTENSION_MAIN
22

3-
#include "vortex_duckdb_extension.hpp"
43
#include "duckdb/common/exception.hpp"
4+
#include "duckdb/common/helper.hpp"
5+
#include "duckdb/common/multi_file_reader.hpp"
6+
#include "duckdb/common/multi_file_reader_function.hpp"
57
#include "duckdb/function/scalar_function.hpp"
8+
#include "duckdb/function/table_function.hpp"
69
#include "duckdb/main/extension_util.hpp"
10+
#include "duckdb/parser/parsed_data/create_table_function_info.hpp"
11+
12+
#include "vortex_duckdb_extension.hpp"
713

814
extern "C" {
915
const char *vortex_duckdb_hello();
1016
}
1117

12-
namespace duckdb {
18+
#ifndef DUCKDB_EXTENSION_MAIN
19+
#error DUCKDB_EXTENSION_MAIN not defined
20+
#endif
1321

14-
inline void VortexDuckdbScalarFun(DataChunk &args, ExpressionState &state, Vector &result) {
15-
auto &name_vector = args.data[0];
16-
UnaryExecutor::Execute<string_t, string_t>(name_vector, result, args.size(), [&](string_t name) {
17-
auto c_str = vortex_duckdb_hello();
18-
auto str = std::string(c_str) + name.GetString() + " 🐥";
19-
return StringVector::AddString(result, str);
20-
});
21-
}
22+
class VortexFileReaderOptions : public duckdb::BaseFileReaderOptions {
23+
public:
24+
bool file_reader_flag {false};
25+
};
2226

23-
static void LoadInternal(DatabaseInstance &instance) {
24-
// Register a dummy table function
25-
auto vortex_duckdb_scalar_function =
26-
ScalarFunction("vortex_hello", {LogicalType::VARCHAR}, LogicalType::VARCHAR, VortexDuckdbScalarFun);
27-
ExtensionUtil::RegisterFunction(instance, vortex_duckdb_scalar_function);
28-
}
27+
struct VortexReadBindData : public duckdb::TableFunctionData {
28+
bool read_bind_data_flag {false};
29+
};
2930

30-
void VortexDuckdbExtension::Load(DuckDB &db) {
31-
LoadInternal(*db.instance);
32-
}
31+
struct VortexInitData : public duckdb::GlobalTableFunctionState {
32+
std::atomic<bool> done {false};
3333

34-
std::string VortexDuckdbExtension::Name() {
35-
return "vortex_duckdb";
36-
}
34+
static duckdb::unique_ptr<GlobalTableFunctionState> Create() {
35+
return duckdb::make_uniq<VortexInitData>();
36+
}
37+
};
3738

38-
std::string VortexDuckdbExtension::Version() const {
39-
#ifdef EXT_VERSION_VORTEX_DUCKDB
40-
return EXT_VERSION_VORTEX_DUCKDB;
41-
#else
42-
return "";
43-
#endif
44-
}
39+
// Vortex file info implementation to integrate with DuckDBs MultiFileReader.
40+
struct VortexMultiFileInfo {
41+
static void GetBindInfo(const duckdb::TableFunctionData &bind_data_p, duckdb::BindInfo &info) {
42+
throw std::runtime_error("called vortex get bind info func");
43+
// auto &bind_data = bind_data_p.Cast<VortexReadBindData>();
44+
// info.type = duckdb::ScanType::EXTERNAL;
45+
// info.InsertOption("setting", duckdb::Value::BOOLEAN(bind_data.read_bind_data_flag));
46+
}
4547

46-
} // namespace duckdb
48+
static std::unique_ptr<duckdb::BaseFileReaderOptions>
49+
InitializeOptions(duckdb::ClientContext &context, duckdb::optional_ptr<duckdb::TableFunctionInfo> info) {
50+
return duckdb::make_uniq<VortexFileReaderOptions>();
51+
}
4752

48-
extern "C" {
53+
static duckdb::optional_idx MaxThreads(const duckdb::MultiFileBindData &bind_data,
54+
const duckdb::MultiFileGlobalState &global_state,
55+
duckdb::FileExpandResult expand_result) {
56+
return 1;
57+
}
58+
59+
static std::unique_ptr<duckdb::TableFunctionData>
60+
InitializeBindData(duckdb::MultiFileBindData &multi_file_data,
61+
std::unique_ptr<duckdb::BaseFileReaderOptions> options_p) {
62+
return duckdb::make_uniq<VortexReadBindData>();
63+
}
64+
};
4965

66+
extern "C" {
67+
// Entrypoint to initialize the DuckDB extension calling `VortexDuckdbExtension::Load`.
5068
DUCKDB_EXTENSION_API void vortex_duckdb_init(duckdb::DatabaseInstance &db) {
5169
duckdb::DuckDB db_wrapper(db);
5270
db_wrapper.LoadExtension<duckdb::VortexDuckdbExtension>();
@@ -57,6 +75,56 @@ DUCKDB_EXTENSION_API const char *vortex_duckdb_version() {
5775
}
5876
}
5977

60-
#ifndef DUCKDB_EXTENSION_MAIN
61-
#error DUCKDB_EXTENSION_MAIN not defined
62-
#endif
78+
void duckdb::VortexDuckdbExtension::Load(DuckDB &db) {
79+
DatabaseInstance &instance = *db.instance;
80+
81+
// Example function to showcase an extension function.
82+
auto vortex_hello =
83+
ScalarFunction("vortex_hello", {LogicalType::VARCHAR}, LogicalType::VARCHAR,
84+
[](duckdb::DataChunk &args, duckdb::ExpressionState &state, duckdb::Vector &result) {
85+
auto &name_vector = args.data[0];
86+
87+
duckdb::UnaryExecutor::Execute<duckdb::string_t, duckdb::string_t>(
88+
name_vector, result, args.size(), [&](duckdb::string_t name) {
89+
auto str = std::string(vortex_duckdb_hello()) + name.GetString();
90+
return duckdb::StringVector::AddString(result, str);
91+
});
92+
});
93+
ExtensionUtil::RegisterFunction(instance, vortex_hello);
94+
95+
auto table_func = [](duckdb::ClientContext &context, duckdb::TableFunctionInput &data,
96+
duckdb::DataChunk &output) {
97+
throw std::runtime_error("table_func");
98+
99+
auto &state = data.global_state->Cast<VortexInitData>();
100+
if (state.done.exchange(true)) {
101+
output.SetCardinality(0);
102+
return;
103+
}
104+
};
105+
106+
auto bind_func = [](duckdb::ClientContext &context, duckdb::TableFunctionBindInput &input,
107+
duckdb::vector<duckdb::LogicalType> &return_types,
108+
duckdb::vector<duckdb::string> &names) -> duckdb::unique_ptr<duckdb::FunctionData> {
109+
return duckdb::make_uniq<VortexReadBindData>();
110+
};
111+
112+
auto init_func =
113+
[](duckdb::ClientContext &context,
114+
duckdb::TableFunctionInitInput &input) -> duckdb::unique_ptr<duckdb::GlobalTableFunctionState> {
115+
return duckdb::make_uniq<VortexInitData>();
116+
};
117+
118+
duckdb::CreateTableFunctionInfo table_func_info(
119+
{"vortex_table_func", {duckdb::LogicalType::VARCHAR}, table_func, bind_func, init_func});
120+
121+
duckdb::ExtensionUtil::RegisterFunction(instance, table_func_info);
122+
}
123+
124+
std::string duckdb::VortexDuckdbExtension::Name() {
125+
return "vortex_duckdb";
126+
}
127+
128+
std::string duckdb::VortexDuckdbExtension::Version() const {
129+
return EXT_VERSION_VORTEX_DUCKDB;
130+
}

0 commit comments

Comments
 (0)