vibhatha commented on code in PR #13613:
URL: https://github.com/apache/arrow/pull/13613#discussion_r939467462
##########
cpp/src/arrow/engine/substrait/extension_set.cc:
##########
@@ -261,155 +389,390 @@ struct ExtensionIdRegistryImpl : ExtensionIdRegistry {
return Status::OK();
}
- util::optional<FunctionRecord> GetFunction(
- util::string_view arrow_function_name) const override {
- if (auto index = GetIndex(function_name_to_index_, arrow_function_name)) {
- return FunctionRecord{function_ids_[*index],
*function_name_ptrs_[*index]};
+ Status CanAddSubstraitCallToArrow(Id substrait_function_id) const override {
+ if (substrait_to_arrow_.find(substrait_function_id) !=
substrait_to_arrow_.end()) {
+ return Status::Invalid(
+ "Cannot register function converter because a converter already
exists");
}
- return {};
+ if (parent_) {
+ return parent_->CanAddSubstraitCallToArrow(substrait_function_id);
+ }
+ return Status::OK();
}
- util::optional<FunctionRecord> GetFunction(Id id) const override {
- if (auto index = GetIndex(function_id_to_index_, id)) {
- return FunctionRecord{function_ids_[*index],
*function_name_ptrs_[*index]};
+ Status CanAddSubstraitAggregateToArrow(Id substrait_function_id) const
override {
+ if (substrait_to_arrow_agg_.find(substrait_function_id) !=
+ substrait_to_arrow_agg_.end()) {
+ return Status::Invalid(
+ "Cannot register function converter because a converter already
exists");
}
- return {};
+ if (parent_) {
+ return parent_->CanAddSubstraitAggregateToArrow(substrait_function_id);
+ }
+ return Status::OK();
+ }
+
+ template <typename ConverterType>
+ Status AddSubstraitToArrowFunc(
+ Id substrait_id, ConverterType conversion_func,
+ std::unordered_map<Id, ConverterType, IdHashEq, IdHashEq>* dest) {
+ // Convert id to view into registry-owned memory
+ Id copied_id = ids_.Emplace(substrait_id);
+
+ auto add_result = dest->emplace(copied_id, std::move(conversion_func));
+ if (!add_result.second) {
+ return Status::Invalid(
+ "Failed to register Substrait to Arrow function converter because a
converter "
+ "already existed");
+ }
+
+ return Status::OK();
+ }
+
+ Status AddSubstraitCallToArrow(Id substrait_function_id,
+ SubstraitCallToArrow conversion_func)
override {
+ if (parent_) {
+
ARROW_RETURN_NOT_OK(parent_->CanAddSubstraitCallToArrow(substrait_function_id));
+ }
+ return AddSubstraitToArrowFunc<SubstraitCallToArrow>(
+ substrait_function_id, std::move(conversion_func),
&substrait_to_arrow_);
}
- Status CanRegisterFunction(Id id,
- const std::string& arrow_function_name) const
override {
- if (function_id_to_index_.find(id) != function_id_to_index_.end()) {
- return Status::Invalid("Function id was already registered");
+ Status AddSubstraitAggregateToArrow(
+ Id substrait_function_id, SubstraitAggregateToArrow conversion_func)
override {
+ if (parent_) {
+ ARROW_RETURN_NOT_OK(
+ parent_->CanAddSubstraitAggregateToArrow(substrait_function_id));
}
- if (function_name_to_index_.find(arrow_function_name) !=
- function_name_to_index_.end()) {
- return Status::Invalid("Function name was already registered");
+ return AddSubstraitToArrowFunc<SubstraitAggregateToArrow>(
+ substrait_function_id, std::move(conversion_func),
&substrait_to_arrow_agg_);
+ }
+
+ template <typename ConverterType>
+ Status AddArrowToSubstraitFunc(std::string arrow_function_name,
ConverterType converter,
+ std::unordered_map<std::string,
ConverterType>* dest) {
+ auto add_result = dest->emplace(std::move(arrow_function_name),
std::move(converter));
+ if (!add_result.second) {
+ return Status::Invalid(
+ "Failed to register Arrow to Substrait function converter for Arrow
function ",
+ arrow_function_name, " because a converter already existed");
}
return Status::OK();
}
- Status RegisterFunction(Id id, std::string arrow_function_name) override {
- DCHECK_EQ(function_ids_.size(), function_name_ptrs_.size());
+ Status AddArrowToSubstraitCall(std::string arrow_function_name,
+ ArrowToSubstraitCall converter) override {
+ if (parent_) {
+
ARROW_RETURN_NOT_OK(parent_->CanAddArrowToSubstraitCall(arrow_function_name));
+ }
+ return AddArrowToSubstraitFunc(std::move(arrow_function_name), converter,
+ &arrow_to_substrait_);
+ }
- Id copied_id{*uris_.emplace(id.uri.to_string()).first,
- *names_.emplace(id.name.to_string()).first};
+ Status AddArrowToSubstraitAggregate(std::string arrow_function_name,
+ ArrowToSubstraitAggregate converter)
override {
+ if (parent_) {
+
ARROW_RETURN_NOT_OK(parent_->CanAddArrowToSubstraitAggregate(arrow_function_name));
+ }
+ return AddArrowToSubstraitFunc(std::move(arrow_function_name), converter,
+ &arrow_to_substrait_agg_);
+ }
- const std::string& copied_function_name{
- *function_names_.emplace(std::move(arrow_function_name)).first};
+ Status CanAddArrowToSubstraitCall(const std::string& function_name) const
override {
+ if (arrow_to_substrait_.find(function_name) != arrow_to_substrait_.end()) {
+ return Status::Invalid(
+ "Cannot register function converter because a converter already
exists");
Review Comment:
Yes this works.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
