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

* add date parsing * support extract year * support int32 return type * apply pr_19 use setted pool Parse and store the file format var into SubstraitVeloxPlanConverter
rui-mo · Oct 20, 2022 · 857da1c · 857da1c
1 parent ebf45e1
commit 857da1c
Show file tree

Hide file tree

Showing 11 changed files with 426 additions and 371 deletions.
diff --git a/velox/core/PlanNode.h b/velox/core/PlanNode.h
@@ -220,10 +220,12 @@ class ArrowStreamNode : public PlanNode {
       const PlanNodeId& id,
       const RowTypePtr& outputType,
       std::shared_ptr<ArrowArrayStream> arrowStream,
+      memory::MemoryPool* pool,
       bool parallelizable = false)
       : PlanNode(id),
         outputType_(outputType),
         arrowStream_(arrowStream),
+        pool_(pool),
         parallelizable_(parallelizable) {
     VELOX_CHECK(arrowStream != nullptr);
   }
@@ -247,6 +249,11 @@ class ArrowStreamNode : public PlanNode {
     return arrowStream_;
   }
 
+
+  memory::MemoryPool* memoryPool() const {
+    return pool_;
+  }
+
   std::string_view name() const override {
     return "ArrowStream";
   }
@@ -256,6 +263,7 @@ class ArrowStreamNode : public PlanNode {
 
   const RowTypePtr outputType_;
   std::shared_ptr<ArrowArrayStream> arrowStream_;
+  memory::MemoryPool* pool_;
   const bool parallelizable_;
 };
 

diff --git a/velox/exec/ArrowStream.cpp b/velox/exec/ArrowStream.cpp
@@ -28,6 +28,7 @@ ArrowStream::ArrowStream(
           arrowStream->id(),
           "Arrow Stream") {
   arrowStream_ = arrowStream->arrowStream();
+  pool_ = arrowStream->memoryPool();
 }
 
 ArrowStream::~ArrowStream() {
@@ -54,7 +55,7 @@ RowVectorPtr ArrowStream::getOutput() {
   }
   // Convert Arrow data into RowVector.
   rowVector_ = std::dynamic_pointer_cast<RowVector>(
-      facebook::velox::importFromArrowAsViewer(arrowSchema, arrowArray));
+      facebook::velox::importFromArrowAsOwner(arrowSchema, arrowArray, pool_));
   return rowVector_;
 }
 

diff --git a/velox/exec/ArrowStream.h b/velox/exec/ArrowStream.h
@@ -48,6 +48,7 @@ class ArrowStream : public SourceOperator {
   bool closed_ = false;
   RowVectorPtr rowVector_;
   std::shared_ptr<ArrowArrayStream> arrowStream_;
+  memory::MemoryPool* pool_;
 
   // For calls from destructor
   bool isFinished0();

diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -139,22 +139,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
@@ -113,6 +113,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
@@ -71,6 +71,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,
@@ -85,12 +115,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>(
       toVeloxType(typeName), 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/SubstraitToVeloxPlan.cpp b/velox/substrait/SubstraitToVeloxPlan.cpp
@@ -295,11 +295,12 @@ core::PlanNodePtr SubstraitVeloxPlanConverter::toVeloxPlan(
 
   const auto& inputType = childNode->outputType();
   const auto& sExpr = filterRel.condition();
+  std::shared_ptr<const core::PlanNode>
 
-  return std::make_shared<core::FilterNode>(
-      nextPlanNodeId(),
-      exprConverter_->toVeloxExpr(sExpr, inputType),
-      childNode);
+      return std::make_shared<core::FilterNode>(
+          nextPlanNodeId(),
+          exprConverter_->toVeloxExpr(sExpr, inputType),
+          childNode);
 }
 
 core::PlanNodePtr SubstraitVeloxPlanConverter::toVeloxPlan(