Pass input types to simple function's 'initialize' method (#9124)

Summary: This change is part of enabling simple functions to process inputs of decimal type. Such processing requires access to decimal type parameters (precision and scale). This change provides full type information via 'initialize' method. See #9096 for the end-to-end workflow. Reviewed By: xiaoxmeng Differential Revision: D55012189
facebookincubator · Mar 19, 2024 · cd6ea36 · cd6ea36
1 parent ac7f58f
commit cd6ea36
Show file tree

Hide file tree

Showing 20 changed files with 82 additions and 15 deletions.
diff --git a/velox/core/SimpleFunctionMetadata.h b/velox/core/SimpleFunctionMetadata.h
@@ -647,9 +647,22 @@ class UDFHolder final
       Fun,
       initialize_method_resolver,
       void,
+      const std::vector<TypePtr>&,
       const core::QueryConfig&,
       const exec_arg_type<TArgs>*...>::value;
 
+  // TODO Remove
+  static constexpr bool udf_has_legacy_initialize = util::has_method<
+      Fun,
+      initialize_method_resolver,
+      void,
+      const core::QueryConfig&,
+      const exec_arg_type<TArgs>*...>::value;
+
+  static_assert(
+      !udf_has_legacy_initialize,
+      "Legacy initialize method! Upgrade.");
+
   static_assert(
       udf_has_call || udf_has_callNullable || udf_has_callNullFree,
       "UDF must implement at least one of `call`, `callNullable`, or `callNullFree` functions.\n"
@@ -706,10 +719,11 @@ class UDFHolder final
   explicit UDFHolder() : Metadata(), instance_{} {}
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& inputTypes,
       const core::QueryConfig& config,
       const typename exec_resolver<TArgs>::in_type*... constantArgs) {
     if constexpr (udf_has_initialize) {
-      return instance_.initialize(config, constantArgs...);
+      return instance_.initialize(inputTypes, config, constantArgs...);
     }
   }
 

diff --git a/velox/docs/develop/scalar-functions.rst b/velox/docs/develop/scalar-functions.rst
@@ -335,6 +335,7 @@ properties and using it when processing inputs.
     const date::time_zone* timeZone_ = nullptr;
 
     FOLLY_ALWAYS_INLINE void initialize(
+        const std::vector<TypePtr>& inputTypes,
         const core::QueryConfig& config,
         const arg_type<Timestamp>* /*timestamp*/) {
       timeZone_ = getTimeZoneFromConfig(config);
@@ -365,6 +366,7 @@ individual rows.
     std::optional<DateTimeUnit> unit_;
 
     FOLLY_ALWAYS_INLINE void initialize(
+        const std::vector<TypePtr>& inputTypes,
         const core::QueryConfig& config,
         const arg_type<Varchar>* unitString,
         const arg_type<Timestamp>* /*timestamp*/) {

diff --git a/velox/examples/SimpleFunctions.cpp b/velox/examples/SimpleFunctions.cpp
@@ -331,6 +331,7 @@ struct MyRegexpMatchFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig&,
       const arg_type<Varchar>*,
       const arg_type<Varchar>* pattern) {

diff --git a/velox/expression/SimpleFunctionAdapter.h b/velox/expression/SimpleFunctionAdapter.h
@@ -205,36 +205,38 @@ class SimpleFunctionAdapter : public VectorFunction {
 
   template <int32_t POSITION, typename... Values>
   void unpackInitialize(
+      const std::vector<TypePtr>& inputTypes,
       const core::QueryConfig& config,
       const std::vector<VectorPtr>& packed,
       const Values*... values) const {
     if constexpr (POSITION == FUNC::num_args) {
-      return (*fn_).initialize(config, values...);
+      return (*fn_).initialize(inputTypes, config, values...);
     } else {
       if (packed.at(POSITION) != nullptr) {
         SelectivityVector rows(1);
         DecodedVector decodedVector(*packed.at(POSITION), rows);
         auto oneReader = VectorReader<arg_at<POSITION>>(&decodedVector);
         auto oneValue = oneReader[0];
 
-        unpackInitialize<POSITION + 1>(config, packed, values..., &oneValue);
+        unpackInitialize<POSITION + 1>(
+            inputTypes, config, packed, values..., &oneValue);
       } else {
         using temp_type = exec_arg_at<POSITION>;
         unpackInitialize<POSITION + 1>(
-            config, packed, values..., (const temp_type*)nullptr);
+            inputTypes, config, packed, values..., (const temp_type*)nullptr);
       }
     }
   }
 
  public:
   SimpleFunctionAdapter(
-      const std::vector<TypePtr>& /*inputTypes*/,
+      const std::vector<TypePtr>& inputTypes,
       const core::QueryConfig& config,
       const std::vector<VectorPtr>& constantInputs)
       : fn_{std::make_unique<FUNC>()} {
     if constexpr (FUNC::udf_has_initialize) {
       try {
-        unpackInitialize<0>(config, constantInputs);
+        unpackInitialize<0>(inputTypes, config, constantInputs);
       } catch (const VeloxRuntimeError&) {
         throw;
       } catch (const std::exception&) {

diff --git a/velox/expression/tests/SimpleFunctionInitTest.cpp b/velox/expression/tests/SimpleFunctionInitTest.cpp
@@ -37,6 +37,7 @@ struct NonDefaultWithArrayInitFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<int32_t>* /*first*/,
       const arg_type<velox::Array<int32_t>>* second) {
@@ -131,6 +132,7 @@ struct NonDefaultWithMapInitFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<int32_t>* /*first*/,
       const arg_type<velox::Map<int32_t, int64_t>>* second) {
@@ -200,6 +202,7 @@ struct InitAlwaysThrowsFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<int32_t>* /*first*/) {
     VELOX_USER_FAIL("Unconditional throw!");

diff --git a/velox/expression/tests/SimpleFunctionTest.cpp b/velox/expression/tests/SimpleFunctionTest.cpp
@@ -1248,6 +1248,7 @@ struct ConstantArgumentFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<int32_t>* /*first*/,
       const arg_type<int32_t>* /*second*/,

diff --git a/velox/functions/lib/Re2Functions.h b/velox/functions/lib/Re2Functions.h
@@ -260,6 +260,7 @@ struct Re2RegexpReplace {
   std::optional<RE2> re_;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* /*string*/,
       const arg_type<Varchar>* pattern,
@@ -282,12 +283,13 @@ struct Re2RegexpReplace {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& inputTypes,
       const core::QueryConfig& config,
       const arg_type<Varchar>* string,
       const arg_type<Varchar>* pattern) {
     StringView emptyReplacement;
 
-    initialize(config, string, pattern, &emptyReplacement);
+    initialize(inputTypes, config, string, pattern, &emptyReplacement);
   }
 
   FOLLY_ALWAYS_INLINE bool call(

diff --git a/velox/functions/lib/TimeUtils.h b/velox/functions/lib/TimeUtils.h
@@ -98,6 +98,7 @@ struct InitSessionTimezone {
   const date::time_zone* timeZone_{nullptr};
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Timestamp>* /*timestamp*/) {
     timeZone_ = getTimeZoneFromConfig(config);

diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -106,18 +106,21 @@ struct DateFunction : public TimestampWithTimezoneSupport<T> {
   const date::time_zone* timeZone_ = nullptr;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* date) {
     timeZone_ = getTimeZoneFromConfig(config);
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Timestamp>* timestamp) {
     timeZone_ = getTimeZoneFromConfig(config);
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<TimestampWithTimezone>* timestampWithTimezone) {
     timeZone_ = getTimeZoneFromConfig(config);
@@ -820,6 +823,7 @@ struct DateTruncFunction : public TimestampWithTimezoneSupport<T> {
   std::optional<DateTimeUnit> unit_;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* unitString,
       const arg_type<Timestamp>* /*timestamp*/) {
@@ -831,6 +835,7 @@ struct DateTruncFunction : public TimestampWithTimezoneSupport<T> {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Varchar>* unitString,
       const arg_type<Date>* /*date*/) {
@@ -840,6 +845,7 @@ struct DateTruncFunction : public TimestampWithTimezoneSupport<T> {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Varchar>* unitString,
       const arg_type<TimestampWithTimezone>* /*timestamp*/) {
@@ -997,6 +1003,7 @@ struct DateAddFunction : public TimestampWithTimezoneSupport<T> {
   std::optional<DateTimeUnit> unit_ = std::nullopt;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* unitString,
       const int64_t* /*value*/,
@@ -1008,6 +1015,7 @@ struct DateAddFunction : public TimestampWithTimezoneSupport<T> {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Varchar>* unitString,
       const int64_t* /*value*/,
@@ -1103,6 +1111,7 @@ struct DateDiffFunction : public TimestampWithTimezoneSupport<T> {
   std::optional<DateTimeUnit> unit_ = std::nullopt;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* unitString,
       const arg_type<Timestamp>* /*timestamp1*/,
@@ -1115,6 +1124,7 @@ struct DateDiffFunction : public TimestampWithTimezoneSupport<T> {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Varchar>* unitString,
       const arg_type<Date>* /*date1*/,
@@ -1125,6 +1135,7 @@ struct DateDiffFunction : public TimestampWithTimezoneSupport<T> {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* unitString,
       const arg_type<TimestampWithTimezone>* /*timestampWithTimezone1*/,
@@ -1195,6 +1206,7 @@ struct DateFormatFunction : public TimestampWithTimezoneSupport<T> {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Timestamp>* /*timestamp*/,
       const arg_type<Varchar>* formatString) {
@@ -1206,6 +1218,7 @@ struct DateFormatFunction : public TimestampWithTimezoneSupport<T> {
   }
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<TimestampWithTimezone>* /*timestamp*/,
       const arg_type<Varchar>* formatString) {
@@ -1271,6 +1284,7 @@ struct DateParseFunction {
   bool isConstFormat_ = false;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* /*input*/,
       const arg_type<Varchar>* formatString) {
@@ -1313,6 +1327,7 @@ struct FormatDateTimeFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Timestamp>* /*timestamp*/,
       const arg_type<Varchar>* formatString) {
@@ -1383,6 +1398,7 @@ struct ParseDateTimeFunction {
   bool isConstFormat_ = false;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* /*input*/,
       const arg_type<Varchar>* format) {
@@ -1427,7 +1443,9 @@ struct CurrentDateFunction {
 
   const date::time_zone* timeZone_ = nullptr;
 
-  FOLLY_ALWAYS_INLINE void initialize(const core::QueryConfig& config) {
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config) {
     timeZone_ = getTimeZoneFromConfig(config);
   }
 

diff --git a/velox/functions/prestosql/HyperLogLogFunctions.h b/velox/functions/prestosql/HyperLogLogFunctions.h
@@ -62,6 +62,7 @@ struct EmptyApproxSetWithMaxErrorFunction {
   std::string serialized_;
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const double* maxStandardError) {
     VELOX_USER_CHECK_NOT_NULL(

diff --git a/velox/functions/prestosql/MapSubset.h b/velox/functions/prestosql/MapSubset.h
@@ -27,6 +27,7 @@ namespace facebook::velox::functions {
     VELOX_DEFINE_FUNCTION_TYPES(TExec);                             \
                                                                     \
     void initialize(                                                \
+        const std::vector<TypePtr>& /*inputTypes*/,                 \
         const core::QueryConfig& /*config*/,                        \
         const arg_type<Map<TType, Generic<T1>>>* /*inputMap*/,      \
         const arg_type<Array<TType>>* keys) {                       \
@@ -98,6 +99,7 @@ struct MapSubsetVarcharFunction {
   VELOX_DEFINE_FUNCTION_TYPES(TExec);
 
   void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Map<Varchar, Generic<T1>>>* /*inputMap*/,
       const arg_type<Array<Varchar>>* keys) {

diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -124,6 +124,7 @@ struct UnixTimestampParseFunction {
   // unix_timestamp(input);
   // If format is not specified, assume kDefaultFormat.
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* /*input*/) {
     format_ = buildJodaDateTimeFormatter(kDefaultFormat_);
@@ -173,6 +174,7 @@ struct UnixTimestampParseWithFormatFunction
   // unix_timestamp(input, format):
   // If format is constant, compile it just once per batch.
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* /*input*/,
       const arg_type<Varchar>* format) {
@@ -227,6 +229,7 @@ struct FromUnixtimeFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<int64_t>* /*unixtime*/,
       const arg_type<Varchar>* format) {
@@ -270,6 +273,7 @@ struct ToUtcTimestampFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Varchar>* /*input*/,
       const arg_type<Varchar>* timezone) {
@@ -299,6 +303,7 @@ struct FromUtcTimestampFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Varchar>* /*input*/,
       const arg_type<Varchar>* timezone) {
@@ -329,6 +334,7 @@ struct GetTimestampFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& config,
       const arg_type<Varchar>* /*input*/,
       const arg_type<Varchar>* format) {
@@ -608,6 +614,7 @@ struct NextDayFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig& /*config*/,
       const arg_type<Date>* /*startDate*/,
       const arg_type<Varchar>* dayOfWeek) {

diff --git a/velox/functions/sparksql/In.cpp b/velox/functions/sparksql/In.cpp
@@ -69,6 +69,7 @@ struct InFunctionOuter {
     VELOX_DEFINE_FUNCTION_TYPES(TExecCtx);
 
     FOLLY_ALWAYS_INLINE void initialize(
+        const std::vector<TypePtr>& /*inputTypes*/,
         const core::QueryConfig& /*config*/,
         const arg_type<TInput>* /*searchTerm*/,
         const arg_type<velox::Array<TInput>>* searchElements) {

diff --git a/velox/functions/sparksql/MightContain.h b/velox/functions/sparksql/MightContain.h
@@ -26,6 +26,7 @@ struct BloomFilterMightContainFunction {
   using Allocator = std::allocator<uint64_t>;
 
   void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
       const core::QueryConfig&,
       const arg_type<Varbinary>* serialized,
       const arg_type<int64_t>*) {