| 77 | } |
| 78 | |
| 79 | Result<std::unique_ptr<substrait::ProjectRel>> CreateProject( |
| 80 | Id function_id, const std::vector<std::string>& arguments, |
| 81 | const std::unordered_map<std::string, std::vector<std::string>> options, |
| 82 | const std::vector<std::shared_ptr<DataType>>& arg_types, const DataType& output_type, |
| 83 | ExtensionSet* ext_set) { |
| 84 | auto project = std::make_unique<substrait::ProjectRel>(); |
| 85 | |
| 86 | auto call = std::make_unique<substrait::Expression::ScalarFunction>(); |
| 87 | ARROW_ASSIGN_OR_RAISE(uint32_t function_anchor, ext_set->EncodeFunction(function_id)); |
| 88 | call->set_function_reference(function_anchor); |
| 89 | |
| 90 | std::size_t arg_index = 0; |
| 91 | std::size_t table_arg_index = 0; |
| 92 | for (const std::shared_ptr<DataType>& arg_type : arg_types) { |
| 93 | substrait::FunctionArgument* argument = call->add_arguments(); |
| 94 | if (arg_type) { |
| 95 | // If it has a type then it's a reference to the input table |
| 96 | auto expression = std::make_unique<substrait::Expression>(); |
| 97 | CreateDirectReference(static_cast<int32_t>(table_arg_index++), expression.get()); |
| 98 | argument->set_allocated_value(expression.release()); |
| 99 | } else { |
| 100 | // If it doesn't have a type then it's an enum |
| 101 | const std::string& enum_value = arguments[arg_index]; |
| 102 | argument->set_enum_(enum_value); |
| 103 | } |
| 104 | arg_index++; |
| 105 | } |
| 106 | for (const auto& opt : options) { |
| 107 | substrait::FunctionOption* option = call->add_options(); |
| 108 | option->set_name(opt.first); |
| 109 | for (const std::string& pref : opt.second) { |
| 110 | option->add_preference(pref); |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | ARROW_ASSIGN_OR_RAISE( |
| 115 | std::unique_ptr<substrait::Type> output_type_substrait, |
| 116 | ToProto(output_type, /*nullable=*/true, ext_set, kPlanBuilderConversionOptions)); |
| 117 | call->set_allocated_output_type(output_type_substrait.release()); |
| 118 | |
| 119 | substrait::Expression* call_expression = project->add_expressions(); |
| 120 | call_expression->set_allocated_scalar_function(call.release()); |
| 121 | |
| 122 | return project; |
| 123 | } |
| 124 | |
| 125 | Result<std::unique_ptr<substrait::AggregateRel>> CreateAgg(Id function_id, |
| 126 | const std::vector<int>& keys, |
no test coverage detected