Skip to content

Commit 355cb08

Browse files
longbinlailongbinlai
authored
Merge branch 'main' into impl_array_parser
2 parents 52ccfb7 + f49b5e0 commit 355cb08

11 files changed

Lines changed: 246 additions & 228 deletions

File tree

‎.github/workflows/neug-extension-test.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ jobs:
146146
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_sniffer.py -k "parquet"
147147
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_copy_temp.py -k "parquet or httpfs"
148148
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_gds_temp_graph.py
149+
NEUG_RUN_EXTENSION_TESTS=true python3 -m pytest -sv tests/test_call_params.py
149150
150151
- name: Run GDS Tests
151152
run: |

‎extension/test_out_of_tree/test_out_of_tree_extension.cc‎

Lines changed: 113 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,127 @@
1414
* limitations under the License.
1515
*/
1616

17+
#include "neug/common/columns/value_columns.h"
1718
#include "neug/compiler/extension/extension_api.h"
19+
#include "neug/compiler/function/neug_call_function.h"
20+
#include "neug/execution/common/context.h"
21+
#include "neug/execution/common/params_map.h"
22+
#include "neug/utils/exception/exception.h"
23+
24+
namespace {
25+
26+
// String literal, or $param (ParamsMap key without '$').
27+
struct DeferredCallArg {
28+
bool is_param = false;
29+
std::string param_name;
30+
std::string literal;
31+
32+
static DeferredCallArg FromLiteral(std::string value) {
33+
DeferredCallArg arg;
34+
arg.literal = std::move(value);
35+
return arg;
36+
}
37+
38+
static DeferredCallArg FromParam(std::string name) {
39+
DeferredCallArg arg;
40+
arg.is_param = true;
41+
arg.param_name = std::move(name);
42+
return arg;
43+
}
44+
45+
std::string resolve(const neug::execution::ParamsMap& params) const {
46+
if (!is_param) {
47+
return literal;
48+
}
49+
auto it = params.find(param_name);
50+
if (it == params.end() || it->second.IsNull()) {
51+
THROW_INVALID_ARGUMENT_EXCEPTION(
52+
"Missing parameter for TEST_ECHO_PARAM: " + param_name);
53+
}
54+
return it->second.GetValue<std::string>();
55+
}
56+
};
57+
58+
struct EchoParamFuncInput : public neug::function::CallFuncInputBase {
59+
DeferredCallArg arg;
60+
std::string value;
61+
62+
std::unique_ptr<neug::function::CallFuncInputBase> bindParams(
63+
const neug::execution::ParamsMap& params) const override {
64+
auto bound = std::make_unique<EchoParamFuncInput>(*this);
65+
bound->value = arg.resolve(params);
66+
return bound;
67+
}
68+
};
69+
70+
struct TestEchoParamFunction {
71+
static constexpr const char* name = "TEST_ECHO_PARAM";
72+
73+
static neug::function::function_set getFunctionSet() {
74+
auto func = std::make_unique<neug::function::NeugCallFunction>(
75+
name,
76+
neug::function::call_input_types{
77+
neug::common::DataType(neug::common::DataTypeId::kVarchar)},
78+
neug::function::call_output_columns{
79+
{"value",
80+
neug::common::DataType(neug::common::DataTypeId::kVarchar)}});
81+
82+
func->bindFunc =
83+
[](const neug::Schema&, const neug::execution::ContextMeta&,
84+
const ::physical::PhysicalPlan& plan,
85+
int op_idx) -> std::unique_ptr<neug::function::CallFuncInputBase> {
86+
const auto& procedure = plan.plan(op_idx).opr().procedure_call();
87+
if (procedure.query().arguments_size() < 1) {
88+
THROW_INVALID_ARGUMENT_EXCEPTION(
89+
"TEST_ECHO_PARAM requires 1 VARCHAR argument");
90+
}
91+
auto input = std::make_unique<EchoParamFuncInput>();
92+
const auto& arg = procedure.query().arguments(0);
93+
if (arg.has_param()) {
94+
input->arg = DeferredCallArg::FromParam(arg.param().name());
95+
} else if (arg.has_const_() && arg.const_().has_str()) {
96+
input->arg = DeferredCallArg::FromLiteral(arg.const_().str());
97+
} else {
98+
THROW_INVALID_ARGUMENT_EXCEPTION(
99+
"TEST_ECHO_PARAM requires a string constant or $param");
100+
}
101+
return input;
102+
};
103+
104+
func->execFunc = [](const neug::function::CallFuncInputBase& input_base,
105+
neug::IStorageInterface&) {
106+
const auto& input = static_cast<const EchoParamFuncInput&>(input_base);
107+
108+
neug::ValueColumnBuilder<std::string> builder;
109+
builder.reserve(1);
110+
builder.push_back_opt(input.value);
111+
112+
neug::execution::Context ctx;
113+
neug::DataChunk chunk;
114+
chunk.set(0, builder.finish());
115+
ctx.append_chunk(std::move(chunk));
116+
ctx.tag_ids = {0};
117+
return ctx;
118+
};
119+
120+
neug::function::function_set set;
121+
set.push_back(std::move(func));
122+
return set;
123+
}
124+
};
125+
126+
} // namespace
18127

19128
extern "C" {
20129

21130
void Init() {
131+
neug::extension::ExtensionAPI::registerFunction<TestEchoParamFunction>(
132+
neug::catalog::CatalogEntryType::TABLE_FUNCTION_ENTRY);
22133
neug::extension::ExtensionAPI::registerExtension(
23134
neug::extension::ExtensionInfo{
24135
"test_out_of_tree",
25-
"Test extension for validating out-of-tree extension builds."});
136+
"Test extension for out-of-tree builds and CALL $params "
137+
"(TEST_ECHO_PARAM)."});
26138
}
27139

28140
const char* Name() { return "TEST_OUT_OF_TREE"; }

‎include/neug/compiler/function/neug_call_function.h‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
#pragma once
1818

19+
#include <memory>
1920
#include <vector>
2021
#include "neug/common/types.h"
2122
#include "neug/compiler/function/function.h"
@@ -35,8 +36,14 @@ class Context;
3536
namespace function {
3637
struct CallFuncInputBase {
3738
virtual ~CallFuncInputBase() = default;
38-
// Resolve deferred $param args against query ParamsMap before execFunc.
39-
virtual void bindParams(const execution::ParamsMap& /*params*/) {}
39+
// Bind deferred $param args into a per-Eval input.
40+
// Opr stores this object unbound after bindFunc; Eval calls bindParams and
41+
// executes against the returned input. Default: no deferred params — return
42+
// nullptr and Opr will exec with the unbound template (must stay immutable).
43+
virtual std::unique_ptr<CallFuncInputBase> bindParams(
44+
const execution::ParamsMap& /*params*/) const {
45+
return nullptr;
46+
}
4047
};
4148

4249
using call_bind_func_t = std::function<std::unique_ptr<CallFuncInputBase>(

‎proto/stored_procedure.proto‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ message Argument {
2929
oneof value {
3030
common.Value const = 3; // real value
3131
common.Variable var = 4;
32+
// Dynamic query param ($name): resolved from ParamsMap at exec time.
33+
common.DynamicParam param = 5;
3234
}
3335
}
3436

‎src/compiler/gopt/g_query_converter.cpp‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1294,11 +1294,23 @@ void GQueryConvertor::convertProcedureCall(
12941294
const auto& param = params.at(pos);
12951295
auto queryArgPB = std::make_unique<::procedure::Argument>();
12961296
queryArgPB->set_param_ind(pos);
1297-
// Dynamic query params ($name): keep the name (no '$') and resolve from
1298-
// ParamsMap at exec time. Literals / variables keep const / var.
1297+
// $param -> typed DynamicParam for params_type / ParamsMap; literals/vars
1298+
// keep const / var.
12991299
if (param->expressionType == common::ExpressionType::PARAMETER) {
1300-
queryArgPB->set_param_name(
1301-
param->constCast<binder::ParameterExpression>().getName());
1300+
const auto& paramExpr = param->constCast<binder::ParameterExpression>();
1301+
if (pos >= callFunc.parameterTypes.size()) {
1302+
THROW_EXCEPTION_WITH_FILE_LINE("CALL $param index out of signature: " +
1303+
paramExpr.getName());
1304+
}
1305+
auto dynPB = std::make_unique<::common::DynamicParam>();
1306+
dynPB->set_name(paramExpr.getName());
1307+
dynPB->set_index(static_cast<int32_t>(pos));
1308+
// CALL $params keep UNKNOWN in the binder; use the procedure signature
1309+
// type for ParamsMap deserialization.
1310+
dynPB->set_allocated_data_type(
1311+
typeConverter->convertLogicalType(callFunc.parameterTypes[pos].copy())
1312+
.release());
1313+
queryArgPB->set_allocated_param(dynPB.release());
13021314
} else {
13031315
auto paramPB = exprConvertor->convert(*param, {});
13041316
if (paramPB->operators_size() == 0) {

‎src/execution/execute/ops/retrieve/gds_algo.cc‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ GDSAlgoOpr::GDSAlgoOpr(std::unique_ptr<function::CallFuncInputBase> algo_input,
3131
neug::result<neug::execution::Context> GDSAlgoOpr::Eval(
3232
IStorageInterface& graph_interface, const ParamsMap& params,
3333
neug::execution::Context&& ctx, neug::execution::OprTimer* timer) {
34+
(void) ctx;
3435
(void) timer;
3536
if (algo_func_ == nullptr) {
3637
THROW_RUNTIME_ERROR("GDSAlgoOpr: GDSAlgoFunction pointer is null");
@@ -42,8 +43,9 @@ neug::result<neug::execution::Context> GDSAlgoOpr::Eval(
4243
if (algo_input_ == nullptr) {
4344
THROW_RUNTIME_ERROR("GDSAlgoOpr: algo input is null");
4445
}
45-
algo_input_->bindParams(params);
46-
return algo_func_->execFunc(*algo_input_, graph_interface);
46+
auto bound_input = algo_input_->bindParams(params);
47+
const auto& input = bound_input ? *bound_input : *algo_input_;
48+
return algo_func_->execFunc(input, graph_interface);
4749
}
4850

4951
neug::result<OpBuildResultT> GDSAlgoOprBuilder::Build(

‎src/execution/execute/ops/retrieve/procedure_call.cc‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,14 @@ namespace execution {
2323
namespace ops {
2424
class ProcedureCallOpr : public IOperator {
2525
private:
26-
std::unique_ptr<neug::function::CallFuncInputBase> callInput;
27-
function::NeugCallFunction* callFunction;
26+
// Unbound template from bindFunc (shared across concurrent Evals).
27+
std::unique_ptr<neug::function::CallFuncInputBase> unboundInput_;
28+
function::NeugCallFunction* callFunction_;
2829

2930
public:
3031
ProcedureCallOpr(std::unique_ptr<neug::function::CallFuncInputBase> input,
3132
function::NeugCallFunction* callFunction)
32-
: callInput(std::move(input)), callFunction(callFunction) {}
33+
: unboundInput_(std::move(input)), callFunction_(callFunction) {}
3334

3435
~ProcedureCallOpr() override = default;
3536

@@ -39,14 +40,25 @@ class ProcedureCallOpr : public IOperator {
3940
IStorageInterface& graph, const ParamsMap& params,
4041
neug::execution::Context&& ctx,
4142
neug::execution::OprTimer* timer) override {
42-
if (callFunction == nullptr) {
43+
(void) ctx;
44+
(void) timer;
45+
if (callFunction_ == nullptr) {
4346
THROW_RUNTIME_ERROR("ProcedureCallOpr: callFunction is nullptr");
4447
}
45-
callInput->bindParams(params);
48+
if (unboundInput_ == nullptr) {
49+
THROW_RUNTIME_ERROR("ProcedureCallOpr: unbound input is nullptr");
50+
}
51+
if (callFunction_->execFunc == nullptr) {
52+
THROW_RUNTIME_ERROR("ProcedureCallOpr: execFunc is nullptr");
53+
}
54+
// bindParams returns a per-Eval bound input; nullptr means no deferred
55+
// params and the unbound template is safe to exec as-is.
56+
auto boundInput = unboundInput_->bindParams(params);
57+
const auto& input = boundInput ? *boundInput : *unboundInput_;
4658
return neug::result<neug::execution::Context>(
47-
callFunction->execFunc(*callInput, graph));
48-
} // namespace ops
49-
}; // namespace execution
59+
callFunction_->execFunc(input, graph));
60+
}
61+
};
5062

5163
neug::result<OpBuildResultT> ProcedureCallOprBuilder::Build(
5264
const neug::Schema& schema, const ContextMeta& ctx_meta,
@@ -77,4 +89,4 @@ neug::result<OpBuildResultT> ProcedureCallOprBuilder::Build(
7789

7890
} // namespace ops
7991
} // namespace execution
80-
} // namespace neug
92+
} // namespace neug

‎src/execution/execute/plan_parser.cc‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,21 @@ static void parse_params_type_impl(const physical::PhysicalPlan& plan,
537537
expression_parse(unfold_opr.input_expr(), params_type);
538538
break;
539539
}
540+
case physical::PhysicalOpr_Operator::OpKindCase::kProcedureCall: {
541+
const auto& query = plan.plan(i).opr().procedure_call().query();
542+
for (int j = 0; j < query.arguments_size(); ++j) {
543+
const auto& arg = query.arguments(j);
544+
if (!arg.has_param()) {
545+
continue;
546+
}
547+
const auto& param = arg.param();
548+
if (params_type.find(param.name()) != params_type.end()) {
549+
continue;
550+
}
551+
params_type[param.name()] = parse_from_ir_data_type(param.data_type());
552+
}
553+
break;
554+
}
540555
default: {
541556
break;
542557
}

‎tests/unittest/CMakeLists.txt‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@ add_neug_test(test_mmap_container test_mmap_container.cc)
2222

2323
add_neug_test(test_connection test_connection.cc)
2424

25-
add_neug_test(test_call_params test_call_params.cc)
26-
2725
add_neug_test(test_query_result test_query_result.cc)
2826

2927
add_neug_test(test_column test_column.cc)

0 commit comments

Comments
 (0)