Skip to content

Commit

Permalink
feat: add explain TableFunction for substrait
Browse files Browse the repository at this point in the history
This adds a new table function, `explain_substrait`, as well as a new
execution kernel for the function and binding kernel. This reuses the
`FromSubstraitFunctionData` structure since it contains all the pieces
we would need.
  • Loading branch information
drin committed Oct 29, 2024
1 parent be71387 commit 5d06c53
Showing 1 changed file with 44 additions and 0 deletions.
44 changes: 44 additions & 0 deletions src/substrait_extension.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,27 @@ static unique_ptr<FunctionData> FromSubstraitBindJSON(ClientContext &context, Ta
return SubstraitBind(context, input, return_types, names, true);
}

static unique_ptr<FunctionData> ExplainSubstraitBind(ClientContext &context, TableFunctionBindInput &input,
vector<LogicalType> &return_types, vector<string> &names) {
if (input.inputs[0].IsNull()) {
throw BinderException("explain_substrait cannot be called with a NULL parameter");
}

constexpr bool is_json = false;
string serialized = input.inputs[0].GetValueUnsafe<string>();

// We can reuse the `FromSubstraitFunctionData` struct to return the explain plan
auto result = make_uniq<FromSubstraitFunctionData>();
result->conn = make_uniq<Connection>(*context.db);
result->plan = SubstraitPlanToDuckDBRel(*result->conn, serialized, is_json);

// But, we need to customize the return type to just be a string
return_types.emplace_back(LogicalType::VARCHAR);
names.emplace_back("Explain Plan");

return std::move(result);
}

static void FromSubFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
auto &data = data_p.bind_data->CastNoConst<FromSubstraitFunctionData>();
if (!data.res) {
Expand All @@ -262,6 +283,18 @@ static void FromSubFunction(ClientContext &context, TableFunctionInput &data_p,
output.Move(*result_chunk);
}

static void ExplainSubFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
auto &data = data_p.bind_data->CastNoConst<FromSubstraitFunctionData>();
if (!data.res) {
data.res = data.plan->Explain();
}
auto result_chunk = data.res->Fetch();
if (!result_chunk) {
return;
}
output.Move(*result_chunk);
}

void InitializeGetSubstrait(const Connection &con) {
auto &catalog = Catalog::GetSystemCatalog(*con.context);

Expand Down Expand Up @@ -296,6 +329,16 @@ void InitializeFromSubstrait(const Connection &con) {
catalog.CreateTableFunction(*con.context, from_sub_info);
}

void InitializeExplainSubstrait(const Connection &con) {
auto &catalog = Catalog::GetSystemCatalog(*con.context);

// create the explain_substrait table function that returns a stringified query
// plan from a substrait plan
TableFunction explain_sub_func("explain_substrait", {LogicalType::BLOB}, ExplainSubFunction, ExplainSubstraitBind);
CreateTableFunctionInfo explain_sub_info(explain_sub_func);
catalog.CreateTableFunction(*con.context, explain_sub_info);
}

void InitializeFromSubstraitJSON(const Connection &con) {
auto &catalog = Catalog::GetSystemCatalog(*con.context);

Expand All @@ -316,6 +359,7 @@ void SubstraitExtension::Load(DuckDB &db) {

InitializeFromSubstrait(con);
InitializeFromSubstraitJSON(con);
InitializeExplainSubstrait(con);

con.Commit();
}
Expand Down

0 comments on commit 5d06c53

Please sign in to comment.