ARROW-17966: [C++] Adjust to new format for Substrait optional arguments

cpp/src/arrow/engine/substrait/expression_internal.cc

@@ -52,22 +52,11 @@ Id NormalizeFunctionName(Id id) {
 
 } // namespace
 
-Status DecodeArg(const substrait::FunctionArgument& arg, uint32_t idx,
-SubstraitCall* call, const ExtensionSet& ext_set,
+Status DecodeArg(const substrait::FunctionArgument& arg, int idx, SubstraitCall* call,
+ const ExtensionSet& ext_set,
  const ConversionOptions& conversion_options) {
  if (arg.has_enum_()) {
- const substrait::FunctionArgument::Enum& enum_val = arg.enum_();
- switch (enum_val.enum_kind_case()) {
- case substrait::FunctionArgument::Enum::EnumKindCase::kSpecified:
- call->SetEnumArg(idx, enum_val.specified());
- break;
- case substrait::FunctionArgument::Enum::EnumKindCase::kUnspecified:
- call->SetEnumArg(idx, std::nullopt);
- break;
- default:
- return Status::Invalid("Unrecognized enum kind case: ",
- enum_val.enum_kind_case());
- }
+ call->SetEnumArg(idx, arg.enum_());
  } else if (arg.has_value()) {
  ARROW_ASSIGN_OR_RAISE(compute::Expression expr,
  FromProto(arg.value(), ext_set, conversion_options));
@@ -80,15 +69,31 @@ Status DecodeArg(const substrait::FunctionArgument& arg, uint32_t idx,
  return Status::OK();
 }
 
+Status DecodeOption(const substrait::FunctionOption& opt, SubstraitCall* call) {
+ std::vector<std::string_view> prefs;
+ if (opt.preference_size() == 0) {
+ return Status::Invalid("Invalid Substrait plan. The option ", opt.name(),
+ " is specified but does not list any choices");
+ }
+ for (const auto& preference : opt.preference()) {
+ prefs.push_back(preference);
+ }
+ call->SetOption(opt.name(), prefs);
+ return Status::OK();
+}
+
 Result<SubstraitCall> DecodeScalarFunction(
  Id id, const substrait::Expression::ScalarFunction& scalar_fn,
  const ExtensionSet& ext_set, const ConversionOptions& conversion_options) {
  ARROW_ASSIGN_OR_RAISE(auto output_type_and_nullable,
  FromProto(scalar_fn.output_type(), ext_set, conversion_options));
  SubstraitCall call(id, output_type_and_nullable.first, output_type_and_nullable.second);
  for (int i = 0; i < scalar_fn.arguments_size(); i++) {
- ARROW_RETURN_NOT_OK(DecodeArg(scalar_fn.arguments(i), static_cast<uint32_t>(i), &call,
- ext_set, conversion_options));
+ ARROW_RETURN_NOT_OK(
+ DecodeArg(scalar_fn.arguments(i), i, &call, ext_set, conversion_options));
+ }
+ for (const auto& opt : scalar_fn.options()) {
+ ARROW_RETURN_NOT_OK(DecodeOption(opt, &call));
  }
  return std::move(call);
 }
@@ -927,17 +932,11 @@ Result<std::unique_ptr<substrait::Expression::ScalarFunction>> EncodeSubstraitCa
  ToProto(*call.output_type(), call.output_nullable(), ext_set, conversion_options));
  scalar_fn->set_allocated_output_type(output_type.release());
 
- for (uint32_t i = 0; i < call.size(); i++) {
+ for (int i = 0; i < call.size(); i++) {
  substrait::FunctionArgument* arg = scalar_fn->add_arguments();
  if (call.HasEnumArg(i)) {
- auto enum_val = std::make_unique<substrait::FunctionArgument::Enum>();
- ARROW_ASSIGN_OR_RAISE(std::optional<std::string_view> enum_arg, call.GetEnumArg(i));
- if (enum_arg) {
- enum_val->set_specified(std::string(*enum_arg));
- } else {
- enum_val->set_allocated_unspecified(new google::protobuf::Empty());
- }
- arg->set_allocated_enum_(enum_val.release());
+ ARROW_ASSIGN_OR_RAISE(std::string_view enum_val, call.GetEnumArg(i));
+ arg->set_enum_(std::string(enum_val));
  } else if (call.HasValueArg(i)) {
  ARROW_ASSIGN_OR_RAISE(compute::Expression value_arg, call.GetValueArg(i));
  ARROW_ASSIGN_OR_RAISE(std::unique_ptr<substrait::Expression> value_expr,

cpp/src/arrow/engine/substrait/extension_set.cc

@@ -25,6 +25,7 @@
 #include "arrow/engine/substrait/expression_internal.h"
 #include "arrow/util/hash_util.h"
 #include "arrow/util/hashing.h"
+#include "arrow/util/string.h"
 
 namespace arrow {
 namespace engine {
@@ -121,7 +122,7 @@ class IdStorageImpl : public IdStorage {
 
 std::unique_ptr<IdStorage> IdStorage::Make() { return std::make_unique<IdStorageImpl>(); }
 
-Result<std::optional<std::string_view>> SubstraitCall::GetEnumArg(uint32_t index) const {
+Result<std::string_view> SubstraitCall::GetEnumArg(int index) const {
  if (index >= size_) {
  return Status::Invalid("Expected Substrait call to have an enum argument at index ",
  index, " but it did not have enough arguments");
@@ -134,16 +135,16 @@ Result<std::optional<std::string_view>> SubstraitCall::GetEnumArg(uint32_t index
  return enum_arg_it->second;
 }
 
-bool SubstraitCall::HasEnumArg(uint32_t index) const {
+bool SubstraitCall::HasEnumArg(int index) const {
  return enum_args_.find(index) != enum_args_.end();
 }
 
-void SubstraitCall::SetEnumArg(uint32_t index, std::optional<std::string> enum_arg) {
+void SubstraitCall::SetEnumArg(int index, std::string enum_arg) {
  size_ = std::max(size_, index + 1);
  enum_args_[index] = std::move(enum_arg);
 }
 
-Result<compute::Expression> SubstraitCall::GetValueArg(uint32_t index) const {
+Result<compute::Expression> SubstraitCall::GetValueArg(int index) const {
  if (index >= size_) {
  return Status::Invalid("Expected Substrait call to have a value argument at index ",
  index, " but it did not have enough arguments");
@@ -156,15 +157,32 @@ Result<compute::Expression> SubstraitCall::GetValueArg(uint32_t index) const {
  return value_arg_it->second;
 }
 
-bool SubstraitCall::HasValueArg(uint32_t index) const {
+bool SubstraitCall::HasValueArg(int index) const {
  return value_args_.find(index) != value_args_.end();
 }
 
-void SubstraitCall::SetValueArg(uint32_t index, compute::Expression value_arg) {
+void SubstraitCall::SetValueArg(int index, compute::Expression value_arg) {
  size_ = std::max(size_, index + 1);
  value_args_[index] = std::move(value_arg);
 }
 
+std::optional<std::vector<std::string> const*> SubstraitCall::GetOption(
+ std::string_view option_name) const {
+ auto opt = options_.find(std::string(option_name));
+ if (opt == options_.end()) {
+ return std::nullopt;
+ }
+ return &opt->second;
+}
+
+void SubstraitCall::SetOption(std::string_view option_name,
+ const std::vector<std::string_view>& option_preferences) {
+ auto& prefs = options_[std::string(option_name)];
+ for (std::string_view pref : option_preferences) {
+ prefs.emplace_back(pref);
+ }
+}
+
 // A builder used when creating a Substrait plan from an Arrow execution plan. In
 // that situation we do not have a set of anchor values already defined so we keep
 // a map of what Ids we have seen.
@@ -645,50 +663,91 @@ struct ExtensionIdRegistryImpl : ExtensionIdRegistry {
 };
 
 template <typename Enum>
-using EnumParser = std::function<Result<Enum>(std::optional<std::string_view>)>;
-
-template <typename Enum>
-EnumParser<Enum> GetEnumParser(const std::vector<std::string>& options) {
-std::unordered_map<std::string, Enum> parse_map;
-for (std::size_t i = 0; i < options.size(); i++) {
-parse_map[options[i]] = static_cast<Enum>(i + 1);
+class EnumParser {
+ public:
+ explicit EnumParser(const std::vector<std::string>& options) {
+ for (std::size_t i = 0; i < options.size(); i++) {
+ parse_map_[options[i]] = static_cast<Enum>(i + 1);
+ reverse_map_[static_cast<Enum>(i + 1)] = options[i];
+}
  }
- return [parse_map](std::optional<std::string_view> enum_val) -> Result<Enum> {
- if (!enum_val) {
- // Assumes 0 is always kUnspecified in Enum
- return static_cast<Enum>(0);
+
+ Result<Enum> Parse(std::string_view enum_val) const {
+ auto it = parse_map_.find(std::string(enum_val));
+ if (it == parse_map_.end()) {
+ return Status::NotImplemented("The value ", enum_val,
+ " is not an expected enum value");
  }
- auto maybe_parsed = parse_map.find(std::string(*enum_val));
- if (maybe_parsed == parse_map.end()) {
- return Status::Invalid("The value ", *enum_val, " is not an expected enum value");
+ return it->second;
+ }
+
+ std::string ImplementedOptionsAsString(
+ const std::vector<Enum>& implemented_opts) const {
+ std::vector<std::string_view> opt_strs;
+ for (const Enum& implemented_opt : implemented_opts) {
+ auto it = reverse_map_.find(implemented_opt);
+ if (it == reverse_map_.end()) {
+ opt_strs.emplace_back("Unknown");
+ } else {
+ opt_strs.emplace_back(it->second);
+ }
  }
- return maybe_parsed->second;
- };
-}
+ return arrow::internal::JoinStrings(opt_strs, ", ");
+ }
+
+ private:
+ std::unordered_map<std::string, Enum> parse_map_;
+ std::unordered_map<Enum, std::string> reverse_map_;
+};
 
 enum class TemporalComponent { kUnspecified = 0, kYear, kMonth, kDay, kSecond };
 static std::vector<std::string> kTemporalComponentOptions = {"YEAR", "MONTH", "DAY",
  "SECOND"};
-static EnumParser<TemporalComponent> kTemporalComponentParser =
- GetEnumParser<TemporalComponent>(kTemporalComponentOptions);
+static EnumParser<TemporalComponent> kTemporalComponentParser(kTemporalComponentOptions);
 
 enum class OverflowBehavior { kUnspecified = 0, kSilent, kSaturate, kError };
 static std::vector<std::string> kOverflowOptions = {"SILENT", "SATURATE", "ERROR"};
-static EnumParser<OverflowBehavior> kOverflowParser =
- GetEnumParser<OverflowBehavior>(kOverflowOptions);
+static EnumParser<OverflowBehavior> kOverflowParser(kOverflowOptions);
 
 template <typename Enum>
-Result<Enum> ParseEnumArg(const SubstraitCall& call, uint32_t arg_index,
+Result<Enum> ParseOptionOrElse(const SubstraitCall& call, std::string_view option_name,
+ const EnumParser<Enum>& parser,
+ const std::vector<Enum>& implemented_options,
+ Enum fallback) {
+ std::optional<std::vector<std::string> const*> enum_arg = call.GetOption(option_name);
+ if (!enum_arg.has_value()) {
+ return fallback;
+ }
+ std::vector<std::string> const* prefs = *enum_arg;
+ for (const std::string& pref : *prefs) {
+ ARROW_ASSIGN_OR_RAISE(Enum parsed, parser.Parse(pref));
+ for (Enum implemented_opt : implemented_options) {
+ if (implemented_opt == parsed) {
+ return parsed;
+ }
+ }
+ }
+
+ // Prepare error message
+ return Status::NotImplemented(
+ "During a call to a function with id ", call.id().uri, "#", call.id().name,
+ " the plan requested the option ", option_name, " to be one of [",
+ arrow::internal::JoinStrings(*prefs, ", "),
+ "] but the only supported options are [",
+ parser.ImplementedOptionsAsString(implemented_options), "]");
+}
+
+template <typename Enum>
+Result<Enum> ParseEnumArg(const SubstraitCall& call, int arg_index,
  const EnumParser<Enum>& parser) {
- ARROW_ASSIGN_OR_RAISE(std::optional<std::string_view> enum_arg,
- call.GetEnumArg(arg_index));
- return parser(enum_arg);
+ ARROW_ASSIGN_OR_RAISE(std::string_view enum_val, call.GetEnumArg(arg_index));
+ return parser.Parse(enum_val);
 }
 
 Result<std::vector<compute::Expression>> GetValueArgs(const SubstraitCall& call,
  int start_index) {
  std::vector<compute::Expression> expressions;
- for (uint32_t index = start_index; index < call.size(); index++) {
+ for (int index = start_index; index < call.size(); index++) {
  ARROW_ASSIGN_OR_RAISE(compute::Expression arg, call.GetValueArg(index));
  expressions.push_back(arg);
  }
@@ -698,13 +757,13 @@ Result<std::vector<compute::Expression>> GetValueArgs(const SubstraitCall& call,
 ExtensionIdRegistry::SubstraitCallToArrow DecodeOptionlessOverflowableArithmetic(
  const std::string& function_name) {
  return [function_name](const SubstraitCall& call) -> Result<compute::Expression> {
- ARROW_ASSIGN_OR_RAISE(OverflowBehavior overflow_behavior,
- ParseEnumArg(call, 0, kOverflowParser));
+ ARROW_ASSIGN_OR_RAISE(
+ OverflowBehavior overflow_behavior,
+ ParseOptionOrElse(call, "overflow", kOverflowParser,
+ {OverflowBehavior::kSilent, OverflowBehavior::kError},
+ OverflowBehavior::kSilent));
  ARROW_ASSIGN_OR_RAISE(std::vector<compute::Expression> value_args,
- GetValueArgs(call, 1));
- if (overflow_behavior == OverflowBehavior::kUnspecified) {
- overflow_behavior = OverflowBehavior::kSilent;
- }
+ GetValueArgs(call, 0));
  if (overflow_behavior == OverflowBehavior::kSilent) {
  return arrow::compute::call(function_name, std::move(value_args));
  } else if (overflow_behavior == OverflowBehavior::kError) {
@@ -736,12 +795,12 @@ ExtensionIdRegistry::ArrowToSubstraitCall EncodeOptionlessOverflowableArithmetic
  SubstraitCall substrait_call(substrait_fn_id, call.type.GetSharedPtr(),
  /*nullable=*/true);
  if (kChecked) {
- substrait_call.SetEnumArg(0, "ERROR");
+ substrait_call.SetOption("overflow", {"ERROR"});
  } else {
- substrait_call.SetEnumArg(0, "SILENT");
+ substrait_call.SetOption("overflow", {"SILENT"});
  }
  for (std::size_t i = 0; i < call.arguments.size(); i++) {
- substrait_call.SetValueArg(static_cast<uint32_t>(i + 1), call.arguments[i]);
+ substrait_call.SetValueArg(static_cast<int>(i), call.arguments[i]);
  }
  return std::move(substrait_call);
  };
@@ -755,14 +814,14 @@ ExtensionIdRegistry::ArrowToSubstraitCall EncodeOptionlessComparison(Id substrai
  SubstraitCall substrait_call(substrait_fn_id, call.type.GetSharedPtr(),
  /*nullable=*/true);
  for (std::size_t i = 0; i < call.arguments.size(); i++) {
- substrait_call.SetValueArg(static_cast<uint32_t>(i), call.arguments[i]);
+ substrait_call.SetValueArg(static_cast<int>(i), call.arguments[i]);
  }
  return std::move(substrait_call);
  };
 }
 
 ExtensionIdRegistry::SubstraitCallToArrow DecodeOptionlessBasicMapping(
- const std::string& function_name, uint32_t max_args) {
+ const std::string& function_name, int max_args) {
  return [function_name,
  max_args](const SubstraitCall& call) -> Result<compute::Expression> {
  if (call.size() > max_args) {

cpp/src/arrow/engine/substrait/extension_set.h

@@ -125,13 +125,17 @@ class SubstraitCall {
  bool output_nullable() const { return output_nullable_; }
  bool is_hash() const { return is_hash_; }
 
- bool HasEnumArg(uint32_t index) const;
- Result<std::optional<std::string_view>> GetEnumArg(uint32_t index) const;
- void SetEnumArg(uint32_t index, std::optional<std::string> enum_arg);
- Result<compute::Expression> GetValueArg(uint32_t index) const;
- bool HasValueArg(uint32_t index) const;
- void SetValueArg(uint32_t index, compute::Expression value_arg);
- uint32_t size() const { return size_; }
+ bool HasEnumArg(int index) const;
+ Result<std::string_view> GetEnumArg(int index) const;
+ void SetEnumArg(int index, std::string enum_arg);
+ Result<compute::Expression> GetValueArg(int index) const;
+ bool HasValueArg(int index) const;
+ void SetValueArg(int index, compute::Expression value_arg);
+ std::optional<std::vector<std::string> const*> GetOption(
+ std::string_view option_name) const;
+ void SetOption(std::string_view option_name,
+ const std::vector<std::string_view>& option_preferences);
+ int size() const { return size_; }
 
  private:
  Id id_;
@@ -140,9 +144,10 @@ class SubstraitCall {
  // Only needed when converting from Substrait -> Arrow aggregates. The
  // Arrow function name depends on whether or not there are any groups
  bool is_hash_;
- std::unordered_map<uint32_t, std::optional<std::string>> enum_args_;
- std::unordered_map<uint32_t, compute::Expression> value_args_;
- uint32_t size_ = 0;
+ std::unordered_map<int, std::string> enum_args_;
+ std::unordered_map<int, compute::Expression> value_args_;
+ std::unordered_map<std::string, std::vector<std::string>> options_;
+ int size_ = 0;
 };
 
 /// Substrait identifies functions and custom data types using a (uri, name) pair.