From 82dc279c99b6a5765b89eab5ea4d213d420ed21c Mon Sep 17 00:00:00 2001
From: Masha Basmanova <mbasmanova@meta.com>
Date: Sun, 17 Mar 2024 23:13:57 -0700
Subject: [PATCH] Add support for DECIMAL types to Simple Function API (#9096)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:
**type/Type.h**

	Add P1, P2, P3, P4, S1, S2, S3, S4 types to specify precision and scale parameters for decimal types during function registration.

	Add LongDecimal<P, S> and ShortDecimal<P, S> templates to specify decimal argument and return types during function registration.

```
  registerFunction<
      DecimalAddFunction,
      LongDecimal<P3, S3>,
      LongDecimal<P1, S1>,
      LongDecimal<P2, S2>>({"plus"}, constraints);
```

**expression/UdfTypeResolver.h**

Define arg_type and out_type for LongDecimal and ShortDecimal.

```
	arg_type<LongDecimal> = int128_t
	out_type<LongDecimal> = int128_t

	arg_type<ShortDecimal> = int64_t
	out_type<ShortDecimal> = int64_t
```

**functions/Registerer.h**

	Add optional ‘constraints’ parameter to registerFunction template. This allows to specify rules for calculating precision and scale for decimal return types.

```
template <template <class> typename Func, typename TReturn, typename... TArgs>
void registerFunction(
    const std::vector<std::string>& aliases = {},
    const std::vector<exec::SignatureVariable>& constraints = {})
```

	Here is how we can specify calculation of precision and scale for the return type of plus(decimal, decimal).

```
  std::vector<exec::SignatureVariable> constraints = {
      exec::SignatureVariable(
          P3::name(),
          fmt::format(
              "min(38, max({a_precision} - {a_scale}, {b_precision} - {b_scale}) + max({a_scale}, {b_scale}) + 1)",
              fmt::arg("a_precision", P1::name()),
              fmt::arg("b_precision", P2::name()),
              fmt::arg("a_scale", S1::name()),
              fmt::arg("b_scale", S2::name())),
          exec::ParameterType::kIntegerParameter),
      exec::SignatureVariable(
          S3::name(),
          fmt::format(
              "max({a_scale}, {b_scale})",
              fmt::arg("a_scale", S1::name()),
              fmt::arg("b_scale", S2::name())),
          exec::ParameterType::kIntegerParameter),
  };
```

**core/SimpleFunctionMetadata.h**

Extend SimpleFunctionMetadata to store physical types (TypeKind) of input arguments and return type in addition to signature.

Decimal “plus” function has a single signature:

```
(decimal(p1, s1), decimal(p2, s2)) -> decimal(p3, s3)
```

But 5 implementations:

```
	(int64_t, int64_t) -> int64_t
	(int64_t, int64_t) -> int128_t
	(int64_t, int128_t) -> int128_t
	(int128_t, int64_t) -> int128_t
	(int128_t, int128_t) -> int128_t
```

We need a way to distinguish between these.

**expression/SimpleFunctionRegistry.h/cpp**

Allow for storing multiple implementations for a single signature.

```
using SignatureMap = std::unordered_map<
    FunctionSignature,
    std::vector<std::unique_ptr<const FunctionEntry>>>;
using FunctionMap = std::unordered_map<std::string, SignatureMap>;
```

Modify SimpleFunctionRegistry::resolveFunction method to find an implementation with matching signature and matching TypeKinds for arguments and return type.

**core/SimpleFunctionMetadata.h**

Introduce optional initializeTypes method for a function to receive input types. Functions that operate on decimal types use this method to get access to precision and scale of the arguments.

```
void initializeTypes(const std::vector<TypePtr>& argTypes)
```

**Example: Decimal ADD**

Here is how a function that adds 2 decimal numbers can be defined. This function supports adding decimal numbers with possibly different precision and scale.

```
template <typename TExec>
struct DecimalAddFunction {
  VELOX_DEFINE_FUNCTION_TYPES(TExec);

  void initializeTypes(const std::vector<TypePtr>& argTypes) {
    auto aType = argTypes[0];
    auto bType = argTypes[1];
    auto [aPrecision, aScale] = getDecimalPrecisionScale(*aType);
    auto [bPrecision, bScale] = getDecimalPrecisionScale(*bType);
    auto [rPrecision, rScale] = Addition::computeResultPrecisionScale(
        aPrecision, aScale, bPrecision, bScale);
    aRescale_ = Addition::computeRescaleFactor(aScale, bScale, rScale);
    bRescale_ = Addition::computeRescaleFactor(bScale, aScale, rScale);
  }

  template <typename R, typename A, typename B>
  void call(R& out, const A& a, const B& b) {
    Addition::template apply<R, A, B>(out, a, b, aRescale_, bRescale_);
  }

 private:
  uint8_t aRescale_;
  uint8_t bRescale_;
};
```

The registration involves specifying a rule for calculating precision and scale for the result based on precision and scale of the inputs and provides 5 implementations with all possible permutations of short and long decimals in the input and result.

```
  std::vector<exec::SignatureVariable> constraints = {
      exec::SignatureVariable(
          P3::name(),
          fmt::format(
              "min(38, max({a_precision} - {a_scale}, {b_precision} - {b_scale}) + max({a_scale}, {b_scale}) + 1)",
              fmt::arg("a_precision", P1::name()),
              fmt::arg("b_precision", P2::name()),
              fmt::arg("a_scale", S1::name()),
              fmt::arg("b_scale", S2::name())),
          exec::ParameterType::kIntegerParameter),
      exec::SignatureVariable(
          S3::name(),
          fmt::format(
              "max({a_scale}, {b_scale})",
              fmt::arg("a_scale", S1::name()),
              fmt::arg("b_scale", S2::name())),
          exec::ParameterType::kIntegerParameter),
  };

  // (long, long) -> long
  registerFunction<
      DecimalAddFunction,
      LongDecimal<P3, S3>,
      LongDecimal<P1, S1>,
      LongDecimal<P2, S2>>({"plus"}, constraints);

  // (short, short) -> short
  registerFunction<
      DecimalAddFunction,
      ShortDecimal<P3, S3>,
      ShortDecimal<P1, S1>,
      ShortDecimal<P2, S2>>({"plus"}, constraints);

  // (short, short) -> long
  registerFunction<
      DecimalAddFunction,
      LongDecimal<P3, S3>,
      ShortDecimal<P1, S1>,
      ShortDecimal<P2, S2>>({"plus"}, constraints);

  // (short, long) -> long
  registerFunction<
      DecimalAddFunction,
      LongDecimal<P3, S3>,
      ShortDecimal<P1, S1>,
      LongDecimal<P2, S2>>({"plus"}, constraints);

  // (long, short) -> long
  registerFunction<
      DecimalAddFunction,
      LongDecimal<P3, S3>,
      LongDecimal<P1, S1>,
      ShortDecimal<P2, S2>>({"plus"}, constraints);
```


Differential Revision: D54953663
---
 velox/core/SimpleFunctionMetadata.h           |  95 ++++++-
 velox/expression/ExprCompiler.cpp             |   2 +-
 velox/expression/SimpleFunctionAdapter.h      |  10 +-
 velox/expression/SimpleFunctionRegistry.cpp   |  53 +++-
 velox/expression/SimpleFunctionRegistry.h     |  24 +-
 velox/expression/UdfTypeResolver.h            |  14 ++
 velox/expression/VectorFunction.h             |   1 +
 velox/expression/tests/SimpleFunctionTest.cpp |   7 +-
 velox/functions/Registerer.h                  |   8 +-
 .../functions/prestosql/tests/CMakeLists.txt  |   1 +
 .../tests/DecimalSimpleFunctionTest.cpp       | 238 ++++++++++++++++++
 velox/type/Type.h                             |  40 +++
 12 files changed, 463 insertions(+), 30 deletions(-)
 create mode 100644 velox/functions/prestosql/tests/DecimalSimpleFunctionTest.cpp

diff --git a/velox/core/SimpleFunctionMetadata.h b/velox/core/SimpleFunctionMetadata.h
index a2ad56dd45f0..85b96f98f82a 100644
--- a/velox/core/SimpleFunctionMetadata.h
+++ b/velox/core/SimpleFunctionMetadata.h
@@ -242,6 +242,36 @@ struct TypeAnalysis<Generic<T, comparable, orderable>> {
   }
 };
 
+template <typename P, typename S>
+struct TypeAnalysis<ShortDecimal<P, S>> {
+  void run(TypeAnalysisResults& results) {
+    results.stats.concreteCount++;
+
+    const auto p = P::name();
+    const auto s = S::name();
+    results.out << fmt::format("decimal({},{})", p, s);
+    results.addVariable(exec::SignatureVariable(
+        p, std::nullopt, exec::ParameterType::kIntegerParameter));
+    results.addVariable(exec::SignatureVariable(
+        s, std::nullopt, exec::ParameterType::kIntegerParameter));
+  }
+};
+
+template <typename P, typename S>
+struct TypeAnalysis<LongDecimal<P, S>> {
+  void run(TypeAnalysisResults& results) {
+    results.stats.concreteCount++;
+
+    const auto p = P::name();
+    const auto s = S::name();
+    results.out << fmt::format("decimal({},{})", p, s);
+    results.addVariable(exec::SignatureVariable(
+        p, std::nullopt, exec::ParameterType::kIntegerParameter));
+    results.addVariable(exec::SignatureVariable(
+        s, std::nullopt, exec::ParameterType::kIntegerParameter));
+  }
+};
+
 template <typename K, typename V>
 struct TypeAnalysis<Map<K, V>> {
   void run(TypeAnalysisResults& results) {
@@ -329,6 +359,8 @@ class ISimpleFunctionMetadata {
   virtual bool isDeterministic() const = 0;
   virtual uint32_t priority() const = 0;
   virtual const std::shared_ptr<exec::FunctionSignature> signature() const = 0;
+  virtual TypeKind resultTypeKind() const = 0;
+  virtual const std::vector<TypeKind>& argTypeKinds() const = 0;
   virtual std::string helpMessage(const std::string& name) const = 0;
   virtual ~ISimpleFunctionMetadata() = default;
 };
@@ -407,10 +439,14 @@ class SimpleFunctionMetadata : public ISimpleFunctionMetadata {
     }
   }
 
-  explicit SimpleFunctionMetadata() {
-    auto analysis = analyzeSignatureTypes();
+  explicit SimpleFunctionMetadata(
+      const std::vector<exec::SignatureVariable>& constraints) {
+    auto analysis = analyzeSignatureTypes(constraints);
+
     buildSignature(analysis);
     priority_ = analysis.stats.computePriority();
+    resultTypeKind_ = analysis.resultTypeKind;
+    argTypeKinds_ = analysis.argTypeKinds;
   }
 
   ~SimpleFunctionMetadata() override = default;
@@ -419,6 +455,14 @@ class SimpleFunctionMetadata : public ISimpleFunctionMetadata {
     return signature_;
   }
 
+  TypeKind resultTypeKind() const override {
+    return resultTypeKind_;
+  }
+
+  const std::vector<TypeKind>& argTypeKinds() const override {
+    return argTypeKinds_;
+  }
+
   std::string helpMessage(const std::string& name) const final {
     // return fmt::format("{}({})", name, signature_->toString());
     std::string s{name};
@@ -446,15 +490,21 @@ class SimpleFunctionMetadata : public ISimpleFunctionMetadata {
     std::string outputType;
     std::map<std::string, exec::SignatureVariable> variables;
     TypeAnalysisResults::Stats stats;
+    TypeKind resultTypeKind;
+    std::vector<TypeKind> argTypeKinds;
   };
 
-  SignatureTypesAnalysisResults analyzeSignatureTypes() {
+  SignatureTypesAnalysisResults analyzeSignatureTypes(
+      const std::vector<exec::SignatureVariable>& constraints) {
     std::vector<std::string> argsTypes;
 
     TypeAnalysisResults results;
     TypeAnalysis<return_type>().run(results);
     std::string outputType = results.typeAsString();
 
+    auto resultTypeKind = SimpleTypeTrait<return_type>::typeKind;
+    std::vector<TypeKind> argTypeKinds;
+
     (
         [&]() {
           // Clear string representation but keep other collected information
@@ -462,14 +512,30 @@ class SimpleFunctionMetadata : public ISimpleFunctionMetadata {
           results.resetTypeString();
           TypeAnalysis<Args>().run(results);
           argsTypes.push_back(results.typeAsString());
+
+          if constexpr (!isVariadicType<Args>::value) {
+            argTypeKinds.push_back(SimpleTypeTrait<Args>::typeKind);
+          }
         }(),
         ...);
 
+    for (const auto& constraint : constraints) {
+      VELOX_CHECK(
+          !constraint.constraint().empty(),
+          "Constraint must be set for variable {}",
+          constraint.name());
+
+      results.variablesInformation.erase(constraint.name());
+      results.variablesInformation.emplace(constraint.name(), constraint);
+    }
+
     return SignatureTypesAnalysisResults{
         std::move(argsTypes),
         std::move(outputType),
         std::move(results.variablesInformation),
-        std::move(results.stats)};
+        std::move(results.stats),
+        resultTypeKind,
+        argTypeKinds};
   }
 
   void buildSignature(const SignatureTypesAnalysisResults& analysis) {
@@ -497,6 +563,8 @@ class SimpleFunctionMetadata : public ISimpleFunctionMetadata {
 
   exec::FunctionSignaturePtr signature_;
   uint32_t priority_;
+  TypeKind resultTypeKind_;
+  std::vector<TypeKind> argTypeKinds_;
 };
 
 // wraps a UDF object to provide the inheritance
@@ -544,6 +612,7 @@ class UDFHolder final
   DECLARE_METHOD_RESOLVER(callNullFree_method_resolver, callNullFree);
   DECLARE_METHOD_RESOLVER(callAscii_method_resolver, callAscii);
   DECLARE_METHOD_RESOLVER(initialize_method_resolver, initialize);
+  DECLARE_METHOD_RESOLVER(initializeTypes_method_resolver, initializeTypes);
 
   // Check which flavor of the call() method is provided by the UDF object. UDFs
   // are required to provide at least one of the following methods:
@@ -650,6 +719,13 @@ class UDFHolder final
       const core::QueryConfig&,
       const exec_arg_type<TArgs>*...>::value;
 
+  // initializeTypes():
+  static constexpr bool udf_has_initializeTypes = util::has_method<
+      Fun,
+      initializeTypes_method_resolver,
+      void,
+      const std::vector<TypePtr>&>::value;
+
   static_assert(
       udf_has_call || udf_has_callNullable || udf_has_callNullFree,
       "UDF must implement at least one of `call`, `callNullable`, or `callNullFree` functions.\n"
@@ -703,7 +779,9 @@ class UDFHolder final
   template <size_t N>
   using exec_type_at = typename std::tuple_element<N, exec_arg_types>::type;
 
-  explicit UDFHolder() : Metadata(), instance_{} {}
+  explicit UDFHolder(
+      const std::vector<exec::SignatureVariable>& constraints = {})
+      : Metadata(constraints), instance_{} {}
 
   FOLLY_ALWAYS_INLINE void initialize(
       const core::QueryConfig& config,
@@ -713,6 +791,13 @@ class UDFHolder final
     }
   }
 
+  FOLLY_ALWAYS_INLINE void initializeTypes(
+      const std::vector<TypePtr>& argTypes) {
+    if constexpr (udf_has_initializeTypes) {
+      return instance_.initializeTypes(argTypes);
+    }
+  }
+
   FOLLY_ALWAYS_INLINE bool call(
       exec_return_type& out,
       const typename exec_resolver<TArgs>::in_type&... args) {
diff --git a/velox/expression/ExprCompiler.cpp b/velox/expression/ExprCompiler.cpp
index 93f6080ee027..df05e7fd5cd3 100644
--- a/velox/expression/ExprCompiler.cpp
+++ b/velox/expression/ExprCompiler.cpp
@@ -435,7 +435,7 @@ ExprPtr compileRewrittenExpression(
           resultType,
           folly::join(", ", inputTypes));
       auto func_2 = simpleFunctionEntry->createFunction()->createVectorFunction(
-          getConstantInputs(compiledInputs), config);
+          inputTypes, getConstantInputs(compiledInputs), config);
       result = std::make_shared<Expr>(
           resultType,
           std::move(compiledInputs),
diff --git a/velox/expression/SimpleFunctionAdapter.h b/velox/expression/SimpleFunctionAdapter.h
index 784576c28d5c..af06c780c7fb 100644
--- a/velox/expression/SimpleFunctionAdapter.h
+++ b/velox/expression/SimpleFunctionAdapter.h
@@ -228,9 +228,10 @@ class SimpleFunctionAdapter : public VectorFunction {
 
  public:
   explicit SimpleFunctionAdapter(
+      const std::vector<TypePtr>& inputTypes,
       const core::QueryConfig& config,
       const std::vector<VectorPtr>& constantInputs)
-      : fn_{std::make_unique<FUNC>()} {
+      : fn_{std::make_unique<FUNC>(std::vector<exec::SignatureVariable>{})} {
     if constexpr (FUNC::udf_has_initialize) {
       try {
         unpackInitialize<0>(config, constantInputs);
@@ -240,6 +241,10 @@ class SimpleFunctionAdapter : public VectorFunction {
         initializeException_ = std::current_exception();
       }
     }
+
+    if constexpr (FUNC::udf_has_initializeTypes) {
+      (*fn_).initializeTypes(inputTypes);
+    }
   }
 
   explicit SimpleFunctionAdapter() {}
@@ -901,10 +906,11 @@ class SimpleFunctionAdapterFactoryImpl : public SimpleFunctionAdapterFactory {
   explicit SimpleFunctionAdapterFactoryImpl() {}
 
   std::unique_ptr<VectorFunction> createVectorFunction(
+      const std::vector<TypePtr>& inputTypes,
       const std::vector<VectorPtr>& constantInputs,
       const core::QueryConfig& config) const override {
     return std::make_unique<SimpleFunctionAdapter<UDFHolder>>(
-        config, constantInputs);
+        inputTypes, config, constantInputs);
   }
 };
 
diff --git a/velox/expression/SimpleFunctionRegistry.cpp b/velox/expression/SimpleFunctionRegistry.cpp
index 36c6849f48b4..a3aff7b6e053 100644
--- a/velox/expression/SimpleFunctionRegistry.cpp
+++ b/velox/expression/SimpleFunctionRegistry.cpp
@@ -40,8 +40,9 @@ void SimpleFunctionRegistry::registerFunctionInternal(
   const auto sanitizedName = sanitizeName(name);
   registeredFunctions_.withWLock([&](auto& map) {
     SignatureMap& signatureMap = map[sanitizedName];
-    signatureMap[*metadata->signature()] =
-        std::make_unique<const FunctionEntry>(metadata, factory);
+    // TODO Check type kinds and avoid adding duplicates.
+    signatureMap[*metadata->signature()].emplace_back(
+        std::make_unique<const FunctionEntry>(metadata, factory));
   });
 }
 
@@ -83,12 +84,48 @@ SimpleFunctionRegistry::resolveFunction(
       for (const auto& [candidateSignature, functionEntry] : *signatureMap) {
         SignatureBinder binder(candidateSignature, argTypes);
         if (binder.tryBind()) {
-          auto* currentCandidate = functionEntry.get();
-          if (!selectedCandidate ||
-              currentCandidate->getMetadata().priority() <
-                  selectedCandidate->getMetadata().priority()) {
-            selectedCandidate = currentCandidate;
-            selectedCandidateType = binder.tryResolveReturnType();
+          for (const auto& currentCandidate : functionEntry) {
+            const auto& m = currentCandidate->getMetadata();
+
+            // Check that TypeKinds of arguments match.
+            bool match = true;
+            for (auto i = 0; i < m.argTypeKinds().size(); ++i) {
+              const auto typeKind = m.argTypeKinds()[i];
+
+              // Generic types do not specify TypeKind. Skip the check.
+              if (typeKind == TypeKind::UNKNOWN) {
+                continue;
+              }
+
+              if (argTypes[i]->kind() != typeKind) {
+                LOG(ERROR) << "Type mismatch for argument " << i << " of "
+                           << name << ": " << argTypes[i]->toString() << " vs "
+                           << mapTypeKindToName(typeKind);
+                match = false;
+              }
+            }
+
+            if (!match) {
+              continue;
+            }
+
+            if (!selectedCandidate ||
+                currentCandidate->getMetadata().priority() <
+                    selectedCandidate->getMetadata().priority()) {
+              auto resultType = binder.tryResolveReturnType();
+              VELOX_CHECK_NOT_NULL(resultType);
+
+              if (m.resultTypeKind() != TypeKind::UNKNOWN &&
+                  resultType->kind() != m.resultTypeKind()) {
+                LOG(ERROR) << "Type mismatch for result of " << name << ": "
+                           << resultType->toString() << " vs "
+                           << mapTypeKindToName(m.resultTypeKind());
+                continue;
+              }
+
+              selectedCandidate = currentCandidate.get();
+              selectedCandidateType = resultType;
+            }
           }
         }
       }
diff --git a/velox/expression/SimpleFunctionRegistry.h b/velox/expression/SimpleFunctionRegistry.h
index 2d1e1a431c83..a334b79fe393 100644
--- a/velox/expression/SimpleFunctionRegistry.h
+++ b/velox/expression/SimpleFunctionRegistry.h
@@ -24,8 +24,9 @@
 namespace facebook::velox::exec {
 
 template <typename T>
-const std::shared_ptr<const T>& singletonUdfMetadata() {
-  static auto instance = std::make_shared<const T>();
+const std::shared_ptr<const T>& singletonUdfMetadata(
+    const std::vector<exec::SignatureVariable>& constraints = {}) {
+  static auto instance = std::make_shared<const T>(constraints);
   return instance;
 }
 
@@ -52,15 +53,19 @@ struct FunctionEntry {
   const FunctionFactory factory_;
 };
 
-using SignatureMap =
-    std::unordered_map<FunctionSignature, std::unique_ptr<const FunctionEntry>>;
+using SignatureMap = std::unordered_map<
+    FunctionSignature,
+    std::vector<std::unique_ptr<const FunctionEntry>>>;
 using FunctionMap = std::unordered_map<std::string, SignatureMap>;
 
 class SimpleFunctionRegistry {
  public:
   template <typename UDF>
-  void registerFunction(const std::vector<std::string>& aliases = {}) {
-    const auto& metadata = singletonUdfMetadata<typename UDF::Metadata>();
+  void registerFunction(
+      const std::vector<std::string>& aliases,
+      const std::vector<exec::SignatureVariable>& constraints) {
+    const auto& metadata =
+        singletonUdfMetadata<typename UDF::Metadata>(constraints);
     const auto factory = [metadata]() { return CreateUdf<UDF>(); };
 
     if (aliases.empty()) {
@@ -139,9 +144,12 @@ SimpleFunctionRegistry& mutableSimpleFunctions();
 
 // This function should be called once and alone.
 template <typename UDFHolder>
-void registerSimpleFunction(const std::vector<std::string>& names) {
+void registerSimpleFunction(
+    const std::vector<std::string>& names,
+    const std::vector<exec::SignatureVariable>& constraints) {
   mutableSimpleFunctions()
-      .registerFunction<SimpleFunctionAdapterFactoryImpl<UDFHolder>>(names);
+      .registerFunction<SimpleFunctionAdapterFactoryImpl<UDFHolder>>(
+          names, constraints);
 }
 
 } // namespace facebook::velox::exec
diff --git a/velox/expression/UdfTypeResolver.h b/velox/expression/UdfTypeResolver.h
index 8d188ef97c64..827be7ab4a37 100644
--- a/velox/expression/UdfTypeResolver.h
+++ b/velox/expression/UdfTypeResolver.h
@@ -101,6 +101,20 @@ struct resolver<Array<V>> {
   using out_type = ArrayWriter<V>;
 };
 
+template <typename P, typename S>
+struct resolver<ShortDecimal<P, S>> {
+  using in_type = int64_t;
+  using null_free_in_type = in_type;
+  using out_type = int64_t;
+};
+
+template <typename P, typename S>
+struct resolver<LongDecimal<P, S>> {
+  using in_type = int128_t;
+  using null_free_in_type = in_type;
+  using out_type = int128_t;
+};
+
 template <>
 struct resolver<Varchar> {
   using in_type = StringView;
diff --git a/velox/expression/VectorFunction.h b/velox/expression/VectorFunction.h
index 6442cfd7b502..323606acf44f 100644
--- a/velox/expression/VectorFunction.h
+++ b/velox/expression/VectorFunction.h
@@ -182,6 +182,7 @@ class ApplyNeverCalled final : public VectorFunction {
 class SimpleFunctionAdapterFactory {
  public:
   virtual std::unique_ptr<VectorFunction> createVectorFunction(
+      const std::vector<TypePtr>& inputTypes,
       const std::vector<VectorPtr>& constantInputs,
       const core::QueryConfig& config) const = 0;
   virtual ~SimpleFunctionAdapterFactory() = default;
diff --git a/velox/expression/tests/SimpleFunctionTest.cpp b/velox/expression/tests/SimpleFunctionTest.cpp
index a77c3de60eca..bc3d17b8c204 100644
--- a/velox/expression/tests/SimpleFunctionTest.cpp
+++ b/velox/expression/tests/SimpleFunctionTest.cpp
@@ -842,7 +842,7 @@ TEST_F(SimpleFunctionTest, stringReuseConstant) {
   // Test reusing the strings from an argument when that argument is in a
   // ConstantVector.  Note that the other 2 arguments are FlatVectors to
   // prevent constant peeling.
-  registerFunction<Substr, Varchar, Varchar, int32_t, int32_t>({"substr"});
+  registerFunction<Substr, Varchar, Varchar, int32_t, int32_t>({"test_substr"});
 
   auto constantVector = vectorMaker_.constantVector<StringView>(
       {"super happy fun string"_sv,
@@ -852,7 +852,8 @@ TEST_F(SimpleFunctionTest, stringReuseConstant) {
   auto lengths = vectorMaker_.flatVector({1, 2, 3});
 
   auto result = evaluate<FlatVector<StringView>>(
-      "substr(c0, c1, c2)", makeRowVector({constantVector, starts, lengths}));
+      "test_substr(c0, c1, c2)",
+      makeRowVector({constantVector, starts, lengths}));
 
   auto expected = vectorMaker_.flatVector({"s"_sv, "up"_sv, "per"_sv});
   assertEqualVectors(expected, result);
@@ -992,7 +993,7 @@ VectorPtr testVariadicArgReuse(
       exec::simpleFunctions()
           .resolveFunction(functionName, {})
           ->createFunction()
-          ->createVectorFunction({}, execCtx->queryCtx()->queryConfig());
+          ->createVectorFunction({}, {}, execCtx->queryCtx()->queryConfig());
 
   // Create a dummy EvalCtx.
   SelectivityVector rows(inputs[0]->size());
diff --git a/velox/functions/Registerer.h b/velox/functions/Registerer.h
index 7b893c6c3928..52e4fbf61b29 100644
--- a/velox/functions/Registerer.h
+++ b/velox/functions/Registerer.h
@@ -43,7 +43,7 @@ void registerFunction(const std::vector<std::string>& aliases = {}) {
       TReturn,
       ConstantChecker<TArgs...>,
       typename UnwrapConstantType<TArgs>::type...>;
-  exec::registerSimpleFunction<holderClass>(aliases);
+  exec::registerSimpleFunction<holderClass>(aliases, {});
 }
 
 // New registration function; mostly a copy from the function above, but taking
@@ -51,7 +51,9 @@ void registerFunction(const std::vector<std::string>& aliases = {}) {
 // a while to maintain backwards compatibility, but the idea is to remove the
 // one above eventually.
 template <template <class> typename Func, typename TReturn, typename... TArgs>
-void registerFunction(const std::vector<std::string>& aliases = {}) {
+void registerFunction(
+    const std::vector<std::string>& aliases = {},
+    const std::vector<exec::SignatureVariable>& constraints = {}) {
   using funcClass = Func<exec::VectorExec>;
   using holderClass = core::UDFHolder<
       funcClass,
@@ -59,7 +61,7 @@ void registerFunction(const std::vector<std::string>& aliases = {}) {
       TReturn,
       ConstantChecker<TArgs...>,
       typename UnwrapConstantType<TArgs>::type...>;
-  exec::registerSimpleFunction<holderClass>(aliases);
+  exec::registerSimpleFunction<holderClass>(aliases, constraints);
 }
 
 } // namespace facebook::velox
diff --git a/velox/functions/prestosql/tests/CMakeLists.txt b/velox/functions/prestosql/tests/CMakeLists.txt
index d0687a849ea0..7f1eccc0be18 100644
--- a/velox/functions/prestosql/tests/CMakeLists.txt
+++ b/velox/functions/prestosql/tests/CMakeLists.txt
@@ -54,6 +54,7 @@ add_executable(
   ComparisonsTest.cpp
   DateTimeFunctionsTest.cpp
   DecimalArithmeticTest.cpp
+  DecimalSimpleFunctionTest.cpp
   ElementAtTest.cpp
   FindFirstTest.cpp
   FromUtf8Test.cpp
diff --git a/velox/functions/prestosql/tests/DecimalSimpleFunctionTest.cpp b/velox/functions/prestosql/tests/DecimalSimpleFunctionTest.cpp
new file mode 100644
index 000000000000..93e00b701a88
--- /dev/null
+++ b/velox/functions/prestosql/tests/DecimalSimpleFunctionTest.cpp
@@ -0,0 +1,238 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cstdint>
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/functions/Macros.h"
+#include "velox/functions/Registerer.h"
+#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
+
+using namespace facebook::velox;
+using namespace facebook::velox::test;
+using namespace facebook::velox::functions::test;
+
+namespace facebook::velox {
+namespace {
+
+class Addition {
+ public:
+  template <typename R, typename A, typename B>
+  inline static void
+  apply(R& r, const A& a, const B& b, uint8_t aRescale, uint8_t bRescale)
+#if defined(__has_feature)
+#if __has_feature(__address_sanitizer__)
+      __attribute__((__no_sanitize__("signed-integer-overflow")))
+#endif
+#endif
+  {
+    int128_t aRescaled;
+    int128_t bRescaled;
+    if (__builtin_mul_overflow(
+            a, DecimalUtil::kPowersOfTen[aRescale], &aRescaled) ||
+        __builtin_mul_overflow(
+            b, DecimalUtil::kPowersOfTen[bRescale], &bRescaled)) {
+      VELOX_ARITHMETIC_ERROR("Decimal overflow: {} + {}", a, b);
+    }
+    r = checkedPlus<R>(R(aRescaled), R(bRescaled));
+    DecimalUtil::valueInRange(r);
+  }
+
+  inline static uint8_t
+  computeRescaleFactor(uint8_t fromScale, uint8_t toScale, uint8_t rScale = 0) {
+    return std::max(0, toScale - fromScale);
+  }
+
+  inline static std::pair<uint8_t, uint8_t> computeResultPrecisionScale(
+      const uint8_t aPrecision,
+      const uint8_t aScale,
+      const uint8_t bPrecision,
+      const uint8_t bScale) {
+    return {
+        std::min(
+            38,
+            std::max(aPrecision - aScale, bPrecision - bScale) +
+                std::max(aScale, bScale) + 1),
+        std::max(aScale, bScale)};
+  }
+};
+
+template <typename TExec>
+struct DecimalAddFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(TExec);
+
+  void initializeTypes(const std::vector<TypePtr>& argTypes) {
+    auto aType = argTypes[0];
+    auto bType = argTypes[1];
+    auto [aPrecision, aScale] = getDecimalPrecisionScale(*aType);
+    auto [bPrecision, bScale] = getDecimalPrecisionScale(*bType);
+    auto [rPrecision, rScale] = Addition::computeResultPrecisionScale(
+        aPrecision, aScale, bPrecision, bScale);
+    aRescale_ = Addition::computeRescaleFactor(aScale, bScale, rScale);
+    bRescale_ = Addition::computeRescaleFactor(bScale, aScale, rScale);
+  }
+
+  template <typename R, typename A, typename B>
+  void call(R& out, const A& a, const B& b) {
+    Addition::template apply<R, A, B>(out, a, b, aRescale_, bRescale_);
+  }
+
+ private:
+  uint8_t aRescale_;
+  uint8_t bRescale_;
+};
+
+class DecimalSimpleFunctionTest : public FunctionBaseTest {
+ protected:
+  DecimalSimpleFunctionTest() {
+    options_.parseDecimalAsDouble = false;
+  }
+
+  template <TypeKind KIND>
+  void testDecimalExpr(
+      const VectorPtr& expected,
+      const std::string& expression,
+      const std::vector<VectorPtr>& input) {
+    using EvalType = typename velox::TypeTraits<KIND>::NativeType;
+    auto result =
+        evaluate<SimpleVector<EvalType>>(expression, makeRowVector(input));
+    assertEqualVectors(expected, result);
+    testOpDictVectors<EvalType>(expression, expected, input);
+  }
+
+  template <typename T>
+  void testOpDictVectors(
+      const std::string& operation,
+      const VectorPtr& expected,
+      const std::vector<VectorPtr>& flatVector) {
+    // Dictionary vectors as arguments.
+    auto newSize = flatVector[0]->size() * 2;
+    std::vector<VectorPtr> dictVectors;
+    for (auto i = 0; i < flatVector.size(); ++i) {
+      auto indices = makeIndices(newSize, [&](int row) { return row / 2; });
+      dictVectors.push_back(
+          VectorTestBase::wrapInDictionary(indices, newSize, flatVector[i]));
+    }
+    auto resultIndices = makeIndices(newSize, [&](int row) { return row / 2; });
+    auto expectedResultDictionary =
+        VectorTestBase::wrapInDictionary(resultIndices, newSize, expected);
+    auto actual =
+        evaluate<SimpleVector<T>>(operation, makeRowVector(dictVectors));
+    assertEqualVectors(expectedResultDictionary, actual);
+  }
+};
+
+TEST_F(DecimalSimpleFunctionTest, add) {
+  std::vector<exec::SignatureVariable> constraints = {
+      exec::SignatureVariable(
+          P3::name(),
+          fmt::format(
+              "min(38, max({a_precision} - {a_scale}, {b_precision} - {b_scale}) + max({a_scale}, {b_scale}) + 1)",
+              fmt::arg("a_precision", P1::name()),
+              fmt::arg("b_precision", P2::name()),
+              fmt::arg("a_scale", S1::name()),
+              fmt::arg("b_scale", S2::name())),
+          exec::ParameterType::kIntegerParameter),
+      exec::SignatureVariable(
+          S3::name(),
+          fmt::format(
+              "max({a_scale}, {b_scale})",
+              fmt::arg("a_scale", S1::name()),
+              fmt::arg("b_scale", S2::name())),
+          exec::ParameterType::kIntegerParameter),
+  };
+
+  registerFunction<
+      DecimalAddFunction,
+      LongDecimal<P3, S3>,
+      LongDecimal<P1, S1>,
+      LongDecimal<P2, S2>>({"foo"}, constraints);
+
+  registerFunction<
+      DecimalAddFunction,
+      ShortDecimal<P3, S3>,
+      ShortDecimal<P1, S1>,
+      ShortDecimal<P2, S2>>({"foo"}, constraints);
+
+  registerFunction<
+      DecimalAddFunction,
+      LongDecimal<P3, S3>,
+      ShortDecimal<P1, S1>,
+      ShortDecimal<P2, S2>>({"foo"}, constraints);
+
+  registerFunction<
+      DecimalAddFunction,
+      LongDecimal<P3, S3>,
+      ShortDecimal<P1, S1>,
+      LongDecimal<P2, S2>>({"foo"}, constraints);
+
+  registerFunction<
+      DecimalAddFunction,
+      LongDecimal<P3, S3>,
+      LongDecimal<P1, S1>,
+      ShortDecimal<P2, S2>>({"foo"}, constraints);
+
+  auto shortFlat = makeFlatVector<int64_t>({1000, 2000}, DECIMAL(18, 3));
+  // Add short and short, returning long.
+  testDecimalExpr<TypeKind::HUGEINT>(
+      makeFlatVector<int128_t>({2000, 4000}, DECIMAL(19, 3)),
+      "foo(c0, c1)",
+      {shortFlat, shortFlat});
+
+  // Add short and long, returning long.
+  auto longFlat = makeFlatVector<int128_t>({1000, 2000}, DECIMAL(19, 3));
+  auto expectedLongFlat =
+      makeFlatVector<int128_t>({2000, 4000}, DECIMAL(20, 3));
+  testDecimalExpr<TypeKind::HUGEINT>(
+      expectedLongFlat, "foo(c0, c1)", {shortFlat, longFlat});
+
+  // Add short and long, returning long.
+  testDecimalExpr<TypeKind::HUGEINT>(
+      expectedLongFlat, "foo(c0, c1)", {longFlat, shortFlat});
+
+  // Add long and long, returning long.
+  testDecimalExpr<TypeKind::HUGEINT>(
+      expectedLongFlat, "c0 + c1", {longFlat, longFlat});
+
+  // Add short and short, returning short.
+  shortFlat = makeFlatVector<int64_t>({1000, 2000}, DECIMAL(10, 3));
+  auto expectedShortFlat =
+      makeFlatVector<int64_t>({2000, 4000}, DECIMAL(11, 3));
+  testDecimalExpr<TypeKind::BIGINT>(
+      expectedShortFlat, "c0 + c1", {shortFlat, shortFlat});
+
+  auto expectedConstantFlat =
+      makeFlatVector<int64_t>({2000, 3000}, DECIMAL(11, 3));
+
+  // Constant and Flat arguments.
+  testDecimalExpr<TypeKind::BIGINT>(
+      expectedConstantFlat, "foo(1.00, c0)", {shortFlat});
+
+  // Flat and Constant arguments.
+  testDecimalExpr<TypeKind::BIGINT>(
+      expectedConstantFlat, "foo(c0,1.00)", {shortFlat});
+
+  testDecimalExpr<TypeKind::BIGINT>(
+      makeNullableFlatVector<int64_t>(
+          {2, 4, std::nullopt, std::nullopt, std::nullopt}, DECIMAL(11, 3)),
+      "foo(c0, c1)",
+      {makeNullableFlatVector<int64_t>(
+           {1, 2, std::nullopt, 6, std::nullopt}, DECIMAL(10, 3)),
+       makeNullableFlatVector<int64_t>(
+           {1, 2, 5, std::nullopt, std::nullopt}, DECIMAL(10, 3))});
+}
+
+} // namespace
+} // namespace facebook::velox
diff --git a/velox/type/Type.h b/velox/type/Type.h
index 5670d1f34b63..8167fffdc853 100644
--- a/velox/type/Type.h
+++ b/velox/type/Type.h
@@ -1741,6 +1741,38 @@ using T6 = TypeVariable<6>;
 using T7 = TypeVariable<7>;
 using T8 = TypeVariable<8>;
 
+template <size_t id>
+struct IntegerVariable {
+  static size_t getId() {
+    return id;
+  }
+
+  static std::string name() {
+    return fmt::format("i{}", id);
+  }
+};
+
+using P1 = IntegerVariable<1>;
+using P2 = IntegerVariable<2>;
+using P3 = IntegerVariable<3>;
+using P4 = IntegerVariable<4>;
+using S1 = IntegerVariable<5>;
+using S2 = IntegerVariable<6>;
+using S3 = IntegerVariable<7>;
+using S4 = IntegerVariable<8>;
+
+template <typename P, typename S>
+struct ShortDecimal {
+ private:
+  ShortDecimal() {}
+};
+
+template <typename P, typename S>
+struct LongDecimal {
+ private:
+  LongDecimal() {}
+};
+
 struct AnyType {};
 
 template <typename T = AnyType, bool comparable = false, bool orderable = false>
@@ -1933,6 +1965,14 @@ struct SimpleTypeTrait<int16_t> : public TypeTraits<TypeKind::SMALLINT> {};
 template <>
 struct SimpleTypeTrait<int8_t> : public TypeTraits<TypeKind::TINYINT> {};
 
+template <typename P, typename S>
+struct SimpleTypeTrait<ShortDecimal<P, S>>
+    : public TypeTraits<TypeKind::BIGINT> {};
+
+template <typename P, typename S>
+struct SimpleTypeTrait<LongDecimal<P, S>>
+    : public TypeTraits<TypeKind::HUGEINT> {};
+
 template <>
 struct SimpleTypeTrait<float> : public TypeTraits<TypeKind::REAL> {};