passing functionSignature instead of name and arguments

Author: chaojun-zhang

src/CMakeLists.txt

@@ -47,7 +47,9 @@ set(SRCS
     Type.cpp
     Function.cpp
     Extension.cpp
-    FunctionMapping.h FunctionLookup.cpp FunctionLookup.h)
+    FunctionMapping.h
+    FunctionLookup.cpp
+    FunctionSignature.h)
 
 add_library(substrait-cpp ${SRCS})
 

src/Extension.cpp

@@ -264,48 +264,6 @@ TypeVariantPtr Extension::lookupType(const std::string &typeName) const {
   return nullptr;
 }
 
-FunctionVariantPtr
-Extension::lookupScalarFunction(const std::string &name,
-                                const std::vector<TypePtr> &types) const {
-  auto functionVariantIter = scalarFunctionVariantMap_.find(name);
-  if (functionVariantIter != scalarFunctionVariantMap_.end()) {
-    for (const auto &candidateFunctionVariant : functionVariantIter->second) {
-      if (candidateFunctionVariant->tryMatch(types)) {
-        return candidateFunctionVariant;
-      }
-    }
-  }
-  return nullptr;
-}
-
-FunctionVariantPtr
-Extension::lookupAggregateFunction(const std::string &name,
-                                   const std::vector<TypePtr> &types) const {
-  auto functionVariantIter = aggregateFunctionVariantMap_.find(name);
-  if (functionVariantIter != aggregateFunctionVariantMap_.end()) {
-    for (const auto &candidateFunctionVariant : functionVariantIter->second) {
-      if (candidateFunctionVariant->tryMatch(types)) {
-        return candidateFunctionVariant;
-      }
-    }
-  }
-  return nullptr;
-}
-
-FunctionVariantPtr
-Extension::lookupWindowFunction(const std::string &name,
-                                const std::vector<TypePtr> &types) const {
-  auto functionVariantIter = windowFunctionVariantMap_.find(name);
-  if (functionVariantIter != windowFunctionVariantMap_.end()) {
-    for (const auto &candidateFunctionVariant : functionVariantIter->second) {
-      if (candidateFunctionVariant->tryMatch(types)) {
-        return candidateFunctionVariant;
-      }
-    }
-  }
-  return nullptr;
-}
-
 void Extension::addScalarFunctionVariant(
     const FunctionVariantPtr &functionVariant) {
   const auto &functionVariants =

src/Extension.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "Function.h"
+#include "FunctionSignature.h"
 #include "Type.h"
 
 namespace io::substrait {
@@ -57,28 +58,22 @@ class Extension {
   /// Add a type variant.
   void addTypeVariant(const TypeVariantPtr &functionVariant);
 
-  /// Lookup scalar function variant by given name and function types.
-  /// @return matched function variant.
-  FunctionVariantPtr
-  lookupScalarFunction(const std::string &name,
-                       const std::vector<TypePtr> &types) const;
-
-  /// Lookup aggregate function variant by given name and function types.
-  /// @return matched function variant.
-  FunctionVariantPtr
-  lookupAggregateFunction(const std::string &name,
-                          const std::vector<TypePtr> &types) const;
-
-  /// Lookup window function variant by given name and function types.
-  /// @return matched function variant.
-  FunctionVariantPtr
-  lookupWindowFunction(const std::string &name,
-                       const std::vector<TypePtr> &types) const;
-
   /// Lookup type variant by given type name.
   /// @return matched type variant
   TypeVariantPtr lookupType(const std::string &typeName) const;
 
+  const FunctionVariantMap &scalaFunctionVariantMap() const {
+    return scalarFunctionVariantMap_;
+  }
+
+  const FunctionVariantMap &windowFunctionVariantMap() const {
+    return windowFunctionVariantMap_;
+  }
+
+  const FunctionVariantMap &aggregateFunctionVariantMap() const {
+    return aggregateFunctionVariantMap_;
+  }
+
 private:
   static std::shared_ptr<Extension> loadDefault();
 

src/Function.cpp

@@ -18,15 +18,14 @@
 namespace io::substrait {
 
 std::string
-FunctionVariant::signature(
-    const std::string& name,
-    const std::vector<FunctionArgumentPtr>& arguments) {
+FunctionVariant::signature(const std::string &name,
+                           const std::vector<FunctionArgumentPtr> &arguments) {
   std::stringstream ss;
   ss << name;
   if (!arguments.empty()) {
     ss << ":";
     for (auto it = arguments.begin(); it != arguments.end(); ++it) {
-      const auto& typeSign = (*it)->toTypeString();
+      const auto &typeSign = (*it)->toTypeString();
       if (it == arguments.end() - 1) {
         ss << typeSign;
       } else {
@@ -38,8 +37,8 @@ FunctionVariant::signature(
   return ss.str();
 }
 
-bool FunctionVariant::tryMatch(
-    const std::vector<TypePtr>& actualTypes) {
+bool FunctionVariant::tryMatch(const FunctionSignature &signature) {
+  const auto &actualTypes = signature.getArguments();
   if (variadic.has_value()) {
     // return false if actual types length less than min of variadic
     const auto max = variadic->max;
@@ -48,24 +47,21 @@ bool FunctionVariant::tryMatch(
       return false;
     }
 
-    const auto& variadicArgument = arguments[0];
+    const auto &variadicArgument = arguments[0];
     // actual type must same as the variadicArgument
-    if (const auto& variadicValueArgument =
-            std::dynamic_pointer_cast<const ValueArgument>(
-                variadicArgument)) {
-      for (auto& actualType : actualTypes) {
+    if (const auto &variadicValueArgument =
+            std::dynamic_pointer_cast<const ValueArgument>(variadicArgument)) {
+      for (auto &actualType : actualTypes) {
         if (!variadicValueArgument->type->isSameAs(actualType)) {
           return false;
         }
       }
     }
-    return true;
   } else {
     std::vector<std::shared_ptr<const ValueArgument>> valueArguments;
-    for (const auto& argument : arguments) {
-      if (const auto& variadicValueArgument =
-              std::dynamic_pointer_cast<const ValueArgument>(
-                  argument)) {
+    for (const auto &argument : arguments) {
+      if (const auto &variadicValueArgument =
+              std::dynamic_pointer_cast<const ValueArgument>(argument)) {
         valueArguments.emplace_back(variadicValueArgument);
       }
     }
@@ -76,19 +72,24 @@ bool FunctionVariant::tryMatch(
     }
 
     for (auto i = 0; i < actualTypes.size(); i++) {
-      const auto& valueArgument = valueArguments[i];
+      const auto &valueArgument = valueArguments[i];
       if (!valueArgument->type->isSameAs(actualTypes[i])) {
         return false;
       }
     }
+  }
+  const auto &sigReturnType = signature.getReturnType();
+  if (this->returnType && sigReturnType) {
+    return returnType->isSameAs(sigReturnType);
+  } else {
     return true;
   }
 }
 
-bool AggregateFunctionVariant::tryMatch(
-    const std::vector<TypePtr>& actualTypes) {
-  bool matched = FunctionVariant::tryMatch(actualTypes);
+bool AggregateFunctionVariant::tryMatch(const FunctionSignature &signature) {
+  bool matched = FunctionVariant::tryMatch(signature);
   if (!matched && intermediate) {
+    const auto &actualTypes = signature.getArguments();
     if (actualTypes.size() == 1) {
       return intermediate->isSameAs(actualTypes[0]);
     }

src/Function.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "FunctionSignature.h"
 #include "Type.h"
 #include <iostream>
 
@@ -71,7 +72,7 @@ struct FunctionVariant {
   std::optional<FunctionVariadic> variadic;
 
   /// Test if the actual types matched with this function variant.
-  virtual bool tryMatch(const std::vector<TypePtr> &actualTypes);
+  virtual bool tryMatch(const FunctionSignature& signature);
 
   /// Create function signature by given function name and arguments.
   static std::string
@@ -89,7 +90,7 @@ struct ScalarFunctionVariant : public FunctionVariant {};
 struct AggregateFunctionVariant : public FunctionVariant {
   TypePtr intermediate;
 
-  bool tryMatch(const std::vector<TypePtr> &actualTypes) override;
+  bool tryMatch(const FunctionSignature& signature) override;
 };
 
 } // namespace io::substrait

src/FunctionLookup.cpp

@@ -17,35 +17,24 @@
 namespace io::substrait {
 
 FunctionVariantPtr
-FunctionLookup::lookupScalarFunction(const std::string &functionName,
-                                     const std::vector<TypePtr> &types) const {
-  const auto &functionMappings = functionMapping_->scalaMapping();
-  const auto &substraitFunctionName =
-      functionMappings.find(functionName) != functionMappings.end()
-          ? functionMappings.at(functionName)
-          : functionName;
-  return extension_->lookupScalarFunction(substraitFunctionName, types);
-}
+FunctionLookup::lookupFunction(const FunctionSignature &signature) const {
+  const auto &functionMappings = getFunctionMap();
 
-FunctionVariantPtr FunctionLookup::lookupAggregateFunction(
-    const std::string &functionName, const std::vector<TypePtr> &types) const {
-  const auto &functionMappings = functionMapping_->aggregateMapping();
   const auto &substraitFunctionName =
-      functionMappings.find(functionName) != functionMappings.end()
-          ? functionMappings.at(functionName)
-          : functionName;
-  return extension_->lookupAggregateFunction(substraitFunctionName, types);
-}
+      functionMappings.find(signature.getName()) != functionMappings.end()
+          ? functionMappings.at(signature.getName())
+          : signature.getName();
 
-FunctionVariantPtr
-FunctionLookup::lookupWindowFunction(const std::string &functionName,
-                                     const std::vector<TypePtr> &types) const {
-  const auto &functionMappings = functionMapping_->windowMapping();
-  const auto &substraitFunctionName =
-      functionMappings.find(functionName) != functionMappings.end()
-          ? functionMappings.at(functionName)
-          : functionName;
-  return extension_->lookupWindowFunction(substraitFunctionName, types);
+  const auto &functionVariants = getFunctionVariants();
+  auto functionVariantIter = functionVariants.find(substraitFunctionName);
+  if (functionVariantIter != functionVariants.end()) {
+    for (const auto &candidateFunctionVariant : functionVariantIter->second) {
+      if (candidateFunctionVariant->tryMatch(signature)) {
+        return candidateFunctionVariant;
+      }
+    }
+  }
+  return nullptr;
 }
 
 } // namespace io::substrait

src/FunctionLookup.h

@@ -16,6 +16,7 @@
 
 #include "Extension.h"
 #include "FunctionMapping.h"
+#include "FunctionSignature.h"
 
 namespace io::substrait {
 
@@ -25,23 +26,69 @@ class FunctionLookup {
                  const FunctionMappingPtr &functionMapping)
       : extension_(extension), functionMapping_(functionMapping) {}
 
-  FunctionVariantPtr
-  lookupScalarFunction(const std::string &functionName,
-                       const std::vector<TypePtr> &types) const;
+  virtual FunctionVariantPtr
+  lookupFunction(const FunctionSignature &signature) const;
 
-  FunctionVariantPtr
-  lookupAggregateFunction(const std::string &functionName,
-                          const std::vector<TypePtr> &types) const;
+  virtual ~FunctionLookup() {}
 
-  FunctionVariantPtr
-  lookupWindowFunction(const std::string &functionName,
-                       const std::vector<TypePtr> &types) const;
+protected:
+  virtual const FunctionMap getFunctionMap() const = 0;
+
+  virtual const FunctionVariantMap &getFunctionVariants() const = 0;
+
+  const FunctionMappingPtr functionMapping_;
 
-private:
   ExtensionPtr extension_;
-  FunctionMappingPtr functionMapping_;
 };
 
 using FunctionLookupPtr = std::shared_ptr<const FunctionLookup>;
 
+class ScalarFunctionLookup : public FunctionLookup {
+public:
+  ScalarFunctionLookup(const ExtensionPtr &extension,
+                       const FunctionMappingPtr &functionMapping)
+      : FunctionLookup(extension, functionMapping) {}
+
+protected:
+  const FunctionMap getFunctionMap() const override {
+    return functionMapping_->scalaMapping();
+  }
+
+  const FunctionVariantMap &getFunctionVariants() const override {
+    return extension_->scalaFunctionVariantMap();
+  }
+};
+
+class AggregateFunctionLookup : public FunctionLookup {
+public:
+  AggregateFunctionLookup(const ExtensionPtr &extension,
+                          const FunctionMappingPtr &functionMapping)
+      : FunctionLookup(extension, functionMapping) {}
+
+protected:
+  const FunctionMap getFunctionMap() const override {
+    return functionMapping_->aggregateMapping();
+  }
+
+  const FunctionVariantMap &getFunctionVariants() const override {
+    return extension_->aggregateFunctionVariantMap();
+  }
+};
+
+class WindowFunctionLookup : public FunctionLookup {
+public:
+  WindowFunctionLookup(const ExtensionPtr &extension,
+                       const FunctionMappingPtr &functionMapping)
+      : FunctionLookup(extension, functionMapping) {}
+
+protected:
+  const FunctionMap getFunctionMap() const override {
+    return functionMapping_->windowMapping();
+  }
+
+  const FunctionVariantMap &getFunctionVariants() const override {
+    return extension_->windowFunctionVariantMap();
+  }
+};
+
 } // namespace io::substrait