diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h index 7d315b9f2014..8e454002e944 100644 --- a/velox/functions/prestosql/DateTimeFunctions.h +++ b/velox/functions/prestosql/DateTimeFunctions.h @@ -136,22 +136,28 @@ struct YearFunction : public InitSessionTimezone, public TimestampWithTimezoneSupport { 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 + FOLLY_ALWAYS_INLINE + void call( + TInput& result, const arg_type& timestamp) { result = getYear(getDateTime(timestamp, this->timeZone_)); } - FOLLY_ALWAYS_INLINE void call(int64_t& result, const arg_type& date) { + template + FOLLY_ALWAYS_INLINE + void call(TInput& result, const arg_type& date) { result = getYear(getDateTime(date)); } - FOLLY_ALWAYS_INLINE void call( - int64_t& result, + template + FOLLY_ALWAYS_INLINE + void call( + TInput& result, const arg_type& 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 index 6c77aa953c46..a0285cffc612 100644 --- a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp @@ -27,6 +27,9 @@ void registerSimpleFunctions() { {"to_unixtime"}); registerFunction({"from_unixtime"}); + registerFunction({"year"}); + registerFunction({"year"}); + registerFunction({"year"}); registerFunction({"year"}); registerFunction({"year"}); registerFunction({"year"}); diff --git a/velox/substrait/SubstraitParser.cpp b/velox/substrait/SubstraitParser.cpp index 8185e0bbc9f9..0a6bf25fe4a3 100644 --- a/velox/substrait/SubstraitParser.cpp +++ b/velox/substrait/SubstraitParser.cpp @@ -65,8 +65,14 @@ std::shared_ptr SubstraitParser::parseType( nullability = substraitType.string().nullability(); break; } + case ::substrait::Type::KindCase::kDate: { + typeName = "DATE"; + nullability = sType.date().nullability(); + break; + } default: - VELOX_NYI("Substrait parsing for type {} not supported.", typeName); + VELOX_NYI( + "Substrait parsing for type {} not supported.", sType.kind_case()); } bool nullable; diff --git a/velox/substrait/SubstraitToVeloxExpr.cpp b/velox/substrait/SubstraitToVeloxExpr.cpp index ae902096588c..dd02f84f7c2c 100644 --- a/velox/substrait/SubstraitToVeloxExpr.cpp +++ b/velox/substrait/SubstraitToVeloxExpr.cpp @@ -73,6 +73,36 @@ SubstraitVeloxExprConverter::toIsNotNullExpr( outputType, std::move(notParams), "not"); } +std::shared_ptr +SubstraitVeloxExprConverter::toExtractExpr( + const std::vector>& params, + const TypePtr& outputType) { + VELOX_CHECK_EQ(params.size(), 2); + auto functionArg = + std::dynamic_pointer_cast(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(); + + // The second parameter is the function parameter. + std::vector> exprParams; + exprParams.reserve(1); + exprParams.emplace_back(params[1]); + if (from == "YEAR") { + return std::make_shared( + 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 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( veloxType, std::move(params), veloxFunction); } diff --git a/velox/substrait/SubstraitToVeloxExpr.h b/velox/substrait/SubstraitToVeloxExpr.h index 62fb7198fbcb..cb4bd1b61fa3 100644 --- a/velox/substrait/SubstraitToVeloxExpr.h +++ b/velox/substrait/SubstraitToVeloxExpr.h @@ -68,6 +68,11 @@ class SubstraitVeloxExprConverter { const std::vector>& params, const TypePtr& outputType); + /// Create expression for extract. + std::shared_ptr toExtractExpr( + const std::vector>& params, + const TypePtr& outputType); + /// Used to convert Substrait Literal into Velox Expression. std::shared_ptr toVeloxExpr( const ::substrait::Expression::Literal& substraitLit); diff --git a/velox/substrait/TypeUtils.cpp b/velox/substrait/TypeUtils.cpp index d2bdc2753e41..67483c0939f2 100644 --- a/velox/substrait/TypeUtils.cpp +++ b/velox/substrait/TypeUtils.cpp @@ -56,6 +56,8 @@ TypePtr toVeloxType(const std::string& typeName) { return DOUBLE(); case TypeKind::VARCHAR: return VARCHAR(); + case TypeKind::DATE: + return DATE(); default: VELOX_NYI("Velox type conversion not supported for type {}.", typeName); }