[OPPRO-104] Support date parsing and extract year (facebookincubator#18)

* add date parsing * support extract year * support int32 return type
rui-mo · Aug 18, 2022 · b31cfd8 · b31cfd8
1 parent dfce6fa
commit b31cfd8
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 6 deletions.
diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -137,22 +137,28 @@ struct YearFunction : public InitSessionTimezone<T>,
                       public TimestampWithTimezoneSupport<T> {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
-  FOLLY_ALWAYS_INLINE int64_t getYear(const std::tm& time) {
+  FOLLY_ALWAYS_INLINE int32_t getYear(const std::tm& time) {
     return 1900 + time.tm_year;
   }
 
-  FOLLY_ALWAYS_INLINE void call(
-      int64_t& result,
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE
+  void call(
+      TInput& result,
       const arg_type<Timestamp>& timestamp) {
     result = getYear(getDateTime(timestamp, this->timeZone_));
   }
 
-  FOLLY_ALWAYS_INLINE void call(int64_t& result, const arg_type<Date>& date) {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE
+  void call(TInput& result, const arg_type<Date>& date) {
     result = getYear(getDateTime(date));
   }
 
-  FOLLY_ALWAYS_INLINE void call(
-      int64_t& result,
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE
+  void call(
+      TInput& result,
       const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
     auto timestamp = this->toTimestamp(timestampWithTimezone);
     result = getYear(getDateTime(timestamp, nullptr));

diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp
@@ -27,6 +27,9 @@ void registerSimpleFunctions() {
       {"to_unixtime"});
   registerFunction<FromUnixtimeFunction, Timestamp, double>({"from_unixtime"});
 
+  registerFunction<YearFunction, int32_t, Timestamp>({"year"});
+  registerFunction<YearFunction, int32_t, Date>({"year"});
+  registerFunction<YearFunction, int32_t, TimestampWithTimezone>({"year"});
   registerFunction<YearFunction, int64_t, Timestamp>({"year"});
   registerFunction<YearFunction, int64_t, Date>({"year"});
   registerFunction<YearFunction, int64_t, TimestampWithTimezone>({"year"});

diff --git a/velox/substrait/SubstraitParser.cpp b/velox/substrait/SubstraitParser.cpp
@@ -94,6 +94,11 @@ std::shared_ptr<SubstraitParser::SubstraitType> SubstraitParser::parseType(
       nullability = substraitType.string().nullability();
       break;
     }
+    case ::substrait::Type::KindCase::kDate: {
+      typeName = "DATE";
+      nullability = sType.date().nullability();
+      break;
+    }
     default:
       VELOX_NYI(
           "Parsing for Substrait type not supported: {}",

diff --git a/velox/substrait/SubstraitToVeloxExpr.cpp b/velox/substrait/SubstraitToVeloxExpr.cpp
@@ -73,6 +73,36 @@ SubstraitVeloxExprConverter::toIsNotNullExpr(
       outputType, std::move(notParams), "not");
 }
 
+std::shared_ptr<const core::ITypedExpr>
+SubstraitVeloxExprConverter::toExtractExpr(
+    const std::vector<std::shared_ptr<const core::ITypedExpr>>& params,
+    const TypePtr& outputType) {
+  VELOX_CHECK_EQ(params.size(), 2);
+  auto functionArg =
+      std::dynamic_pointer_cast<const core::ConstantTypedExpr>(params[0]);
+  if (functionArg) {
+    // Get the function argument.
+    auto variant = functionArg->value();
+    if (!variant.hasValue()) {
+      VELOX_FAIL("Value expected in variant.");
+    }
+    // The first parameter specifies extracting from which field.
+    // Only year is supported currently.
+    std::string from = variant.value<std::string>();
+
+    // The second parameter is the function parameter.
+    std::vector<std::shared_ptr<const core::ITypedExpr>> exprParams;
+    exprParams.reserve(1);
+    exprParams.emplace_back(params[1]);
+    if (from == "YEAR") {
+      return std::make_shared<const core::CallTypedExpr>(
+          outputType, std::move(exprParams), "year");
+    }
+    VELOX_NYI("Extract from {} not supported.", from);
+  }
+  VELOX_FAIL("Constant is expected to be the first parameter in extract.");
+}
+
 std::shared_ptr<const core::ITypedExpr>
 SubstraitVeloxExprConverter::toVeloxExpr(
     const ::substrait::Expression::ScalarFunction& substraitFunc,
@@ -87,12 +117,16 @@ SubstraitVeloxExprConverter::toVeloxExpr(
   const auto& veloxType =
       toVeloxType(subParser_->parseType(sFunc.output_type())->type);
 
+  if (veloxFunction == "extract") {
+    return toExtractExpr(params, veloxType);
+  }
   if (veloxFunction == "alias") {
     return toAliasExpr(params);
   }
   if (veloxFunction == "is_not_null") {
     return toIsNotNullExpr(params, veloxType);
   }
+
   return std::make_shared<const core::CallTypedExpr>(
       veloxType, std::move(params), veloxFunction);
 }

diff --git a/velox/substrait/SubstraitToVeloxExpr.h b/velox/substrait/SubstraitToVeloxExpr.h
@@ -68,6 +68,11 @@ class SubstraitVeloxExprConverter {
       const std::vector<std::shared_ptr<const core::ITypedExpr>>& params,
       const TypePtr& outputType);
 
+  /// Create expression for extract.
+  std::shared_ptr<const core::ITypedExpr> toExtractExpr(
+      const std::vector<std::shared_ptr<const core::ITypedExpr>>& params,
+      const TypePtr& outputType);
+
   /// Used to convert Substrait Literal into Velox Expression.
   std::shared_ptr<const core::ConstantTypedExpr> toVeloxExpr(
       const ::substrait::Expression::Literal& substraitLit);

diff --git a/velox/substrait/TypeUtils.cpp b/velox/substrait/TypeUtils.cpp
@@ -106,6 +106,8 @@ TypePtr toVeloxType(const std::string& typeName) {
     }
     case TypeKind::UNKNOWN:
       return UNKNOWN();
+    case TypeKind::DATE:
+      return DATE();
     default:
       VELOX_NYI("Velox type conversion not supported for type {}.", typeName);
   }