1 | #include "duckdb/function/scalar/math_functions.hpp" |
2 | #include "duckdb/common/vector_operations/vector_operations.hpp" |
3 | #include "duckdb/execution/expression_executor.hpp" |
4 | #include "duckdb/main/client_context.hpp" |
5 | #include "duckdb/planner/expression/bound_function_expression.hpp" |
6 | #include <random> |
7 | |
8 | using namespace duckdb; |
9 | using namespace std; |
10 | |
11 | struct RandomBindData : public FunctionData { |
12 | ClientContext &context; |
13 | uniform_real_distribution<double> dist; |
14 | |
15 | RandomBindData(ClientContext &context, uniform_real_distribution<double> dist) : context(context), dist(dist) { |
16 | } |
17 | |
18 | unique_ptr<FunctionData> Copy() override { |
19 | return make_unique<RandomBindData>(context, dist); |
20 | } |
21 | }; |
22 | |
23 | static void random_function(DataChunk &args, ExpressionState &state, Vector &result) { |
24 | assert(args.column_count() == 0); |
25 | auto &func_expr = (BoundFunctionExpression &)state.expr; |
26 | auto &info = (RandomBindData &)*func_expr.bind_info; |
27 | |
28 | result.vector_type = VectorType::FLAT_VECTOR; |
29 | auto result_data = FlatVector::GetData<double>(result); |
30 | for (idx_t i = 0; i < args.size(); i++) { |
31 | result_data[i] = info.dist(info.context.random_engine); |
32 | } |
33 | } |
34 | |
35 | unique_ptr<FunctionData> random_bind(BoundFunctionExpression &expr, ClientContext &context) { |
36 | uniform_real_distribution<double> dist(0, 1); |
37 | return make_unique<RandomBindData>(context, move(dist)); |
38 | } |
39 | |
40 | void RandomFun::RegisterFunction(BuiltinFunctions &set) { |
41 | set.AddFunction(ScalarFunction("random" , {}, SQLType::DOUBLE, random_function, true, random_bind)); |
42 | } |
43 | |