feat: add support for enum arguments in textplans

substrait-io · Jun 16, 2023 · d8ac0eb · d8ac0eb
1 parent fee0829
commit d8ac0eb
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 7 deletions.
diff --git a/src/substrait/textplan/converter/PlanPrinterVisitor.cpp b/src/substrait/textplan/converter/PlanPrinterVisitor.cpp
@@ -59,6 +59,12 @@ std::string invocationToString(
  return "unspecified";
 }
 
+std::string visitEnumArgument(const std::string& str) {
+ std::stringstream text;
+ text << str << "_enum";
+ return text.str();
+}
+
 } // namespace
 
 std::string PlanPrinterVisitor::printRelation(const SymbolInfo& symbol) {
@@ -393,10 +399,7 @@ std::any PlanPrinterVisitor::visitScalarFunction(
  }
  switch (arg.arg_type_case()) {
  case ::substrait::proto::FunctionArgument::kEnum:
- errorListener_->addError(
- "Enum arguments not yet supported in scalar functions: " +
- arg.ShortDebugString());
- text << "ENUM_NOT_SUPPORTED";
+ text << visitEnumArgument(arg.enum_());
  break;
  case ::substrait::proto::FunctionArgument::kType:
  text << ANY_CAST(std::string, visitType(arg.type()));
@@ -500,8 +503,9 @@ std::any PlanPrinterVisitor::visitNested(
 std::any PlanPrinterVisitor::visitEnum(
  const ::substrait::proto::Expression_Enum& value) {
  errorListener_->addError(
- "Enum expressions are not yet supported: " + value.ShortDebugString());
- return std::string("ENUM_NOT_YET_IMPLEMENTED");
+ "Enum expressions are deprecated and not supported: " +
+ value.ShortDebugString());
+ return std::string("ENUM_EXPRESSION_DEPRECATED");
 }
 
 std::any PlanPrinterVisitor::visitStructSelect(
@@ -569,7 +573,7 @@ std::any PlanPrinterVisitor::visitAggregateFunction(
  }
  switch (arg.arg_type_case()) {
  case ::substrait::proto::FunctionArgument::kEnum:
- text << "ENUM_NOT_SUPPORTED";
+ text << visitEnumArgument(arg.enum_());
  break;
  case ::substrait::proto::FunctionArgument::kType:
  text << ANY_CAST(std::string, visitType(arg.type()));

diff --git a/src/substrait/textplan/parser/SubstraitPlanRelationVisitor.cpp b/src/substrait/textplan/parser/SubstraitPlanRelationVisitor.cpp
@@ -11,6 +11,7 @@
 #include "SubstraitPlanParser/SubstraitPlanParser.h"
 #include "SubstraitPlanTypeVisitor.h"
 #include "absl/strings/numbers.h"
+#include "absl/strings/strip.h"
 #include "date/tz.h"
 #include "substrait/expression/DecimalLiteral.h"
 #include "substrait/proto/algebra.pb.h"
@@ -48,6 +49,12 @@ bool startsWith(const std::string& haystack, std::string_view needle) {
  return strncmp(haystack.c_str(), needle.data(), needle.size()) == 0;
 }
 
+// Returns true if the string 'haystack' ends with the string 'needle'.
+bool endsWith(const std::string& haystack, const std::string& needle) {
+ return haystack.size() > needle.size() &&
+ haystack.substr(haystack.size() - needle.size()) == needle;
+}
+
 void setNullable(::substrait::proto::Type* type) {
  switch (type->kind_case()) {
  case ::substrait::proto::Type::kBool:
@@ -748,6 +755,13 @@ std::any SubstraitPlanRelationVisitor::visitExpressionFunctionUse(
 
  expr.mutable_scalar_function()->set_function_reference(funcReference);
  for (const auto& exp : ctx->expression()) {
+ if (endsWith(exp->getText(), "_enum")) {
+ auto str = exp->getText();
+ str = absl::StripSuffix(str, "_enum");
+ expr.mutable_scalar_function()->add_arguments()->set_enum_(str);
+ continue;
+ }
+
  auto result = visitExpression(exp);
  if (result.type() != typeid(::substrait::proto::Expression)) {
  errorListener_->addError(