diff --git a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp index 58288df27bd..9b3a80b0750 100644 --- a/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp +++ b/velox/experimental/cudf/connectors/hive/CudfHiveDataSource.cpp @@ -113,8 +113,16 @@ CudfHiveDataSource::CudfHiveDataSource( } auto const remainingFilterType = getTableRowType(); + // The connector exposes the session timezone directly rather than through a + // QueryConfig. sessionStartTimeMs is unused for filter pushdown (no now() / + // current_timestamp here), so it is left at 0. + const velox::cudf_velox::CudfDateTimeContext context{ + connectorQueryCtx_->sessionTimezone(), + connectorQueryCtx_->adjustTimestampToTimezone(), + 0, + }; cudfExpressionEvaluator_ = velox::cudf_velox::createCudfExpression( - remainingFilterExprSet_->exprs()[0], remainingFilterType); + remainingFilterExprSet_->exprs()[0], remainingFilterType, context); // TODO(kn): Get column names and subfields from remaining filter and add to // readColumnNames_ } diff --git a/velox/experimental/cudf/exec/CudfFilterProject.cpp b/velox/experimental/cudf/exec/CudfFilterProject.cpp index 7576ce68c3a..97555b097b8 100644 --- a/velox/experimental/cudf/exec/CudfFilterProject.cpp +++ b/velox/experimental/cudf/exec/CudfFilterProject.cpp @@ -18,9 +18,7 @@ #include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/exec/CudfFilterProject.h" #include "velox/experimental/cudf/exec/GpuResources.h" -#include "velox/experimental/cudf/exec/Validation.h" #include "velox/experimental/cudf/exec/VeloxCudfInterop.h" -#include "velox/experimental/cudf/expression/DateTruncFunction.h" #include "velox/experimental/cudf/vector/CudfVector.h" #include "velox/common/memory/Memory.h" @@ -53,28 +51,6 @@ void debugPrintTree( } } -bool isTimezoneSensitiveDateTrunc( - const std::shared_ptr& expr) { - const auto dateTruncName = - CudfConfig::getInstance().functionNamePrefix + "date_trunc"; - return expr->name() == dateTruncName && - DateTruncFunction::isTimezoneSensitive(expr); -} - -bool containsTimezoneSensitiveDateTrunc( - const std::shared_ptr& expr) { - if (isTimezoneSensitiveDateTrunc(expr)) { - return true; - } - - for (const auto& input : expr->inputs()) { - if (containsTimezoneSensitiveDateTrunc(input)) { - return true; - } - } - return false; -} - bool checkAddIdentityProjection( const core::TypedExprPtr& projection, const RowTypePtr& inputType, @@ -145,19 +121,7 @@ bool canBeEvaluatedByCudf( std::unique_ptr exprSet = exec::makeExprSetFromFlag( std::move(exprsCopy), &precompileCtx, lazyDereference); - const core::QueryConfig defaultQueryConfig = core::QueryConfig({}); - const core::QueryConfig& queryConfig = - queryCtx ? queryCtx->queryConfig() : defaultQueryConfig; - const bool adjustTimestampToTimezone = - queryConfig.adjustTimestampToTimezone(); - for (const auto& e : exprSet->exprs()) { - if (adjustTimestampToTimezone && containsTimezoneSensitiveDateTrunc(e)) { - LOG_FALLBACK( - "date_trunc(timestamp) requires CPU evaluation when " - "adjust_timestamp_to_session_timezone is enabled"); - return false; - } if (!canBeEvaluatedByCudf(e)) { return false; } @@ -233,6 +197,12 @@ void CudfFilterProject::initialize() { const auto inputType = project_ ? project_->sources()[0]->outputType() : filter_->sources()[0]->outputType(); + // Capture the session timezone so timezone-aware GPU functions (date/time + // extraction, the TIMESTAMP WITH TIME ZONE family, temporal casts) match the + // CPU path. + const auto exprContext = + contextFromConfig(operatorCtx_->driverCtx()->queryConfig()); + // convert to AST if (CudfConfig::getInstance().debugEnabled) { int i = 0; @@ -243,21 +213,22 @@ void CudfFilterProject::initialize() { } if (hasFilter_) { // First expr is Filter, rest are Project - filterEvaluator_ = createCudfExpression(expr->exprs()[0], inputType); + filterEvaluator_ = + createCudfExpression(expr->exprs()[0], inputType, exprContext); std::transform( expr->exprs().begin() + 1, expr->exprs().end(), std::back_inserter(projectEvaluators_), - [inputType](const auto& expr) { - return createCudfExpression(expr, inputType); + [&](const auto& expr) { + return createCudfExpression(expr, inputType, exprContext); }); } else { std::transform( expr->exprs().begin(), expr->exprs().end(), std::back_inserter(projectEvaluators_), - [inputType](const auto& expr) { - return createCudfExpression(expr, inputType); + [&](const auto& expr) { + return createCudfExpression(expr, inputType, exprContext); }); } diff --git a/velox/experimental/cudf/exec/CudfHashJoin.cpp b/velox/experimental/cudf/exec/CudfHashJoin.cpp index 95d92d7b96d..afe8014e39b 100644 --- a/velox/experimental/cudf/exec/CudfHashJoin.cpp +++ b/velox/experimental/cudf/exec/CudfHashJoin.cpp @@ -26,6 +26,7 @@ #include "velox/common/testutil/TestValue.h" #include "velox/core/PlanNode.h" +#include "velox/exec/Driver.h" #include "velox/exec/Task.h" // NOLINT(misc-unused-headers) #include "velox/type/TypeUtil.h" @@ -490,9 +491,16 @@ void CudfHashJoinProbe::initialize() { // Create a reusable evaluator for the filter column. This is expensive to // build, and the expression + input schema are stable for the lifetime of // the operator instance. + // Resolve the session timezone once so timezone-sensitive CudfFunctions in + // the join filter receive it at construction. + const auto context = + contextFromConfig(operatorCtx_->driverCtx()->queryConfig()); + std::vector filterRowTypes{probeType_, buildType_}; filterEvaluator_ = createCudfExpression( - exprs.exprs()[0], facebook::velox::type::concatRowTypes(filterRowTypes)); + exprs.exprs()[0], + facebook::velox::type::concatRowTypes(filterRowTypes), + context); // Check if the filter expression spans both join sides (e.g., switch // expressions referencing columns from both probe and build). If so, we @@ -521,7 +529,8 @@ void CudfHashJoinProbe::initialize() { buildType_, probeType_, rightPrecomputeInstructions_, - leftPrecomputeInstructions_); + leftPrecomputeInstructions_, + context); } else { createAstTree( exprs.exprs()[0], @@ -530,7 +539,8 @@ void CudfHashJoinProbe::initialize() { probeType_, buildType_, leftPrecomputeInstructions_, - rightPrecomputeInstructions_); + rightPrecomputeInstructions_, + context); } } } diff --git a/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp b/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp index 82a88bd8950..d6bca6073c8 100644 --- a/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp +++ b/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp @@ -25,6 +25,7 @@ #include "velox/experimental/cudf/expression/AstExpressionUtils.h" #include "velox/experimental/cudf/expression/PrecomputeInstruction.h" +#include "velox/exec/Driver.h" #include "velox/exec/Task.h" #include @@ -294,6 +295,11 @@ void CudfNestedLoopJoinProbe::initialize() { exec::ExprSet exprs({joinNode_->joinCondition()}, operatorCtx_->execCtx()); VELOX_CHECK_EQ(exprs.exprs().size(), 1); + // Resolve the session timezone once so timezone-sensitive CudfFunctions built + // on the precompute path receive it at construction. + const auto context = + contextFromConfig(operatorCtx_->driverCtx()->queryConfig()); + // Convert Velox expression to cuDF AST expression tree. // The AST will be passed to cudf::conditional_inner_join() for GPU // evaluation. @@ -304,7 +310,8 @@ void CudfNestedLoopJoinProbe::initialize() { probeType_, buildType_, leftPrecomputeInstructions_, - rightPrecomputeInstructions_); + rightPrecomputeInstructions_, + context); // Set hasFilter_ only after the AST has been fully built so that a throw // from createAstTree() does not leave the operator marked as having a filter diff --git a/velox/experimental/cudf/expression/ArrayAccessFunctions.cpp b/velox/experimental/cudf/expression/ArrayAccessFunctions.cpp index 308dd99672d..728b1ea576f 100644 --- a/velox/experimental/cudf/expression/ArrayAccessFunctions.cpp +++ b/velox/experimental/cudf/expression/ArrayAccessFunctions.cpp @@ -519,6 +519,7 @@ class ArrayAccessFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // Case 1: constant array, variable index. diff --git a/velox/experimental/cudf/expression/AstExpression.cpp b/velox/experimental/cudf/expression/AstExpression.cpp index 28d4aa87409..63f2f89c01a 100644 --- a/velox/experimental/cudf/expression/AstExpression.cpp +++ b/velox/experimental/cudf/expression/AstExpression.cpp @@ -35,10 +35,11 @@ cudf::ast::expression const& createAstTree( cudf::ast::tree& tree, std::vector>& scalars, const RowTypePtr& inputRowSchema, - std::vector& precomputeInstructions) { - AstContext context{ - tree, scalars, {inputRowSchema}, {precomputeInstructions}, expr}; - return context.pushExprToTree(expr); + std::vector& precomputeInstructions, + const CudfDateTimeContext& context) { + AstContext astContext{ + tree, scalars, {inputRowSchema}, {precomputeInstructions}, expr, context}; + return astContext.pushExprToTree(expr); } cudf::ast::expression const& createAstTree( @@ -48,22 +49,30 @@ cudf::ast::expression const& createAstTree( const RowTypePtr& leftRowSchema, const RowTypePtr& rightRowSchema, std::vector& leftPrecomputeInstructions, - std::vector& rightPrecomputeInstructions) { - AstContext context{ + std::vector& rightPrecomputeInstructions, + const CudfDateTimeContext& context) { + AstContext astContext{ tree, scalars, {leftRowSchema, rightRowSchema}, {leftPrecomputeInstructions, rightPrecomputeInstructions}, - expr}; - return context.pushExprToTree(expr); + expr, + context}; + return astContext.pushExprToTree(expr); } ASTExpression::ASTExpression( std::shared_ptr expr, - const RowTypePtr& inputRowSchema) + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context) : expr_(expr), inputRowSchema_(inputRowSchema) { createAstTree( - expr, cudfTree_, scalars_, inputRowSchema, precomputeInstructions_); + expr, + cudfTree_, + scalars_, + inputRowSchema, + precomputeInstructions_, + context); } void ASTExpression::close() { @@ -134,8 +143,10 @@ void registerAstEvaluator(int priority) { [](std::shared_ptr expr) { return ASTExpression::canEvaluate(expr); }, - [](std::shared_ptr expr, const RowTypePtr& row) { - return std::make_shared(std::move(expr), row); + [](std::shared_ptr expr, + const RowTypePtr& row, + const CudfDateTimeContext& context) { + return std::make_shared(std::move(expr), row, context); }, /*overwrite=*/false); } diff --git a/velox/experimental/cudf/expression/AstExpression.h b/velox/experimental/cudf/expression/AstExpression.h index 47f88404c84..ecbdd2e0e44 100644 --- a/velox/experimental/cudf/expression/AstExpression.h +++ b/velox/experimental/cudf/expression/AstExpression.h @@ -29,7 +29,8 @@ cudf::ast::expression const& createAstTree( cudf::ast::tree& tree, std::vector>& scalars, const RowTypePtr& inputRowSchema, - std::vector& precomputeInstructions); + std::vector& precomputeInstructions, + const CudfDateTimeContext& context); cudf::ast::expression const& createAstTree( const std::shared_ptr& expr, @@ -38,7 +39,8 @@ cudf::ast::expression const& createAstTree( const RowTypePtr& leftRowSchema, const RowTypePtr& rightRowSchema, std::vector& leftPrecomputeInstructions, - std::vector& rightPrecomputeInstructions); + std::vector& rightPrecomputeInstructions, + const CudfDateTimeContext& context); // Evaluates the expression tree class ASTExpression : public CudfExpression { @@ -48,7 +50,8 @@ class ASTExpression : public CudfExpression { // precompute instructions and stores them ASTExpression( std::shared_ptr expr, - const RowTypePtr& inputRowSchema); + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context); // Evaluates the expression tree for the given input columns ColumnOrView eval( diff --git a/velox/experimental/cudf/expression/AstExpressionUtils.h b/velox/experimental/cudf/expression/AstExpressionUtils.h index 794820ab616..f62483624dc 100644 --- a/velox/experimental/cudf/expression/AstExpressionUtils.h +++ b/velox/experimental/cudf/expression/AstExpressionUtils.h @@ -415,6 +415,10 @@ struct AstContext { precomputeInstructions; const std::shared_ptr rootExpr; // Track the root expression + // Query-scoped context threaded into timezone-sensitive functions built on + // the precompute path (e.g. date_format or a VARCHAR->TIMESTAMP cast inside a + // join condition). + CudfDateTimeContext context; bool allowPureAstOnly; cudf::ast::expression const& pushExprToTree( @@ -580,7 +584,7 @@ cudf::ast::expression const& AstContext::pushExprToTree( if (sideIdx < 0) { sideIdx = 0; // Default to left side if no fields found } - auto node = createCudfExpression(expr, inputRowSchema[sideIdx]); + auto node = createCudfExpression(expr, inputRowSchema[sideIdx], context); return addPrecomputeInstructionOnSide(sideIdx, 0, name, "", node); } VELOX_FAIL("Unsupported expression: {}", name); diff --git a/velox/experimental/cudf/expression/CMakeLists.txt b/velox/experimental/cudf/expression/CMakeLists.txt index c85a51f355f..f5c6f787132 100644 --- a/velox/experimental/cudf/expression/CMakeLists.txt +++ b/velox/experimental/cudf/expression/CMakeLists.txt @@ -26,17 +26,21 @@ add_library( PrestoFunctions.cpp prestosql/DateAddFunction.cpp prestosql/DatePlusIntervalFunction.cpp + prestosql/TimezoneFunctions.cpp SparkFunctions.cpp sparksql/DateAddFunction.cpp sparksql/HashFunction.cpp sparksql/SubStringFunction.cpp SubfieldFiltersToAst.cpp + TimestampWithTimeZoneColumn.cpp + TimezoneConversion.cpp ) target_link_libraries( velox_cudf_expression PUBLIC cudf::cudf PRIVATE arrow velox_common_base velox_cudf_vector velox_exception + velox_presto_types ) target_compile_options(velox_cudf_expression PRIVATE -Wno-missing-field-initializers) diff --git a/velox/experimental/cudf/expression/DateTruncFunction.cpp b/velox/experimental/cudf/expression/DateTruncFunction.cpp index 56209657951..ad0c4987742 100644 --- a/velox/experimental/cudf/expression/DateTruncFunction.cpp +++ b/velox/experimental/cudf/expression/DateTruncFunction.cpp @@ -16,9 +16,12 @@ #include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/expression/AstUtils.h" #include "velox/experimental/cudf/expression/DateTruncFunction.h" +#include "velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h" +#include "velox/experimental/cudf/expression/TimezoneConversion.h" #include "velox/expression/ConstantExpr.h" #include "velox/functions/lib/TimeUtils.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include #include @@ -29,6 +32,29 @@ namespace facebook::velox::cudf_velox { using functions::DateTimeUnit; +namespace { + +// Maps a timestamp data type to the duration type of the same resolution, used +// to subtract the hour-truncation remainder from the UTC instant. +cudf::data_type durationTypeForTimestamp(cudf::data_type timestampType) { + switch (timestampType.id()) { + case cudf::type_id::TIMESTAMP_SECONDS: + return cudf::data_type(cudf::type_id::DURATION_SECONDS); + case cudf::type_id::TIMESTAMP_MILLISECONDS: + return cudf::data_type(cudf::type_id::DURATION_MILLISECONDS); + case cudf::type_id::TIMESTAMP_MICROSECONDS: + return cudf::data_type(cudf::type_id::DURATION_MICROSECONDS); + case cudf::type_id::TIMESTAMP_NANOSECONDS: + return cudf::data_type(cudf::type_id::DURATION_NANOSECONDS); + default: + VELOX_FAIL( + "date_trunc hour requires a timestamp column, got cudf type id: {}", + static_cast(timestampType.id())); + } +} + +} // namespace + bool DateTruncFunction::canEvaluate( const std::shared_ptr& expr) { if (expr->inputs().size() != 2) { @@ -48,12 +74,13 @@ bool DateTruncFunction::canEvaluate( const auto& inputType = expr->inputs()[1]->type(); const bool isTimestamp = inputType->isTimestamp(); const bool isDate = inputType->isDate(); - if (!isTimestamp && !isDate) { + const bool isTswtz = isTimestampWithTimeZoneType(inputType); + if (!isTimestamp && !isDate && !isTswtz) { return false; } if (*unit == DateTimeUnit::kSecond || *unit == DateTimeUnit::kMinute || *unit == DateTimeUnit::kHour) { - return isTimestamp; + return isTimestamp || isTswtz; } if (*unit == DateTimeUnit::kDay || *unit == DateTimeUnit::kWeek || *unit == DateTimeUnit::kMonth || *unit == DateTimeUnit::kQuarter || @@ -63,19 +90,6 @@ bool DateTruncFunction::canEvaluate( return false; } -bool DateTruncFunction::isTimezoneSensitive( - const std::shared_ptr& expr) { - if (!canEvaluate(expr) || !expr->inputs()[1]->type()->isTimestamp()) { - return false; - } - - const auto unitString = constantVarcharValue(expr->inputs()[0]); - const auto unit = functions::fromDateTimeUnitString(*unitString, false); - return *unit == DateTimeUnit::kHour || *unit == DateTimeUnit::kDay || - *unit == DateTimeUnit::kWeek || *unit == DateTimeUnit::kMonth || - *unit == DateTimeUnit::kQuarter || *unit == DateTimeUnit::kYear; -} - DateTruncFunction::DateTruncFunction( const std::shared_ptr& expr) { VELOX_CHECK_EQ( @@ -86,18 +100,21 @@ DateTruncFunction::DateTruncFunction( auto inputType = expr->inputs()[1]->type(); const bool isTimestamp = inputType->isTimestamp(); const bool isDate = inputType->isDate(); + isTimestampWithTimeZone_ = isTimestampWithTimeZoneType(inputType); VELOX_CHECK( - isTimestamp || isDate, - "date_trunc only supports date or timestamp inputs"); + isTimestamp || isDate || isTimestampWithTimeZone_, + "date_trunc only supports date, timestamp, or timestamp with time zone inputs"); auto parsed = functions::fromDateTimeUnitString(*unitString, true); VELOX_CHECK(parsed.has_value(), "Invalid date_trunc unit: {}", *unitString); unit_ = *parsed; - // Validate time-only units require timestamp input. + // Validate time-only units require an instant (timestamp or TSWTZ) input. if (unit_ == DateTimeUnit::kSecond || unit_ == DateTimeUnit::kMinute || unit_ == DateTimeUnit::kHour) { VELOX_CHECK( - isTimestamp, "date_trunc {} requires timestamp input", *unitString); + isTimestamp || isTimestampWithTimeZone_, + "date_trunc {} requires timestamp input", + *unitString); } auto stream = cudf::get_default_stream(cudf::allow_default_stream); @@ -111,12 +128,10 @@ DateTruncFunction::DateTruncFunction( stream.synchronize(); } -ColumnOrView DateTruncFunction::eval( - std::vector& inputColumns, +ColumnOrView DateTruncFunction::truncateOnColumn( + cudf::column_view inputCol, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const { - VELOX_CHECK_EQ(inputColumns.size(), 1, "date_trunc expects one column input"); - auto inputCol = asView(inputColumns[0]); auto outputType = inputCol.type(); auto dayType = cudf::data_type(cudf::type_id::TIMESTAMP_DAYS); auto intType = cudf::data_type(cudf::type_id::INT32); @@ -252,4 +267,105 @@ ColumnOrView DateTruncFunction::eval( VELOX_UNREACHABLE(); } +ColumnOrView DateTruncFunction::eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + VELOX_CHECK_EQ(inputColumns.size(), 1, "date_trunc expects one column input"); + + if (isTimestampWithTimeZone_) { + // TIMESTAMP WITH TIME ZONE carries its own zone per row; truncate on each + // row's embedded wall clock (per-row multi-zone), independent of the + // session zone. Matches CPU DateTruncFunction::call(TSWTZ). + auto packed = asView(inputColumns[0]); + auto zoneKey = tswtzZoneKey(packed, stream, mr); + auto distinct = tswtzDistinctZoneKeys(zoneKey->view(), stream, mr); + auto local = tswtzLocalWallClock(packed, stream, mr); + + std::unique_ptr truncatedUtcMillis; + if (unit_ == DateTimeUnit::kSecond || unit_ == DateTimeUnit::kMinute || + unit_ == DateTimeUnit::kHour) { + // unit < day: take the local-to-truncated delta and subtract it from the + // UTC instant. Whole-minute offsets make this exact for second/minute, + // and it is the DST-safe form for the hour branch. + ColumnOrView flooredLocal = truncateOnColumn(local->view(), stream, mr); + auto delta = cudf::binary_operation( + local->view(), + asView(flooredLocal), + cudf::binary_operator::SUB, + cudf::data_type(cudf::type_id::DURATION_MILLISECONDS), + stream, + mr); + auto utcInstant = tswtzUtcInstant(packed, stream, mr); + truncatedUtcMillis = cudf::binary_operation( + utcInstant->view(), + delta->view(), + cudf::binary_operator::SUB, + cudf::data_type(cudf::type_id::TIMESTAMP_MILLISECONDS), + stream, + mr); + } else { + // day and above: truncate the local wall clock, then convert back to UTC + // per row's zone (a spring-forward gap throws, matching toGMT). + ColumnOrView truncatedLocal = truncateOnColumn(local->view(), stream, mr); + truncatedUtcMillis = tswtzLocalToUtc( + asView(truncatedLocal), + zoneKey->view(), + distinct, + /*correctForward=*/false, + stream, + mr); + } + return tswtzPack(truncatedUtcMillis->view(), zoneKey->view(), stream, mr); + } + + auto inputCol = asView(inputColumns[0]); + const auto outputType = inputCol.type(); + + // DATE (TIMESTAMP_DAYS) is zone-free, and under a UTC session no conversion + // is needed; both use the raw truncation directly. + const bool applyTimezone = outputType.id() != cudf::type_id::TIMESTAMP_DAYS && + context_.appliesSessionTimezone(); + if (!applyTimezone || unit_ == DateTimeUnit::kSecond || + unit_ == DateTimeUnit::kMinute) { + // second/minute truncate the UTC epoch directly (every zone offset is a + // whole number of minutes), matching CPU truncateTimestamp. + return truncateOnColumn(inputCol, stream, mr); + } + + const std::string& zone = context_.sessionTimezone; + + if (unit_ == DateTimeUnit::kHour) { + // Compute the local-to-truncated-hour delta and subtract it from the UTC + // instant. This reproduces CPU truncateTimestamp's DST-safe hour branch + // (which avoids the ambiguous local->UTC roundtrip) and handles + // fractional-offset zones such as Asia/Kolkata (+05:30). + auto local = toLocalTimestamp(inputCol, zone, stream, mr); + auto flooredLocal = cudf::datetime::floor_datetimes( + local->view(), cudf::datetime::rounding_frequency::HOUR, stream, mr); + auto delta = cudf::binary_operation( + local->view(), + flooredLocal->view(), + cudf::binary_operator::SUB, + durationTypeForTimestamp(outputType), + stream, + mr); + return cudf::binary_operation( + inputCol, + delta->view(), + cudf::binary_operator::SUB, + outputType, + stream, + mr); + } + + // day and above: truncate on the local wall clock, then convert back to UTC. + // Matches CPU truncateTimestamp (truncate local, then toGMT); a local time in + // a spring-forward gap raises in toUtcTimestamp, which is the correct parity. + auto local = toLocalTimestamp(inputCol, zone, stream, mr); + ColumnOrView truncatedLocal = truncateOnColumn(local->view(), stream, mr); + return toUtcTimestamp(asView(truncatedLocal), zone, stream, mr); +} + } // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/DateTruncFunction.h b/velox/experimental/cudf/expression/DateTruncFunction.h index ebbcd8869dd..5dbafe771ca 100644 --- a/velox/experimental/cudf/expression/DateTruncFunction.h +++ b/velox/experimental/cudf/expression/DateTruncFunction.h @@ -30,18 +30,28 @@ class DateTruncFunction : public CudfFunction { public: static bool canEvaluate(const std::shared_ptr& expr); - static bool isTimezoneSensitive( - const std::shared_ptr& expr); - explicit DateTruncFunction(const std::shared_ptr& expr); ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override; private: + // Truncates inputCol to unit_ on the values as given, with no timezone + // conversion. The timezone-aware eval wraps this with toLocalTimestamp / + // toUtcTimestamp for day-and-above units under a session timezone. + ColumnOrView truncateOnColumn( + cudf::column_view inputCol, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + functions::DateTimeUnit unit_{}; + // True when the input is TIMESTAMP WITH TIME ZONE; eval then truncates on + // each row's embedded zone (per-row multi-zone), independent of the session + // zone. + bool isTimestampWithTimeZone_{false}; std::unique_ptr oneScalar_; std::unique_ptr threeScalar_; std::unique_ptr negOneScalar_; diff --git a/velox/experimental/cudf/expression/ExpressionEvaluator.cpp b/velox/experimental/cudf/expression/ExpressionEvaluator.cpp index b229bae2254..535b5c819bf 100644 --- a/velox/experimental/cudf/expression/ExpressionEvaluator.cpp +++ b/velox/experimental/cudf/expression/ExpressionEvaluator.cpp @@ -19,9 +19,12 @@ #include "velox/experimental/cudf/expression/DecimalExpressionKernels.h" #include "velox/experimental/cudf/expression/ExpressionEvaluator.h" #include "velox/experimental/cudf/expression/NullMask.h" +#include "velox/experimental/cudf/expression/TimezoneConversion.h" +#include "velox/experimental/cudf/expression/prestosql/TimezoneFunctions.h" #include "velox/common/base/Exceptions.h" #include "velox/common/memory/Memory.h" +#include "velox/core/QueryConfig.h" #include "velox/core/QueryCtx.h" #include "velox/expression/ConstantExpr.h" #include "velox/expression/EvalCtx.h" @@ -227,8 +230,10 @@ static void ensureBuiltinExpressionEvaluatorsRegistered() { [](std::shared_ptr expr) { return FunctionExpression::canEvaluate(std::move(expr)); }, - [](std::shared_ptr expr, const RowTypePtr& row) { - return FunctionExpression::create(std::move(expr), row); + [](std::shared_ptr expr, + const RowTypePtr& row, + const CudfDateTimeContext& context) { + return FunctionExpression::create(std::move(expr), row, context); }, /*overwrite=*/false); @@ -280,6 +285,14 @@ getCudfFunctionRegistry() { return registry; } +CudfDateTimeContext contextFromConfig(const core::QueryConfig& config) { + return CudfDateTimeContext{ + config.sessionTimezone(), + config.adjustTimestampToTimezone(), + config.sessionStartTimeMs(), + }; +} + namespace { static bool matchCallAgainstSignatures( @@ -341,6 +354,7 @@ class SplitFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -371,6 +385,7 @@ class CastFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -392,6 +407,7 @@ class CardinalityFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -407,6 +423,7 @@ class IsNullFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { VELOX_CHECK_EQ(inputColumns.size(), 1, "is_null expects 1 input"); @@ -422,6 +439,7 @@ class IsNotNullFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { VELOX_CHECK_EQ(inputColumns.size(), 1, "isnotnull expects 1 input"); @@ -455,6 +473,7 @@ class RoundFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -561,6 +580,7 @@ class BinaryFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto isComparisonOp = [](cudf::binary_operator op) { @@ -875,6 +895,7 @@ class LogicalFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // If there are no input columns, the result is a scalar. @@ -986,6 +1007,7 @@ class UnaryFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { return cudf::unary_operation(asView(inputColumns[0]), op_, stream, mr); @@ -1024,6 +1046,7 @@ class BetweenFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // return (value >= min) && (value <= max) @@ -1131,6 +1154,7 @@ class GreatestLeastFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // All inputs were constant -- return the pre-folded scalar as a column. @@ -1192,6 +1216,7 @@ class SwitchFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { if (left_ == nullptr && right_ == nullptr) { @@ -1257,6 +1282,7 @@ class CoalesceFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // Coalesce is practically a cudf::replace_nulls over multiple columns. @@ -1295,6 +1321,35 @@ class CoalesceFunction : public CudfFunction { std::unique_ptr literalScalar_; }; +// Returns true for timestamp types whose calendar fields depend on the session +// timezone. DATE / TIMESTAMP_DAYS are timezone-naive on the CPU path and are +// excluded. +bool isSubDayTimestamp(cudf::data_type type) { + switch (type.id()) { + case cudf::type_id::TIMESTAMP_SECONDS: + case cudf::type_id::TIMESTAMP_MILLISECONDS: + case cudf::type_id::TIMESTAMP_MICROSECONDS: + case cudf::type_id::TIMESTAMP_NANOSECONDS: + return true; + default: + return false; + } +} + +// Converts a timestamp column to the session-local wall clock when the context +// requests it, so a following extraction reads local fields like the CPU path. +// Returns nullptr when no conversion applies; callers then use the input view. +std::unique_ptr maybeConvertToSessionLocal( + const cudf::column_view& input, + const CudfDateTimeContext& context, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (!context.appliesSessionTimezone() || !isSubDayTimestamp(input.type())) { + return nullptr; + } + return toLocalTimestamp(input, context.sessionTimezone, stream, mr); +} + class ExtractComponentFunction : public CudfFunction { public: ExtractComponentFunction( @@ -1307,11 +1362,21 @@ class ExtractComponentFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); + // second and millisecond are sub-minute fields: every timezone offset is a + // whole number of minutes, so they are unaffected by the session timezone. + // The CPU path extracts them without applying the timezone, so skip the + // conversion here to match. + std::unique_ptr local; + if (component_ != cudf::datetime::datetime_component::SECOND && + component_ != cudf::datetime::datetime_component::MILLISECOND) { + local = maybeConvertToSessionLocal(inputCol, context_, stream, mr); + } return cudf::datetime::extract_datetime_component( - inputCol, component_, stream, mr); + local ? local->view() : inputCol, component_, stream, mr); } private: @@ -1339,10 +1404,13 @@ class QuarterFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); - return cudf::datetime::extract_quarter(inputCol, stream, mr); + auto local = maybeConvertToSessionLocal(inputCol, context_, stream, mr); + return cudf::datetime::extract_quarter( + local ? local->view() : inputCol, stream, mr); } }; @@ -1355,10 +1423,13 @@ class DayOfYearFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); - return cudf::datetime::day_of_year(inputCol, stream, mr); + auto local = maybeConvertToSessionLocal(inputCol, context_, stream, mr); + return cudf::datetime::day_of_year( + local ? local->view() : inputCol, stream, mr); } }; @@ -1371,11 +1442,17 @@ class WeekFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); + auto local = maybeConvertToSessionLocal(inputCol, context_, stream, mr); auto weekStrings = cudf::strings::from_timestamps( - inputCol, "%V", cudf::strings_column_view{}, stream, mr); + local ? local->view() : inputCol, + "%V", + cudf::strings_column_view{}, + stream, + mr); return cudf::strings::to_integers( cudf::strings_column_view(weekStrings->view()), cudf::data_type(cudf::type_id::INT32), @@ -1395,11 +1472,17 @@ class YearOfWeekFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); + auto local = maybeConvertToSessionLocal(inputCol, context_, stream, mr); auto yearStrings = cudf::strings::from_timestamps( - inputCol, "%G", cudf::strings_column_view{}, stream, mr); + local ? local->view() : inputCol, + "%G", + cudf::strings_column_view{}, + stream, + mr); return cudf::strings::to_integers( cudf::strings_column_view(yearStrings->view()), cudf::data_type(cudf::type_id::INT32), @@ -1417,6 +1500,7 @@ class LengthFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -1433,6 +1517,7 @@ class LowerFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -1449,6 +1534,7 @@ class UpperFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -1512,6 +1598,7 @@ class LikeFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { size_t nextInput = 0; @@ -1767,6 +1854,7 @@ class StringPatternPredicateFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { size_t nextInput = 0; @@ -1912,6 +2000,7 @@ class ConcatFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { // Validate sizes. @@ -1991,6 +2080,7 @@ class RowConstructorFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { VELOX_CHECK( @@ -2072,7 +2162,8 @@ void registerCudfFunctions( std::shared_ptr createCudfFunction( const std::string& name, - const std::shared_ptr& expr) { + const std::shared_ptr& expr, + const CudfDateTimeContext& context) { auto& registry = getCudfFunctionRegistry(); auto it = registry.find(name); if (it == registry.end()) { @@ -2088,7 +2179,11 @@ std::shared_ptr createCudfFunction( if (spec.canEvaluate && !spec.canEvaluate(expr)) { continue; } - return spec.factory(name, expr); + auto function = spec.factory(name, expr); + if (function) { + function->setContext(context); + } + return function; } return nullptr; } @@ -2683,6 +2778,11 @@ bool registerBuiltinFunctions(const std::string& prefix) { .variableArity("decimal(p,s)") .build()}); + // TIMESTAMP WITH TIME ZONE function family (from_unixtime, to_unixtime, + // at_timezone, timezone_hour/minute, to_iso8601, format_datetime, + // parse_datetime, from_iso8601_timestamp, now/current_timestamp). + registerTimezoneFunctions(prefix); + // Note: Spark and Presto functions are now registered separately via // registerSparkFunctions() and registerPrestoFunctions() return true; @@ -2690,7 +2790,8 @@ bool registerBuiltinFunctions(const std::string& prefix) { std::shared_ptr FunctionExpression::create( const std::shared_ptr& expr, - const RowTypePtr& inputRowSchema) { + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context) { using velox::exec::FieldReference; auto node = std::make_shared(); @@ -2698,7 +2799,7 @@ std::shared_ptr FunctionExpression::create( node->inputRowSchema_ = inputRowSchema; auto name = expr->name(); - node->function_ = createCudfFunction(name, expr); + node->function_ = createCudfFunction(name, expr, context); if (auto fieldExpr = std::dynamic_pointer_cast(expr)) { if (!fieldExpr->inputs().empty()) { @@ -2719,7 +2820,7 @@ std::shared_ptr FunctionExpression::create( for (const auto& input : expr->inputs()) { if (!std::dynamic_pointer_cast(input)) { node->subexpressions_.push_back( - createCudfExpression(input, inputRowSchema)); + createCudfExpression(input, inputRowSchema, context)); } } } @@ -2804,7 +2905,16 @@ ColumnOrView FunctionExpression::eval( subexprResults.push_back(subexpr->eval(inputColumnViews, stream, mr)); } - auto result = function_->eval(subexprResults, stream, mr); + // The batch row count, threaded to eval so a zero-argument function (e.g. + // now()) can size its constant output. Every argument and input column + // carries it; the GPU path always has at least one input column + // (zero-column projections fall back to CPU and empty batches short-circuit + // upstream). + const auto numRows = !subexprResults.empty() + ? asView(subexprResults.front()).size() + : (inputColumnViews.empty() ? 0 : inputColumnViews.front().size()); + + auto result = function_->eval(subexprResults, numRows, stream, mr); if (finalize) { const auto requestedType = cudf_velox::veloxToCudfDataType(expr_->type()); auto resultView = asView(result); @@ -2892,7 +3002,8 @@ bool canBeEvaluatedByCudf(std::shared_ptr expr, bool deep) { std::shared_ptr createCudfExpression( std::shared_ptr expr, - const RowTypePtr& inputRowSchema) { + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context) { ensureBuiltinExpressionEvaluatorsRegistered(); const auto& registry = getCudfExpressionEvaluatorRegistry(); @@ -2906,10 +3017,10 @@ std::shared_ptr createCudfExpression( } if (best != nullptr) { - return best->create(expr, inputRowSchema); + return best->create(expr, inputRowSchema, context); } - return FunctionExpression::create(expr, inputRowSchema); + return FunctionExpression::create(expr, inputRowSchema, context); } void unregisterFunctions() { diff --git a/velox/experimental/cudf/expression/ExpressionEvaluator.h b/velox/experimental/cudf/expression/ExpressionEvaluator.h index ec5af17c952..9d30e194073 100644 --- a/velox/experimental/cudf/expression/ExpressionEvaluator.h +++ b/velox/experimental/cudf/expression/ExpressionEvaluator.h @@ -30,6 +30,10 @@ #include #include +namespace facebook::velox::core { +class QueryConfig; +} + namespace facebook::velox::cudf_velox { // Holds either a non-owning cudf::column_view (zero-copy) or an owning @@ -71,13 +75,56 @@ void checkAllTrue( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/// Carries query-scoped evaluation settings that individual GPU functions need +/// but that are not part of the expression tree, most notably the session +/// timezone. Populated from the QueryConfig at expression-creation time and +/// attached to every CudfFunction so timezone-aware functions can match the CPU +/// path. Defaults represent "no session timezone" (UTC/GMT), matching the CPU +/// behavior when adjust_timestamp_to_session_timezone is off. +struct CudfDateTimeContext { + /// Session timezone name (QueryConfig::sessionTimezone), e.g. + /// "America/Los_Angeles". Empty means none. + std::string sessionTimezone; + /// Whether timezone-less timestamp conversions honor the session timezone + /// (QueryConfig::adjustTimestampToTimezone). + bool adjustTimestampToTimezone{false}; + /// Session start time in milliseconds since epoch + /// (QueryConfig::sessionStartTimeMs); used by now()/current_timestamp. + int64_t sessionStartTimeMs{0}; + + /// Returns true when extraction functions must convert the instant to the + /// session-local wall clock before reading a calendar field. + bool appliesSessionTimezone() const { + return adjustTimestampToTimezone && !sessionTimezone.empty(); + } +}; + +/// Builds a CudfDateTimeContext from the query config, copying the session +/// timezone, the adjust-to-session-timezone flag, and the session start time. +/// Operators that construct cuDF expressions build the context here so the +/// derivation lives in one place and timezone-aware functions match the CPU +/// path. +CudfDateTimeContext contextFromConfig(const core::QueryConfig& config); + class CudfFunction { public: virtual ~CudfFunction() = default; virtual ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const = 0; + + /// Attaches the query-scoped evaluation context. Called once after the + /// function is created. Functions that do not need it simply ignore context_. + void setContext(const CudfDateTimeContext& context) { + context_ = context; + } + + protected: + // Query-scoped evaluation context (session timezone and start time), attached + // via setContext. Timezone-aware functions read it; others ignore it. + CudfDateTimeContext context_; }; using CudfFunctionFactory = std::function( @@ -114,10 +161,12 @@ void registerCudfFunctions( /// Create a CudfFunction for the given name and expression. /// Returns nullptr if no registered function matches the expression's -/// signature. +/// signature. The context is attached to the created function so +/// timezone-aware functions can read the session timezone. std::shared_ptr createCudfFunction( const std::string& name, - const std::shared_ptr& expr); + const std::shared_ptr& expr, + const CudfDateTimeContext& context = {}); bool registerBuiltinFunctions(const std::string& prefix); @@ -142,7 +191,8 @@ using CudfExpressionEvaluatorCanEvaluate = using CudfExpressionEvaluatorCreate = std::function( std::shared_ptr expr, - const RowTypePtr& inputRowSchema)>; + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context)>; // Register a CudfExpression evaluator. // - name: unique identifier (e.g., "ast", "function", "my_custom"). @@ -161,7 +211,8 @@ class FunctionExpression : public CudfExpression { public: static std::shared_ptr create( const std::shared_ptr& expr, - const RowTypePtr& inputRowSchema); + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context = {}); // TODO (dm): A storage for keeping results in case this is a multiply // referenced subexpression (to do CSE) @@ -196,7 +247,8 @@ class FunctionExpression : public CudfExpression { std::shared_ptr createCudfExpression( std::shared_ptr expr, - const RowTypePtr& inputRowSchema); + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context = {}); /// Lightweight check if an expression tree is supported by any CUDF evaluator /// without initializing CudfExpression objects. diff --git a/velox/experimental/cudf/expression/JitExpression.cpp b/velox/experimental/cudf/expression/JitExpression.cpp index 3b0b5b2c64b..ef181a1c5b8 100644 --- a/velox/experimental/cudf/expression/JitExpression.cpp +++ b/velox/experimental/cudf/expression/JitExpression.cpp @@ -20,8 +20,9 @@ namespace facebook::velox::cudf_velox { JitExpression::JitExpression( std::shared_ptr expr, - const RowTypePtr& inputRowSchema) - : expr_{expr, inputRowSchema} {} + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context) + : expr_{expr, inputRowSchema, context} {} void JitExpression::close() { expr_.close(); @@ -86,8 +87,10 @@ void registerJitEvaluator(int priority) { [](std::shared_ptr expr) { return JitExpression::canEvaluate(expr); }, - [](std::shared_ptr expr, const RowTypePtr& row) { - return std::make_shared(std::move(expr), row); + [](std::shared_ptr expr, + const RowTypePtr& row, + const CudfDateTimeContext& context) { + return std::make_shared(std::move(expr), row, context); }, /*overwrite=*/false); } diff --git a/velox/experimental/cudf/expression/JitExpression.h b/velox/experimental/cudf/expression/JitExpression.h index 73ca9ccff09..ea9e143462c 100644 --- a/velox/experimental/cudf/expression/JitExpression.h +++ b/velox/experimental/cudf/expression/JitExpression.h @@ -32,7 +32,8 @@ class JitExpression : public CudfExpression { // precompute instructions and stores them JitExpression( std::shared_ptr expr, - const RowTypePtr& inputRowSchema); + const RowTypePtr& inputRowSchema, + const CudfDateTimeContext& context); // Evaluates the expression tree for the given input columns ColumnOrView eval( diff --git a/velox/experimental/cudf/expression/PrestoFunctions.cpp b/velox/experimental/cudf/expression/PrestoFunctions.cpp index b2ec4b716fd..5ca7bded5ab 100644 --- a/velox/experimental/cudf/expression/PrestoFunctions.cpp +++ b/velox/experimental/cudf/expression/PrestoFunctions.cpp @@ -24,6 +24,7 @@ #include "velox/common/base/Exceptions.h" #include "velox/expression/ConstantExpr.h" #include "velox/expression/FunctionSignature.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/vector/BaseVector.h" #include @@ -99,6 +100,7 @@ class SubstrFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { auto inputCol = asView(inputColumns[0]); @@ -153,7 +155,15 @@ void registerPrestoFunctions(const std::string& prefix) { registerCudfFunction( prefix + "date_add", - [](const std::string&, const std::shared_ptr& expr) { + [](const std::string&, const std::shared_ptr& expr) + -> std::shared_ptr { + if (isTimestampWithTimeZoneType(expr->inputs()[2]->type())) { + return std::make_shared< + prestosql::DateAddTimestampWithTimeZoneFunction>(expr); + } + if (expr->inputs()[2]->type()->isTimestamp()) { + return std::make_shared(expr); + } return std::make_shared(expr); }, {FunctionSignatureBuilder() @@ -161,9 +171,25 @@ void registerPrestoFunctions(const std::string& prefix) { .constantArgumentType("varchar") .argumentType("bigint") .argumentType("date") + .build(), + FunctionSignatureBuilder() + .returnType("timestamp") + .constantArgumentType("varchar") + .argumentType("bigint") + .argumentType("timestamp") + .build(), + FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .constantArgumentType("varchar") + .argumentType("bigint") + .argumentType("timestamp with time zone") .build()}, true, - prestosql::DateAddFunction::canEvaluate); + [](const std::shared_ptr& expr) { + return prestosql::DateAddFunction::canEvaluate(expr) || + prestosql::DateAddTimestampFunction::canEvaluate(expr) || + prestosql::DateAddTimestampWithTimeZoneFunction::canEvaluate(expr); + }); registerCudfFunction( prefix + "date_trunc", @@ -179,6 +205,11 @@ void registerPrestoFunctions(const std::string& prefix) { .returnType("date") .constantArgumentType("varchar") .argumentType("date") + .build(), + FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .constantArgumentType("varchar") + .argumentType("timestamp with time zone") .build()}, true, DateTruncFunction::canEvaluate); diff --git a/velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.cpp b/velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.cpp new file mode 100644 index 00000000000..5e3d80fc150 --- /dev/null +++ b/velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.cpp @@ -0,0 +1,311 @@ +/* + * 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 "velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h" +#include "velox/experimental/cudf/expression/TimezoneConversion.h" + +#include "velox/common/base/Exceptions.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" +#include "velox/type/tz/TimeZoneMap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +namespace facebook::velox::cudf_velox { +namespace { + +constexpr cudf::type_id kInt64 = cudf::type_id::INT64; +constexpr cudf::type_id kBool8 = cudf::type_id::BOOL8; +constexpr cudf::type_id kTsMillis = cudf::type_id::TIMESTAMP_MILLISECONDS; + +cudf::data_type int64Type() { + return cudf::data_type{kInt64}; +} + +cudf::numeric_scalar int64Scalar( + int64_t value, + rmm::cuda_stream_view stream) { + return cudf::numeric_scalar(value, true, stream); +} + +// Reinterprets an 8-byte-wide column (timestamp/duration/int64) as another +// 8-byte type without copying. +cudf::column_view bitcastColumn( + const cudf::column_view& view, + cudf::type_id id) { + return cudf::column_view{ + cudf::data_type{id}, + view.size(), + view.head(), + view.null_mask(), + view.null_count(), + view.offset()}; +} + +// Mirrors the CPU pack() range check: throws if any non-null millis value falls +// outside [kMinMillisUtc, kMaxMillisUtc]. +void checkMillisInRange( + const cudf::column_view& millis, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (millis.size() == 0 || millis.null_count() == millis.size()) { + return; + } + auto minScalar = cudf::reduce( + millis, + *cudf::make_min_aggregation(), + int64Type(), + stream, + mr); + auto maxScalar = cudf::reduce( + millis, + *cudf::make_max_aggregation(), + int64Type(), + stream, + mr); + const auto lo = static_cast*>(minScalar.get()) + ->value(stream); + const auto hi = static_cast*>(maxScalar.get()) + ->value(stream); + VELOX_USER_CHECK( + lo >= kMinMillisUtc && hi <= kMaxMillisUtc, + "TimestampWithTimeZone overflow: [{}, {}] ms", + lo, + hi); +} + +} // namespace + +std::unique_ptr tswtzZoneKey( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return cudf::binary_operation( + packed, + int64Scalar(kTimezoneMask, stream), + cudf::binary_operator::BITWISE_AND, + int64Type(), + stream, + mr); +} + +std::vector tswtzDistinctZoneKeys( + const cudf::column_view& perRowZoneKey, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto unique = cudf::distinct( + cudf::table_view{{perRowZoneKey}}, + {0}, + cudf::duplicate_keep_option::KEEP_ANY, + cudf::null_equality::EQUAL, + cudf::nan_equality::ALL_EQUAL, + stream, + mr); + auto uniqueKeys = unique->view().column(0); + auto uniqueValid = cudf::is_valid(uniqueKeys, stream, mr); + std::vector hostKeys(uniqueKeys.size()); + std::vector hostValid(uniqueKeys.size()); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostKeys.data(), + uniqueKeys.data(), + hostKeys.size() * sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream.value())); + CUDF_CUDA_TRY(cudaMemcpyAsync( + hostValid.data(), + uniqueValid->view().data(), + hostValid.size() * sizeof(int8_t), + cudaMemcpyDeviceToHost, + stream.value())); + stream.synchronize(); + + std::vector keys; + keys.reserve(uniqueKeys.size()); + for (cudf::size_type i = 0; i < uniqueKeys.size(); ++i) { + if (hostValid[i]) { // Skip the null zone key. + keys.push_back(static_cast(hostKeys[i])); + } + } + return keys; +} + +std::unique_ptr tswtzUtcInstant( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto millis = cudf::binary_operation( + packed, + int64Scalar(kMillisShift, stream), + cudf::binary_operator::SHIFT_RIGHT, + int64Type(), + stream, + mr); + return std::make_unique( + bitcastColumn(millis->view(), kTsMillis), stream, mr); +} + +std::unique_ptr tswtzOffsetSeconds( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto utcInstant = tswtzUtcInstant(packed, stream, mr); + auto perRowKey = tswtzZoneKey(packed, stream, mr); + auto keys = tswtzDistinctZoneKeys(perRowKey->view(), stream, mr); + + // Start all-null; fill each zone's rows. A null key matches no real key, so + // its rows keep the null default (CPU propagates null). + auto result = cudf::make_numeric_column( + int64Type(), packed.size(), cudf::mask_state::ALL_NULL, stream, mr); + for (const int16_t zoneKey : keys) { + auto offsetDuration = utcOffsetSeconds( + utcInstant->view(), tz::getTimeZoneName(zoneKey), stream, mr); + auto offsetSeconds = std::make_unique( + bitcastColumn(offsetDuration->view(), kInt64), stream, mr); + auto isThisZone = cudf::binary_operation( + perRowKey->view(), + int64Scalar(zoneKey, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + result = cudf::copy_if_else( + offsetSeconds->view(), result->view(), isThisZone->view(), stream, mr); + } + return result; +} + +std::unique_ptr tswtzLocalWallClock( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto utcInstant = tswtzUtcInstant(packed, stream, mr); + auto millis = bitcastColumn(utcInstant->view(), kInt64); + auto offsetSeconds = tswtzOffsetSeconds(packed, stream, mr); + auto offsetMillis = cudf::binary_operation( + offsetSeconds->view(), + int64Scalar(1'000, stream), + cudf::binary_operator::MUL, + int64Type(), + stream, + mr); + auto localMillis = cudf::binary_operation( + millis, + offsetMillis->view(), + cudf::binary_operator::ADD, + int64Type(), + stream, + mr); + return std::make_unique( + bitcastColumn(localMillis->view(), kTsMillis), stream, mr); +} + +std::unique_ptr tswtzLocalToUtc( + const cudf::column_view& localMillisTs, + const cudf::column_view& perRowZoneKey, + const std::vector& distinctKeys, + bool correctForward, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto result = cudf::make_timestamp_column( + cudf::data_type{kTsMillis}, + localMillisTs.size(), + cudf::mask_state::ALL_NULL, + stream, + mr); + for (const int16_t zoneKey : distinctKeys) { + auto isThisZone = cudf::binary_operation( + perRowZoneKey, + int64Scalar(zoneKey, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + // Mask out other zones' rows (null) so this zone's gap check ignores them. + auto nullTs = cudf::make_timestamp_column( + cudf::data_type{kTsMillis}, + localMillisTs.size(), + cudf::mask_state::ALL_NULL, + stream, + mr); + auto maskedLocal = cudf::copy_if_else( + localMillisTs, nullTs->view(), isThisZone->view(), stream, mr); + const auto zoneName = tz::getTimeZoneName(zoneKey); + // date_trunc snaps to local midnight (never a gap) and uses the throwing + // path; date_add can land a same-wall-clock time in a spring-forward gap + // and uses the correcting path (matches addToTimestampWithTimezone). + auto utc = correctForward + ? toUtcTimestampCorrecting(maskedLocal->view(), zoneName, stream, mr) + : toUtcTimestamp(maskedLocal->view(), zoneName, stream, mr); + result = cudf::copy_if_else( + utc->view(), result->view(), isThisZone->view(), stream, mr); + } + return result; +} + +std::unique_ptr tswtzPack( + const cudf::column_view& utcInstant, + const cudf::column_view& perRowZoneKey, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + // Normalize to a millisecond instant, then bit-cast to raw int64 millis. + std::unique_ptr millisTs; + cudf::column_view millisView; + if (utcInstant.type().id() == kTsMillis) { + millisView = utcInstant; + } else { + millisTs = cudf::cast(utcInstant, cudf::data_type{kTsMillis}, stream, mr); + millisView = millisTs->view(); + } + auto millis = std::make_unique( + bitcastColumn(millisView, kInt64), stream, mr); + checkMillisInRange(millis->view(), stream, mr); + + auto shifted = cudf::binary_operation( + millis->view(), + int64Scalar(kMillisShift, stream), + cudf::binary_operator::SHIFT_LEFT, + int64Type(), + stream, + mr); + auto maskedKey = cudf::binary_operation( + perRowZoneKey, + int64Scalar(kTimezoneMask, stream), + cudf::binary_operator::BITWISE_AND, + int64Type(), + stream, + mr); + return cudf::binary_operation( + shifted->view(), + maskedKey->view(), + cudf::binary_operator::BITWISE_OR, + int64Type(), + stream, + mr); +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h b/velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h new file mode 100644 index 00000000000..f09b8390f6e --- /dev/null +++ b/velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h @@ -0,0 +1,98 @@ +/* + * 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. + */ + +#pragma once + +#include +#include + +#include +#include + +#include +#include +#include + +namespace facebook::velox::cudf_velox { + +/// Column-level primitives for the packed Presto TIMESTAMP WITH TIME ZONE +/// representation (upper 52 bits UTC millis, lower 12 bits time-zone key). All +/// operate per row and support columns that mix zone keys, matching the rest of +/// the GPU TSWTZ family. Shared by the timezone functions and the +/// date_trunc/date_add TSWTZ overloads. + +/// Returns the per-row zone key (INT64, nulls preserved) of a packed column: +/// packed & kTimezoneMask. A null packed row yields a null key. +std::unique_ptr tswtzZoneKey( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Returns the distinct non-null zone keys present in a per-row zone-key column +/// (as produced by tswtzZoneKey). Performs one device-to-host synchronization. +std::vector tswtzDistinctZoneKeys( + const cudf::column_view& perRowZoneKey, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Returns the UTC instant (TIMESTAMP_MILLISECONDS) of a packed column: +/// arithmetic (packed >> 12) bit-cast to a timestamp column. +std::unique_ptr tswtzUtcInstant( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Returns the per-row UT offset in whole seconds (INT64) for a packed column +/// that may mix zone keys. Null rows stay null. O(number of distinct zones) +/// device passes. +std::unique_ptr tswtzOffsetSeconds( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Returns the per-row local wall clock (TIMESTAMP_MILLISECONDS) of a packed +/// column, applying each row's own zone offset (multi-zone). +std::unique_ptr tswtzLocalWallClock( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Converts a per-row local wall-clock column back to UTC instants +/// (TIMESTAMP_MILLISECONDS), applying each row's own zone. localMillisTs is a +/// TIMESTAMP_MILLISECONDS local column; perRowZoneKey/distinctKeys identify +/// each row's zone (from tswtzZoneKey/tswtzDistinctZoneKeys). When +/// correctForward is false a spring-forward-gap local time throws (matches +/// Timestamp::toGMT); when true it resolves forward without throwing (local +/// minus the pre-transition offset, matching addToTimestampWithTimezone). +/// Overlaps always resolve to the earliest instant. +std::unique_ptr tswtzLocalToUtc( + const cudf::column_view& localMillisTs, + const cudf::column_view& perRowZoneKey, + const std::vector& distinctKeys, + bool correctForward, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Repacks a UTC-instant column (any timestamp resolution; cast to millis) plus +/// a per-row zone-key column into a packed TSWTZ INT64 column. Throws if any +/// non-null instant falls outside the representable millis range. +std::unique_ptr tswtzPack( + const cudf::column_view& utcInstant, + const cudf::column_view& perRowZoneKey, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/TimezoneConversion.cpp b/velox/experimental/cudf/expression/TimezoneConversion.cpp new file mode 100644 index 00000000000..4508dfb7eed --- /dev/null +++ b/velox/experimental/cudf/expression/TimezoneConversion.cpp @@ -0,0 +1,511 @@ +/* + * 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 "velox/experimental/cudf/expression/TimezoneConversion.h" + +#include "velox/common/base/Exceptions.h" +#include "velox/external/date/date.h" +#include "velox/external/tzdb/time_zone.h" +#include "velox/type/tz/TimeZoneMap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { + +// Maps a cudf timestamp type to the duration type of the same resolution, used +// to add a seconds offset back at the input's precision. +cudf::type_id durationTypeIdForTimestamp(cudf::type_id timestampType) { + switch (timestampType) { + case cudf::type_id::TIMESTAMP_SECONDS: + return cudf::type_id::DURATION_SECONDS; + case cudf::type_id::TIMESTAMP_MILLISECONDS: + return cudf::type_id::DURATION_MILLISECONDS; + case cudf::type_id::TIMESTAMP_MICROSECONDS: + return cudf::type_id::DURATION_MICROSECONDS; + case cudf::type_id::TIMESTAMP_NANOSECONDS: + return cudf::type_id::DURATION_NANOSECONDS; + default: + VELOX_FAIL( + "Unsupported timestamp resolution for timezone conversion: {}", + static_cast(timestampType)); + } +} + +// Re-applies the input's null mask onto an offset column. The gather that +// produces the offset yields a fully-valid column regardless of the input's +// validity, so this is the single place that restores it -- a null instant must +// yield a null offset so callers (timezone_hour/minute, to_iso8601, +// format_datetime) propagate it. +std::unique_ptr withInputNullMask( + std::unique_ptr offset, + const cudf::column_view& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (input.null_count() > 0) { + offset->set_null_mask( + cudf::copy_bitmask(input, stream, mr), input.null_count()); + } + return offset; +} + +// A single offset interval: at `instant` (UTC seconds) the zone's UTC offset +// becomes `offset` (seconds) and stays constant until the next transition. +struct Transition { + int64_t instant; + int64_t offset; +}; + +// Walks the zone's daylight-savings transitions from 1700 to 2400. That window +// covers the representable range of nanosecond timestamps (~1678-2262) with no +// folding; instants beyond it reuse the last interval's offset. +std::vector enumerateTransitions(const tz::TimeZone* timeZone) { + using std::chrono::seconds; + + // Offset-only zones (e.g. "+05:30") have a single constant offset. + if (auto fixed = timeZone->offset(); fixed.has_value()) { + return {Transition{0, std::chrono::duration_cast(*fixed).count()}}; + } + + const auto* zone = timeZone->tz(); + VELOX_CHECK_NOT_NULL( + zone, + "Time zone has neither a fixed offset nor a database entry: {}", + timeZone->name()); + + const auto yearStart = [](int year) { + return std::chrono::duration_cast( + date::sys_days{date::year{year} / date::January / 1} + .time_since_epoch()) + .count(); + }; + const int64_t horizonSeconds = yearStart(2400); + + std::vector transitions; + int64_t probe = yearStart(1700); + while (true) { + auto info = zone->get_info(date::sys_seconds{seconds{probe}}); + transitions.push_back( + {info.begin.time_since_epoch().count(), info.offset.count()}); + + const int64_t endSeconds = info.end.time_since_epoch().count(); + if (info.end == date::sys_seconds::max() || endSeconds >= horizonSeconds || + endSeconds <= probe) { + break; + } + probe = endSeconds; + } + return transitions; +} + +// Copies a host vector to a new device column of the given type. The element +// type T must match the column's physical representation (int64_t for +// TIMESTAMP_SECONDS/DURATION_SECONDS, int8_t for BOOL8). +template +std::unique_ptr makeDeviceColumn( + const std::vector& host, + cudf::type_id typeId, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto size = static_cast(host.size()); + auto column = cudf::make_fixed_width_column( + cudf::data_type{typeId}, size, cudf::mask_state::UNALLOCATED, stream, mr); + if (size > 0) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + column->mutable_view().data(), + host.data(), + host.size() * sizeof(T), + cudaMemcpyHostToDevice, + stream.value())); + } + return column; +} + +// Builds the forward (UTC-keyed) table [instant (TIMESTAMP_SECONDS), offset +// (DURATION_SECONDS)] from the zone's transitions. The first key is forced to +// INT64_MIN so the active-interval index (upper_bound - 1) is never out of +// range; instants after the last transition reuse its offset. Synchronizes the +// stream before returning so the host vectors outlive the async uploads. +std::unique_ptr buildForwardTable( + const std::vector& transitions, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + std::vector keys; + std::vector offsets; + keys.reserve(transitions.size()); + offsets.reserve(transitions.size()); + for (const auto& transition : transitions) { + keys.push_back(transition.instant); + offsets.push_back(transition.offset); + } + keys.front() = std::numeric_limits::min(); + + std::vector> columns; + columns.push_back( + makeDeviceColumn(keys, cudf::type_id::TIMESTAMP_SECONDS, stream, mr)); + columns.push_back( + makeDeviceColumn(offsets, cudf::type_id::DURATION_SECONDS, stream, mr)); + stream.synchronize(); + return std::make_unique(std::move(columns)); +} + +// Builds the local-keyed inverse table [localInstant (TIMESTAMP_SECONDS), +// offset (DURATION_SECONDS), gap (BOOL8)] from the zone's transitions. A +// transition from prevOffset to curOffset at UTC instant `inst` shifts the wall +// clock between inst+prevOffset and inst+curOffset. A forward shift (curOffset +// > prevOffset, spring forward) makes that local range nonexistent, so it is +// flagged as a gap; a backward shift (fall back) makes it ambiguous, and +// keeping the pre-transition offset over the overlap matches toGMT's kEarliest +// choice (so only the later local boundary needs a breakpoint). Synchronizes +// the stream before returning so the host vectors outlive the async uploads. +std::unique_ptr buildInverseTable( + const std::vector& transitions, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + // Force the first key below every representable input so the active-interval + // index (upper_bound - 1) is never out of range. + constexpr int64_t kFloor = std::numeric_limits::min(); + + struct Breakpoint { + int64_t key; + int64_t offset; + int8_t gap; + }; + std::vector breakpoints; + breakpoints.push_back({kFloor, transitions.front().offset, 0}); + for (size_t i = 1; i < transitions.size(); ++i) { + const int64_t prevOffset = transitions[i - 1].offset; + const int64_t curOffset = transitions[i].offset; + const int64_t localPrev = transitions[i].instant + prevOffset; + const int64_t localCur = transitions[i].instant + curOffset; + if (curOffset > prevOffset) { + breakpoints.push_back({localPrev, prevOffset, 1}); + breakpoints.push_back({localCur, curOffset, 0}); + } else if (curOffset < prevOffset) { + breakpoints.push_back({localPrev, curOffset, 0}); + } + } + std::stable_sort( + breakpoints.begin(), + breakpoints.end(), + [](const Breakpoint& a, const Breakpoint& b) { return a.key < b.key; }); + + std::vector keys; + std::vector offsets; + std::vector gaps; + for (const auto& breakpoint : breakpoints) { + // On duplicate local keys keep the last; the stable sort preserves the + // emission order, which is the intended precedence. + if (!keys.empty() && keys.back() == breakpoint.key) { + offsets.back() = breakpoint.offset; + gaps.back() = breakpoint.gap; + } else { + keys.push_back(breakpoint.key); + offsets.push_back(breakpoint.offset); + gaps.push_back(breakpoint.gap); + } + } + + std::vector> columns; + columns.push_back( + makeDeviceColumn(keys, cudf::type_id::TIMESTAMP_SECONDS, stream, mr)); + columns.push_back( + makeDeviceColumn(offsets, cudf::type_id::DURATION_SECONDS, stream, mr)); + columns.push_back(makeDeviceColumn(gaps, cudf::type_id::BOOL8, stream, mr)); + stream.synchronize(); + return std::make_unique(std::move(columns)); +} + +// Returns the active-interval row index (upper_bound - 1, INT32) for each +// timestamp against a table's sorted key column. The key is truncated to whole +// seconds because offsets only change on second boundaries. Both tables force +// the first key to INT64_MIN, so the result is always a valid index. +std::unique_ptr activeIntervalIndices( + const cudf::column_view& transitionKeys, + const cudf::column_view& timestamps, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto key = cudf::cast( + timestamps, + cudf::data_type{cudf::type_id::TIMESTAMP_SECONDS}, + stream, + mr); + auto positions = cudf::upper_bound( + cudf::table_view{{transitionKeys}}, + cudf::table_view{{key->view()}}, + {cudf::order::ASCENDING}, + {cudf::null_order::AFTER}, + stream, + mr); + auto one = cudf::numeric_scalar(1, true, stream); + return cudf::binary_operation( + positions->view(), + one, + cudf::binary_operator::SUB, + cudf::data_type{cudf::type_id::INT32}, + stream, + mr); +} + +// Per-zone forward (UTC-keyed) and inverse (local-keyed, gap-flagged) offset +// tables, built once from Velox's time zone database and cached for the process +// lifetime. The forward table answers UTC->local and per-row offset queries; +// the inverse table answers local->UTC with the daylight-savings policy of +// Timestamp::toGMT baked in. +class OffsetTable { + public: + OffsetTable( + std::unique_ptr forward, + std::unique_ptr inverse) + : forward_(std::move(forward)), inverse_(std::move(inverse)) {} + + // Returns the table for `timeZone`, building it on first use and caching it + // by zone id for the process lifetime. Thread-safe. + static std::shared_ptr get(const tz::TimeZone* timeZone); + + // Per-row UT offset (DURATION_SECONDS) at each UTC instant; the input null + // mask is re-applied so a null instant yields a null offset. + std::unique_ptr utcOffset( + const cudf::column_view& utcTimestamps, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + // utc + offset, at the input's resolution. + std::unique_ptr toLocal( + const cudf::column_view& utcTimestamps, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + // local - offset. When correctForward is false, raises a user error on a + // nonexistent (spring-forward gap) local time; when true, keeps the computed + // instant (local minus the pre-transition offset the gap interval stores), + // matching addToTimestampWithTimezone. An ambiguous (fall-back overlap) local + // always resolves to the earliest instant. Null rows are never treated as + // gaps. + std::unique_ptr toUtc( + const cudf::column_view& localTimestamps, + bool correctForward, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + private: + // [instant (TIMESTAMP_SECONDS), offset (DURATION_SECONDS)], UTC-keyed. + std::unique_ptr forward_; + // [instant (TIMESTAMP_SECONDS), offset (DURATION_SECONDS), gap (BOOL8)], + // local-keyed. + std::unique_ptr inverse_; +}; + +// static +std::shared_ptr OffsetTable::get( + const tz::TimeZone* timeZone) { + VELOX_CHECK_NOT_NULL(timeZone, "Time zone must not be null"); + static folly::Synchronized< + std::unordered_map>> + cache; + + const auto id = timeZone->id(); + { + auto locked = cache.rlock(); + if (auto it = locked->find(id); it != locked->end()) { + return it->second; + } + } + // Build on the default stream and current resource so the cached device + // tables do not depend on any caller's stream or memory resource. + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + auto transitions = enumerateTransitions(timeZone); + VELOX_CHECK(!transitions.empty()); + auto table = std::make_shared( + buildForwardTable(transitions, stream, mr), + buildInverseTable(transitions, stream, mr)); + // Another thread may have inserted the same zone meanwhile; emplace keeps the + // existing entry and discards this build. + return cache.wlock()->emplace(id, std::move(table)).first->second; +} + +std::unique_ptr OffsetTable::utcOffset( + const cudf::column_view& utcTimestamps, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + auto indices = activeIntervalIndices( + forward_->view().column(0), utcTimestamps, stream, mr); + auto gathered = cudf::gather( + cudf::table_view{{forward_->view().column(1)}}, + indices->view(), + cudf::out_of_bounds_policy::DONT_CHECK, + stream, + mr); + auto columns = gathered->release(); + return withInputNullMask(std::move(columns[0]), utcTimestamps, stream, mr); +} + +std::unique_ptr OffsetTable::toLocal( + const cudf::column_view& utcTimestamps, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + auto offsetSeconds = utcOffset(utcTimestamps, stream, mr); + + // Add the offset at the input's resolution so sub-second precision survives. + const auto durationType = + durationTypeIdForTimestamp(utcTimestamps.type().id()); + std::unique_ptr offsetConverted; + cudf::column_view offsetView = offsetSeconds->view(); + if (durationType != cudf::type_id::DURATION_SECONDS) { + offsetConverted = cudf::cast( + offsetSeconds->view(), cudf::data_type{durationType}, stream, mr); + offsetView = offsetConverted->view(); + } + return cudf::binary_operation( + utcTimestamps, + offsetView, + cudf::binary_operator::ADD, + utcTimestamps.type(), + stream, + mr); +} + +std::unique_ptr OffsetTable::toUtc( + const cudf::column_view& localTimestamps, + bool correctForward, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + auto indices = activeIntervalIndices( + inverse_->view().column(0), localTimestamps, stream, mr); + auto gathered = cudf::gather( + cudf::table_view{ + {inverse_->view().column(1), inverse_->view().column(2)}}, + indices->view(), + cudf::out_of_bounds_policy::DONT_CHECK, + stream, + mr); + auto gatheredView = gathered->view(); + + // utc = local - offset, at the input's resolution so sub-second precision + // survives. + const auto durationType = + durationTypeIdForTimestamp(localTimestamps.type().id()); + auto offset = cudf::cast( + gatheredView.column(0), cudf::data_type{durationType}, stream, mr); + auto result = cudf::binary_operation( + localTimestamps, + offset->view(), + cudf::binary_operator::SUB, + localTimestamps.type(), + stream, + mr); + + // A nonexistent local time (spring-forward gap) has no UTC instant. The + // throwing path matches CPU's toGMT and fails; the correcting path keeps + // `result` (local minus the pre-transition offset the gap interval stores), + // matching addToTimestampWithTimezone. Null rows are not gaps, so mask them + // out before the gap check. + if (!correctForward) { + cudf::column_view gap = gatheredView.column(1); + std::unique_ptr maskedGap; + if (localTimestamps.nullable() && localTimestamps.null_count() > 0) { + auto valid = cudf::is_valid(localTimestamps, stream, mr); + maskedGap = cudf::binary_operation( + gap, + valid->view(), + cudf::binary_operator::LOGICAL_AND, + cudf::data_type{cudf::type_id::BOOL8}, + stream, + mr); + gap = maskedGap->view(); + } + auto anyGap = cudf::reduce( + gap, + *cudf::make_any_aggregation(), + cudf::data_type{cudf::type_id::BOOL8}, + stream, + mr); + auto& anyGapScalar = static_cast&>(*anyGap); + if (anyGapScalar.is_valid(stream) && anyGapScalar.value(stream)) { + VELOX_USER_FAIL( + "Cannot convert local time to UTC: the time does not exist in the " + "time zone (daylight savings gap)"); + } + } + return result; +} + +} // namespace + +std::unique_ptr utcOffsetSeconds( + const cudf::column_view& utcTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return OffsetTable::get(tz::locateZone(timezoneName)) + ->utcOffset(utcTimestamps, stream, mr); +} + +std::unique_ptr toLocalTimestamp( + const cudf::column_view& utcTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return OffsetTable::get(tz::locateZone(timezoneName)) + ->toLocal(utcTimestamps, stream, mr); +} + +std::unique_ptr toUtcTimestamp( + const cudf::column_view& localTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return OffsetTable::get(tz::locateZone(timezoneName)) + ->toUtc(localTimestamps, /*correctForward=*/false, stream, mr); +} + +std::unique_ptr toUtcTimestampCorrecting( + const cudf::column_view& localTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return OffsetTable::get(tz::locateZone(timezoneName)) + ->toUtc(localTimestamps, /*correctForward=*/true, stream, mr); +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/TimezoneConversion.h b/velox/experimental/cudf/expression/TimezoneConversion.h new file mode 100644 index 00000000000..82e74c4045a --- /dev/null +++ b/velox/experimental/cudf/expression/TimezoneConversion.h @@ -0,0 +1,98 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include + +#include + +#include +#include + +namespace facebook::velox::cudf_velox { + +/// Converts a column of UTC timestamps to the local wall-clock instants of +/// `timezoneName`, DST-aware. The returned column has the same timestamp type +/// (and resolution) as the input; each value is shifted by that instant's UTC +/// offset so that a subsequent cudf::datetime::extract_datetime_component or +/// cudf::strings::from_timestamps reads the local calendar fields -- matching +/// the Velox CPU path, which converts the instant to the session timezone +/// before extracting. +/// +/// The offset comes from a per-zone transition table built from Velox's own +/// time zone database (the same source the CPU path uses) and cached for the +/// process lifetime; a sorted search (cudf::upper_bound) + cudf::gather selects +/// each row's offset and cudf::binary_operation adds it. Instants after the +/// last codified transition reuse its offset. +/// +/// Null rows propagate. This is the inverse of toUtcTimestamp. +std::unique_ptr toLocalTimestamp( + const cudf::column_view& utcTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Converts a column of wall-clock local timestamps in `timezoneName` to the +/// UTC instants they denote, DST-aware and matching the Velox CPU path +/// (Timestamp::toGMT). A local time that falls in a spring-forward gap does not +/// exist, so the conversion raises a user error; a local time in a fall-back +/// overlap is ambiguous and resolves to the earliest instant. The returned +/// column keeps the input's timestamp resolution and null mask. Null rows are +/// never treated as gaps, so a caller that converts only some rows can null out +/// the rest to exclude them from the gap check. +/// +/// This is the inverse of toLocalTimestamp and reads the same cached, +/// tzdb-sourced transition table, in its local-keyed form with a gap flag per +/// breakpoint. Building the table from Velox's own time zone database -- the +/// source the CPU path uses -- makes the gap and overlap boundaries match +/// exactly. The conversion is a sorted search (cudf::upper_bound) plus an +/// offset subtract. +std::unique_ptr toUtcTimestamp( + const cudf::column_view& localTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Like toUtcTimestamp, but a local time in a spring-forward gap resolves +/// forward instead of raising, matching addToTimestampWithTimezone's day-and- +/// above path (correct_nonexistent_time followed by to_sys(kEarliest)). That +/// combination shifts the nonexistent local by the gap size and then subtracts +/// the post-transition offset, which reduces to subtracting the pre-transition +/// offset -- exactly the offset the local-keyed transition table already stores +/// across the gap. Overlaps still resolve to the earliest instant, and null +/// rows propagate. Reads the same cached table as toUtcTimestamp. +std::unique_ptr toUtcTimestampCorrecting( + const cudf::column_view& localTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +/// Returns the per-row UT offset (DURATION_SECONDS), DST-aware, for the given +/// timezone at each UTC instant -- i.e. local = utc + offset. This is the +/// primitive behind toLocalTimestamp; it is also used directly to render +/// timezone offsets (timezone_hour/minute, to_iso8601, format_datetime). Reads +/// the same cached, tzdb-sourced transition table as toLocalTimestamp. Null +/// rows in the input propagate to the result. +std::unique_ptr utcOffsetSeconds( + const cudf::column_view& utcTimestamps, + std::string_view timezoneName, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/prestosql/DateAddFunction.cpp b/velox/experimental/cudf/expression/prestosql/DateAddFunction.cpp index 85bb538766e..a1be88f5190 100644 --- a/velox/experimental/cudf/expression/prestosql/DateAddFunction.cpp +++ b/velox/experimental/cudf/expression/prestosql/DateAddFunction.cpp @@ -15,10 +15,13 @@ */ #include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/expression/AstUtils.h" +#include "velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h" +#include "velox/experimental/cudf/expression/TimezoneConversion.h" #include "velox/experimental/cudf/expression/prestosql/DateAddFunction.h" #include "velox/expression/ConstantExpr.h" #include "velox/functions/prestosql/DateTimeFunctions.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/vector/ConstantVector.h" #include @@ -146,6 +149,124 @@ std::unique_ptr scaleToInt32( mr); } +// Maps a timestamp data type to the duration type of the same resolution. +cudf::data_type durationTypeForTimestamp(cudf::data_type timestampType) { + switch (timestampType.id()) { + case cudf::type_id::TIMESTAMP_SECONDS: + return cudf::data_type(cudf::type_id::DURATION_SECONDS); + case cudf::type_id::TIMESTAMP_MILLISECONDS: + return cudf::data_type(cudf::type_id::DURATION_MILLISECONDS); + case cudf::type_id::TIMESTAMP_MICROSECONDS: + return cudf::data_type(cudf::type_id::DURATION_MICROSECONDS); + case cudf::type_id::TIMESTAMP_NANOSECONDS: + return cudf::data_type(cudf::type_id::DURATION_NANOSECONDS); + default: + VELOX_FAIL( + "date_add requires a timestamp column, got cudf type id: {}", + static_cast(timestampType.id())); + } +} + +// Month, quarter, and year add whole calendar months. +bool isMonthBasedUnit(DateTimeUnit unit) { + return unit == DateTimeUnit::kMonth || unit == DateTimeUnit::kQuarter || + unit == DateTimeUnit::kYear; +} + +// Adds value units to an instant column (timestamp of any resolution), +// preserving the instant's resolution. valueCol carries the per-row increment +// when present; otherwise literalValue/literalValid supply a constant. Sub-day +// and day/week units add a duration; month/quarter/year add calendar months. +// The raw value is validated to int32 (matching CPU checkValueInInt32Range). +std::unique_ptr addUnitToInstant( + cudf::column_view instant, + std::optional valueCol, + int64_t literalValue, + bool literalValid, + DateTimeUnit unit, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (isMonthBasedUnit(unit)) { + const int32_t scale = unit == DateTimeUnit::kQuarter + ? 3 + : (unit == DateTimeUnit::kYear ? 12 : 1); + if (valueCol.has_value()) { + auto months = scaleToInt32(*valueCol, instant, scale, stream, mr); + return cudf::datetime::add_calendrical_months( + instant, months->view(), stream, mr); + } + cudf::numeric_scalar months( + checkedScaleValue(literalValue, scale), literalValid, stream, mr); + return cudf::datetime::add_calendrical_months(instant, months, stream, mr); + } + + const auto nativeDuration = durationTypeForTimestamp(instant.type()); + std::unique_ptr duration; + + if (unit == DateTimeUnit::kDay || unit == DateTimeUnit::kWeek) { + // Day/week use a whole-day (int32-rep) duration, matching the DATE path. + const int32_t scale = unit == DateTimeUnit::kWeek ? 7 : 1; + std::unique_ptr daysInt; + if (valueCol.has_value()) { + daysInt = scaleToInt32(*valueCol, instant, scale, stream, mr); + } else { + cudf::numeric_scalar days( + checkedScaleValue(literalValue, scale), literalValid, stream, mr); + daysInt = cudf::make_column_from_scalar(days, instant.size(), stream, mr); + } + auto durationDays = cudf::cast( + daysInt->view(), + cudf::data_type(cudf::type_id::DURATION_DAYS), + stream, + mr); + duration = cudf::cast(durationDays->view(), nativeDuration, stream, mr); + } else { + // Sub-day: build an int64 tick count (value * multiplier) in seconds (or + // milliseconds for the millisecond unit), then rescale to the instant. + const bool millis = unit == DateTimeUnit::kMillisecond; + const auto subDayDuration = cudf::data_type( + millis ? cudf::type_id::DURATION_MILLISECONDS + : cudf::type_id::DURATION_SECONDS); + const int64_t multiplier = unit == DateTimeUnit::kHour + ? 3600 + : (unit == DateTimeUnit::kMinute ? 60 : 1); + + std::unique_ptr countInt64; + if (valueCol.has_value()) { + checkValueRange(*valueCol, instant, stream, mr); + if (multiplier == 1) { + countInt64 = cudf::cast( + *valueCol, cudf::data_type(cudf::type_id::INT64), stream, mr); + } else { + cudf::numeric_scalar multScalar(multiplier, true, stream, mr); + countInt64 = cudf::binary_operation( + *valueCol, + multScalar, + cudf::binary_operator::MUL, + cudf::data_type(cudf::type_id::INT64), + stream, + mr); + } + } else { + checkDateAddValueInInt32Range(literalValue); + cudf::numeric_scalar count( + literalValue * multiplier, literalValid, stream, mr); + countInt64 = + cudf::make_column_from_scalar(count, instant.size(), stream, mr); + } + auto subDay = cudf::cast(countInt64->view(), subDayDuration, stream, mr); + duration = cudf::cast(subDay->view(), nativeDuration, stream, mr); + } + + return cudf::binary_operation( + instant, + duration->view(), + cudf::binary_operator::ADD, + instant.type(), + stream, + mr); +} + } // namespace bool DateAddFunction::canEvaluate( @@ -200,6 +321,7 @@ DateAddFunction::DateAddFunction( ColumnOrView DateAddFunction::eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const { // Walk the non-literal inputs in argument order. Constants were captured at @@ -283,4 +405,226 @@ ColumnOrView DateAddFunction::evalMonthBased( dateCol, months->view(), stream, mr); } +bool DateAddTimestampFunction::canEvaluate( + const std::shared_ptr& expr) { + if (expr->inputs().size() != 3 || !expr->type()->isTimestamp() || + !expr->inputs()[2]->type()->isTimestamp()) { + return false; + } + + auto valueExpr = + std::dynamic_pointer_cast(expr->inputs()[1]); + auto timestampExpr = + std::dynamic_pointer_cast(expr->inputs()[2]); + if (valueExpr && timestampExpr) { + return false; + } + + auto unitString = constantVarcharValue(expr->inputs()[0]); + if (!unitString.has_value()) { + return false; + } + return functions::fromDateTimeUnitString(*unitString, false).has_value(); +} + +DateAddTimestampFunction::DateAddTimestampFunction( + const std::shared_ptr& expr) { + using velox::exec::ConstantExpr; + VELOX_CHECK( + canEvaluate(expr), + "date_add expression cannot be evaluated by " + "prestosql::DateAddTimestampFunction"); + + auto unitString = constantVarcharValue(expr->inputs()[0]); + unit_ = *functions::fromDateTimeUnitString(*unitString, true); + + auto valueExpr = std::dynamic_pointer_cast(expr->inputs()[1]); + valueIsLiteral_ = valueExpr != nullptr; + timestampIsLiteral_ = + std::dynamic_pointer_cast(expr->inputs()[2]) != nullptr; + + if (valueIsLiteral_) { + literalValueIsValid_ = !valueExpr->value()->isNullAt(0); + if (literalValueIsValid_) { + literalValue_ = + valueExpr->value()->as>()->value(); + } + } + if (timestampIsLiteral_) { + literalTimestamp_ = makeScalarFromConstantExpr(expr->inputs()[2]); + } +} + +ColumnOrView DateAddTimestampFunction::eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + size_t idx = 0; + + std::optional valueCol; + if (!valueIsLiteral_) { + valueCol = asView(inputColumns[idx++]); + } + + std::unique_ptr literalTimestampColumn; + cudf::column_view timestampCol; + if (!timestampIsLiteral_) { + timestampCol = asView(inputColumns[idx++]); + } else { + VELOX_CHECK_NOT_NULL(literalTimestamp_); + VELOX_CHECK( + valueCol.has_value(), + "date_add with only literal inputs is not supported"); + literalTimestampColumn = cudf::make_column_from_scalar( + *literalTimestamp_, valueCol->size(), stream, mr); + timestampCol = literalTimestampColumn->view(); + } + + // Under a UTC session, add directly to the instant. When the session applies + // a timezone, addToTimestamp adds on the local wall clock for ALL units (a + // sub-day add can cross a DST boundary), then converts back to UTC. + if (!context_.appliesSessionTimezone()) { + return addUnitToInstant( + timestampCol, + valueCol, + literalValue_, + literalValueIsValid_, + unit_, + stream, + mr); + } + + const std::string& zone = context_.sessionTimezone; + auto local = toLocalTimestamp(timestampCol, zone, stream, mr); + auto added = addUnitToInstant( + local->view(), + valueCol, + literalValue_, + literalValueIsValid_, + unit_, + stream, + mr); + // A spring-forward gap raises in toUtcTimestamp, matching CPU toGMT parity. + return toUtcTimestamp(added->view(), zone, stream, mr); +} + +bool DateAddTimestampWithTimeZoneFunction::canEvaluate( + const std::shared_ptr& expr) { + if (expr->inputs().size() != 3 || + !isTimestampWithTimeZoneType(expr->type()) || + !isTimestampWithTimeZoneType(expr->inputs()[2]->type())) { + return false; + } + + auto valueExpr = + std::dynamic_pointer_cast(expr->inputs()[1]); + auto timestampExpr = + std::dynamic_pointer_cast(expr->inputs()[2]); + if (valueExpr && timestampExpr) { + return false; + } + + auto unitString = constantVarcharValue(expr->inputs()[0]); + if (!unitString.has_value()) { + return false; + } + return functions::fromDateTimeUnitString(*unitString, false).has_value(); +} + +DateAddTimestampWithTimeZoneFunction::DateAddTimestampWithTimeZoneFunction( + const std::shared_ptr& expr) { + using velox::exec::ConstantExpr; + VELOX_CHECK( + canEvaluate(expr), + "date_add expression cannot be evaluated by " + "prestosql::DateAddTimestampWithTimeZoneFunction"); + + auto unitString = constantVarcharValue(expr->inputs()[0]); + unit_ = *functions::fromDateTimeUnitString(*unitString, true); + + auto valueExpr = std::dynamic_pointer_cast(expr->inputs()[1]); + valueIsLiteral_ = valueExpr != nullptr; + timestampIsLiteral_ = + std::dynamic_pointer_cast(expr->inputs()[2]) != nullptr; + + if (valueIsLiteral_) { + literalValueIsValid_ = !valueExpr->value()->isNullAt(0); + if (literalValueIsValid_) { + literalValue_ = + valueExpr->value()->as>()->value(); + } + } + if (timestampIsLiteral_) { + // A TSWTZ constant is physically a bigint (the packed value), so the + // dispatched scalar is a numeric int64 scalar holding the packed instant. + literalTimestamp_ = makeScalarFromConstantExpr(expr->inputs()[2]); + } +} + +ColumnOrView DateAddTimestampWithTimeZoneFunction::eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + size_t idx = 0; + + std::optional valueCol; + if (!valueIsLiteral_) { + valueCol = asView(inputColumns[idx++]); + } + + std::unique_ptr literalPackedColumn; + cudf::column_view packedCol; + if (!timestampIsLiteral_) { + packedCol = asView(inputColumns[idx++]); + } else { + VELOX_CHECK_NOT_NULL(literalTimestamp_); + VELOX_CHECK( + valueCol.has_value(), + "date_add with only literal inputs is not supported"); + literalPackedColumn = cudf::make_column_from_scalar( + *literalTimestamp_, valueCol->size(), stream, mr); + packedCol = literalPackedColumn->view(); + } + + auto perRowZoneKey = tswtzZoneKey(packedCol, stream, mr); + + // Sub-day units add directly to the UTC instant (DST-agnostic); day-and-above + // units add on each row's local wall clock, then convert back to UTC, + // resolving a spring-forward gap forward instead of throwing. This matches + // addToTimestampWithTimezone. + if (functions::isTimeUnit(unit_)) { + auto utcInstant = tswtzUtcInstant(packedCol, stream, mr); + auto added = addUnitToInstant( + utcInstant->view(), + valueCol, + literalValue_, + literalValueIsValid_, + unit_, + stream, + mr); + return tswtzPack(added->view(), perRowZoneKey->view(), stream, mr); + } + + auto local = tswtzLocalWallClock(packedCol, stream, mr); + auto added = addUnitToInstant( + local->view(), + valueCol, + literalValue_, + literalValueIsValid_, + unit_, + stream, + mr); + auto distinctKeys = tswtzDistinctZoneKeys(perRowZoneKey->view(), stream, mr); + auto utc = tswtzLocalToUtc( + added->view(), + perRowZoneKey->view(), + distinctKeys, + /*correctForward=*/true, + stream, + mr); + return tswtzPack(utc->view(), perRowZoneKey->view(), stream, mr); +} + } // namespace facebook::velox::cudf_velox::prestosql diff --git a/velox/experimental/cudf/expression/prestosql/DateAddFunction.h b/velox/experimental/cudf/expression/prestosql/DateAddFunction.h index f0dede70792..7047cd1cd2a 100644 --- a/velox/experimental/cudf/expression/prestosql/DateAddFunction.h +++ b/velox/experimental/cudf/expression/prestosql/DateAddFunction.h @@ -42,6 +42,7 @@ class DateAddFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override; @@ -77,4 +78,84 @@ class DateAddFunction : public CudfFunction { std::unique_ptr literalDate_; }; +/// date_add(unit, value, timestamp) -> TIMESTAMP. +/// Adds value units to a timestamp. unit is a constant string from the full +/// Presto set (millisecond, second, minute, hour, day, week, month, quarter, +/// year). Sub-day units add directly to the UTC instant (DST-agnostic); when +/// the session applies a timezone, day-and-above units convert to the session +/// wall clock, add, then convert back to UTC (a spring-forward gap throws, +/// matching addToTimestamp/toGMT). value (bigint) and timestamp may each be a +/// constant or a column, but at least one must be a column. +class DateAddTimestampFunction : public CudfFunction { + public: + /// Returns true if expr is date_add with 3 inputs, TIMESTAMP return type, + /// TIMESTAMP third argument, a constant parseable unit, and not both value + /// and timestamp constant. + static bool canEvaluate(const std::shared_ptr& expr); + + explicit DateAddTimestampFunction( + const std::shared_ptr& expr); + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override; + + private: + // Increment unit; any of the Presto date/time units including sub-day. + functions::DateTimeUnit unit_{}; + // True if the value (second) argument is a constant ConstantExpr. + bool valueIsLiteral_{}; + // True if the timestamp (third) argument is a constant ConstantExpr. + bool timestampIsLiteral_{}; + // True if literalValue_ is non-null (else the increment is a null scalar). + bool literalValueIsValid_{}; + // Constant value when valueIsLiteral_ is true. + int64_t literalValue_{}; + // Pre-built scalar of the constant timestamp input, when + // timestampIsLiteral_ is true. + std::unique_ptr literalTimestamp_; +}; + +/// date_add(unit, value, timestamp with time zone) -> TIMESTAMP WITH TIME ZONE. +/// Adds value units to a packed TSWTZ instant, preserving each row's zone key. +/// unit is a constant string from the full Presto set. Matching +/// addToTimestampWithTimezone: sub-day units (millisecond, second, minute, +/// hour) add directly to the UTC instant; day-and-above units add on each row's +/// local wall clock and convert back to UTC, resolving a spring-forward gap +/// forward instead of throwing. value (bigint) and the TSWTZ argument may each +/// be a constant or a column, but at least one must be a column. +class DateAddTimestampWithTimeZoneFunction : public CudfFunction { + public: + /// Returns true if expr is date_add with 3 inputs, TIMESTAMP WITH TIME ZONE + /// return type, TIMESTAMP WITH TIME ZONE third argument, a constant parseable + /// unit, and not both value and the TSWTZ argument constant. + static bool canEvaluate(const std::shared_ptr& expr); + + explicit DateAddTimestampWithTimeZoneFunction( + const std::shared_ptr& expr); + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override; + + private: + // Increment unit; any of the Presto date/time units including sub-day. + functions::DateTimeUnit unit_{}; + // True if the value (second) argument is a constant ConstantExpr. + bool valueIsLiteral_{}; + // True if the TSWTZ (third) argument is a constant ConstantExpr. + bool timestampIsLiteral_{}; + // True if literalValue_ is non-null (else the increment is a null scalar). + bool literalValueIsValid_{}; + // Constant value when valueIsLiteral_ is true. + int64_t literalValue_{}; + // Pre-built scalar of the constant TSWTZ input (packed int64), when + // timestampIsLiteral_ is true. + std::unique_ptr literalTimestamp_; +}; + } // namespace facebook::velox::cudf_velox::prestosql diff --git a/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.cpp b/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.cpp index 4cdd57fe7c9..2ce0f9a8cd4 100644 --- a/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.cpp +++ b/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.cpp @@ -75,6 +75,7 @@ DatePlusIntervalFunction::DatePlusIntervalFunction( ColumnOrView DatePlusIntervalFunction::eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const { auto dateCol = asView(inputColumns[0]); diff --git a/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.h b/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.h index eda012b8518..f79994e2078 100644 --- a/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.h +++ b/velox/experimental/cudf/expression/prestosql/DatePlusIntervalFunction.h @@ -31,6 +31,7 @@ class DatePlusIntervalFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override; diff --git a/velox/experimental/cudf/expression/prestosql/TimezoneFunctions.cpp b/velox/experimental/cudf/expression/prestosql/TimezoneFunctions.cpp new file mode 100644 index 00000000000..a82b397a7dd --- /dev/null +++ b/velox/experimental/cudf/expression/prestosql/TimezoneFunctions.cpp @@ -0,0 +1,1694 @@ +/* + * 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 "velox/experimental/cudf/expression/ExpressionEvaluator.h" +#include "velox/experimental/cudf/expression/TimestampWithTimeZoneColumn.h" +#include "velox/experimental/cudf/expression/TimezoneConversion.h" +#include "velox/experimental/cudf/expression/prestosql/TimezoneFunctions.h" + +#include "velox/common/base/CheckedArithmetic.h" +#include "velox/common/base/Exceptions.h" +#include "velox/expression/ConstantExpr.h" +#include "velox/expression/Expr.h" +#include "velox/expression/FunctionSignature.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneRegistration.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" +#include "velox/type/tz/TimeZoneMap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace facebook::velox::cudf_velox { +namespace { + +using velox::exec::ConstantExpr; + +constexpr cudf::type_id kInt64 = cudf::type_id::INT64; +constexpr cudf::type_id kBool8 = cudf::type_id::BOOL8; + +cudf::data_type int64Type() { + return cudf::data_type{kInt64}; +} + +// Reads a required constant string argument (e.g. a timezone name or format). +std::string constStringArg( + const std::shared_ptr& expr, + int32_t index) { + auto constant = + std::dynamic_pointer_cast(expr->inputs()[index]); + VELOX_CHECK_NOT_NULL( + constant, "Expected a constant argument at index {}", index); + return constant->value()->toString(0); +} + +// Reads a required constant integer argument (e.g. an hour/minute offset). +int64_t constIntArg( + const std::shared_ptr& expr, + int32_t index) { + auto constant = + std::dynamic_pointer_cast(expr->inputs()[index]); + VELOX_CHECK_NOT_NULL( + constant, "Expected a constant argument at index {}", index); + return std::stoll(constant->value()->toString(0)); +} + +// Reinterprets an 8-byte-wide column (timestamp/duration/int64) as another +// 8-byte type without copying. Used to move between the packed int64 +// representation and timestamp/duration columns. +cudf::column_view bitcastColumn( + const cudf::column_view& view, + cudf::type_id id) { + return cudf::column_view{ + cudf::data_type{id}, + view.size(), + view.head(), + view.null_mask(), + view.null_count(), + view.offset()}; +} + +cudf::numeric_scalar int64Scalar( + int64_t value, + rmm::cuda_stream_view stream) { + return cudf::numeric_scalar(value, true, stream); +} + +std::unique_ptr binaryOp( + const cudf::column_view& lhs, + const cudf::scalar& rhs, + cudf::binary_operator op, + cudf::data_type outType, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return cudf::binary_operation(lhs, rhs, op, outType, stream, mr); +} + +// Unpacks the UTC millis (arithmetic >> 12) from a packed TIMESTAMP WITH TIME +// ZONE column. +std::unique_ptr unpackMillis( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return binaryOp( + packed, + int64Scalar(kMillisShift, stream), + cudf::binary_operator::SHIFT_RIGHT, + int64Type(), + stream, + mr); +} + +// The per-row zone key of a packed column plus the distinct non-null keys +// present, computed once and shared by the numeric-offset and zone-name paths. +struct DistinctZones { + // packed & kTimezoneMask (INT64); nulls preserved (a null packed row yields a + // null key, which matches no real key in the per-zone selects below). + std::unique_ptr perRowKey; + // Distinct valid zone keys present in the column (the null key excluded). + std::vector keys; +}; + +// Extracts the per-row zone key, finds the distinct set on device +// (cudf::distinct), and copies the (small) distinct keys plus their validity to +// host so each zone's name/transition lookup runs once. One device->host sync. +DistinctZones distinctZones( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto perRowKey = tswtzZoneKey(packed, stream, mr); + auto keys = tswtzDistinctZoneKeys(perRowKey->view(), stream, mr); + return {std::move(perRowKey), std::move(keys)}; +} + +// Per-row UT offset in whole seconds (INT64) for a packed column that may mix +// zone keys. For each distinct key, computes utcOffsetSeconds over the whole +// column and selects the rows carrying that key. Null rows stay null (their +// null key matches no real key, so copy_if_else keeps the null default). +// O(#distinct zones) device passes. +std::unique_ptr perRowOffsetSeconds( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return tswtzOffsetSeconds(packed, stream, mr); +} + +// Per-row zone *name* (STRING) for a packed column that may mix zone keys, for +// the format_datetime 'ZZZ' zone-id token. Starts all-null and fills each +// distinct zone's rows with tz::getTimeZoneName(key) via a string-scalar +// copy_if_else. Null rows stay null. O(#distinct zones) device passes. +std::unique_ptr perRowZoneName( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto zones = distinctZones(packed, stream, mr); + + // Start all-null strings; fill each zone's rows with its name. An invalid + // string_scalar builds an all-null strings column; a null key's rows are + // never selected, so they keep that null (CPU propagates null). + auto result = cudf::make_column_from_scalar( + cudf::string_scalar("", false, stream), packed.size(), stream, mr); + for (const auto zoneKey : zones.keys) { + auto isThisZone = binaryOp( + zones.perRowKey->view(), + int64Scalar(zoneKey, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + // string_scalar-lhs / column-rhs overload: true -> the zone-name scalar, + // false or null-mask -> the accumulated result. + result = cudf::copy_if_else( + cudf::string_scalar(tz::getTimeZoneName(zoneKey), true, stream), + result->view(), + isThisZone->view(), + stream, + mr); + } + return result; +} + +// Renders a column of UT offsets (INT64 seconds) as a time-zone token. With +// includeColon the form is "+HH:MM" (Joda 'ZZ'); otherwise "+HHMM" (Joda 'Z'). +// When zeroOffsetText is set, rows with a zero offset render that text instead +// (e.g. "Z" for to_iso8601's ISO8601 output). +std::unique_ptr formatOffsetStrings( + const cudf::column_view& offsetSeconds, + bool includeColon, + const std::optional& zeroOffsetText, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto isNegative = binaryOp( + offsetSeconds, + int64Scalar(0, stream), + cudf::binary_operator::LESS, + cudf::data_type{kBool8}, + stream, + mr); + // abs(offset) = isNegative ? -offset : offset. + auto negated = cudf::binary_operation( + int64Scalar(0, stream), + offsetSeconds, + cudf::binary_operator::SUB, + int64Type(), + stream, + mr); + auto absolute = cudf::copy_if_else( + negated->view(), offsetSeconds, isNegative->view(), stream, mr); + auto hours = binaryOp( + absolute->view(), + int64Scalar(3'600, stream), + cudf::binary_operator::DIV, + int64Type(), + stream, + mr); + auto totalMinutes = binaryOp( + absolute->view(), + int64Scalar(60, stream), + cudf::binary_operator::DIV, + int64Type(), + stream, + mr); + auto minutes = binaryOp( + totalMinutes->view(), + int64Scalar(60, stream), + cudf::binary_operator::MOD, + int64Type(), + stream, + mr); + + auto hoursStr = cudf::strings::from_integers(hours->view(), stream, mr); + auto hoursPadded = cudf::strings::zfill( + cudf::strings_column_view(hoursStr->view()), 2, stream, mr); + auto minutesStr = cudf::strings::from_integers(minutes->view(), stream, mr); + auto minutesPadded = cudf::strings::zfill( + cudf::strings_column_view(minutesStr->view()), 2, stream, mr); + + auto sign = cudf::copy_if_else( + cudf::string_scalar("-", true, stream), + cudf::string_scalar("+", true, stream), + isNegative->view(), + stream, + mr); + + // "+/-" + "HH", then join with ":" before "MM". + auto signHour = cudf::strings::concatenate( + cudf::table_view{{sign->view(), hoursPadded->view()}}, + cudf::string_scalar("", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + auto offsetStr = cudf::strings::concatenate( + cudf::table_view{{signHour->view(), minutesPadded->view()}}, + cudf::string_scalar(includeColon ? ":" : "", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + if (!zeroOffsetText.has_value()) { + return offsetStr; + } + // Render the zero-offset rows as the supplied text (e.g. "Z"). + auto isZero = binaryOp( + offsetSeconds, + int64Scalar(0, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + return cudf::copy_if_else( + cudf::string_scalar(*zeroOffsetText, true, stream), + offsetStr->view(), + isZero->view(), + stream, + mr); +} + +// Computes the local wall-clock timestamp (TIMESTAMP_MILLISECONDS) and the UT +// offset (INT64 seconds) for a packed column that may mix zone keys, applying +// each row's own offset. +struct LocalAndOffset { + std::unique_ptr localMillis; + std::unique_ptr offsetSeconds; +}; + +LocalAndOffset localAndOffset( + const cudf::column_view& packed, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + return { + tswtzLocalWallClock(packed, stream, mr), + tswtzOffsetSeconds(packed, stream, mr)}; +} + +// Classifies the trailing Joda time-zone token so the caller can render it: an +// offset (with or without a colon), the zone id, or the zone name. Matches CPU +// DateTimeFormatter: single 'Z' has no colon, 'ZZ' has a colon, 'ZZZ' or more +// is the zone id, and lowercase 'z' is the zone abbreviation/name. +enum class TrailingZone { + kNone, + kOffsetNoColon, + kOffsetColon, + kZoneId, + kZoneName, +}; + +// Translates the subset of Joda DateTimeFormat pattern letters used by the +// covered functions to cuDF strftime/strptime specifiers. A trailing time-zone +// token (Z/z) is classified via trailing and rendered separately. +std::string jodaToStrftime(const std::string& joda, TrailingZone& trailing) { + trailing = TrailingZone::kNone; + std::string out; + size_t i = 0; + while (i < joda.size()) { + const char c = joda[i]; + if (c == '\'') { + ++i; + while (i < joda.size() && joda[i] != '\'') { + out += joda[i++]; + } + if (i < joda.size()) { + ++i; + } + continue; + } + if (std::isalpha(static_cast(c))) { + size_t j = i; + while (j < joda.size() && joda[j] == c) { + ++j; + } + const size_t runLength = j - i; + switch (c) { + case 'y': + case 'Y': + out += runLength >= 3 ? "%Y" : "%y"; + break; + case 'M': + out += runLength >= 4 ? "%B" : (runLength == 3 ? "%b" : "%m"); + break; + case 'd': + out += "%d"; + break; + case 'H': + out += "%H"; + break; + case 'h': + out += "%I"; + break; + case 'm': + out += "%M"; + break; + case 's': + out += "%S"; + break; + case 'S': + // The 'S' run length is the fractional-second digit count ('S' -> 1 + // digit, 'SSSSSS' -> 6), matching CPU's formatFractionOfSecond. cuDF + // renders "%f" for n in 1..9; nothing finer than nanoseconds is + // representable. + if (runLength > 9) { + VELOX_NYI( + "format_datetime supports at most 9 fractional-second digits " + "on GPU, got {}", + runLength); + } + out += "%" + std::to_string(runLength) + "f"; + break; + case 'a': + out += "%p"; + break; + case 'E': + out += runLength >= 4 ? "%A" : "%a"; + break; + case 'Z': + case 'z': + VELOX_CHECK_EQ( + j, + joda.size(), + "cuDF datetime format supports a time zone token only at the end"); + if (c == 'z') { + trailing = TrailingZone::kZoneName; + } else if (runLength == 1) { + trailing = TrailingZone::kOffsetNoColon; + } else if (runLength == 2) { + trailing = TrailingZone::kOffsetColon; + } else { + trailing = TrailingZone::kZoneId; + } + break; + default: + VELOX_NYI( + "Unsupported datetime format letter on GPU: {}", + std::string(1, c)); + } + i = j; + continue; + } + out += c; + ++i; + } + return out; +} + +// to_unixtime(timestamp with time zone) -> double. +class ToUnixtimeFunction : public CudfFunction { + public: + explicit ToUnixtimeFunction(const std::shared_ptr& expr) { + VELOX_CHECK_EQ( + expr->inputs().size(), 1, "to_unixtime expects exactly 1 input"); + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto packed = asView(inputColumns[0]); + auto millis = unpackMillis(packed, stream, mr); + auto millisDouble = cudf::cast( + millis->view(), cudf::data_type{cudf::type_id::FLOAT64}, stream, mr); + auto thousand = cudf::numeric_scalar(1000.0, true, stream); + return cudf::binary_operation( + millisDouble->view(), + thousand, + cudf::binary_operator::DIV, + cudf::data_type{cudf::type_id::FLOAT64}, + stream, + mr); + } +}; + +// at_timezone(timestamp with time zone, varchar) -> timestamp with time zone. +class AtTimezoneFunction : public CudfFunction { + public: + explicit AtTimezoneFunction(const std::shared_ptr& expr) { + VELOX_CHECK_EQ( + expr->inputs().size(), 2, "at_timezone expects exactly 2 inputs"); + targetZoneId_ = tz::getTimeZoneID(constStringArg(expr, 1)); + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto packed = asView(inputColumns[0]); + // Keep the UTC millis bits, replace the low 12 zone bits with the new key. + auto cleared = binaryOp( + packed, + int64Scalar(~static_cast(kTimezoneMask), stream), + cudf::binary_operator::BITWISE_AND, + int64Type(), + stream, + mr); + return binaryOp( + cleared->view(), + int64Scalar(targetZoneId_ & kTimezoneMask, stream), + cudf::binary_operator::BITWISE_OR, + int64Type(), + stream, + mr); + } + + private: + int16_t targetZoneId_; +}; + +// timezone_hour / timezone_minute (timestamp with time zone) -> bigint. +class TimezoneFieldFunction : public CudfFunction { + public: + TimezoneFieldFunction( + const std::shared_ptr& expr, + bool minuteField) + : minuteField_(minuteField) { + VELOX_CHECK_EQ( + expr->inputs().size(), + 1, + "timezone_hour/timezone_minute expects exactly 1 input"); + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto packed = asView(inputColumns[0]); + auto offsetSeconds = perRowOffsetSeconds(packed, stream, mr); + if (minuteField_) { + auto perMinute = binaryOp( + offsetSeconds->view(), + int64Scalar(60, stream), + cudf::binary_operator::DIV, + int64Type(), + stream, + mr); + return binaryOp( + perMinute->view(), + int64Scalar(60, stream), + cudf::binary_operator::MOD, + int64Type(), + stream, + mr); + } + return binaryOp( + offsetSeconds->view(), + int64Scalar(3'600, stream), + cudf::binary_operator::DIV, + int64Type(), + stream, + mr); + } + + private: + bool minuteField_; +}; + +// to_iso8601(timestamp with time zone) -> varchar. +class ToIso8601Function : public CudfFunction { + public: + explicit ToIso8601Function(const std::shared_ptr& expr) { + VELOX_CHECK_EQ( + expr->inputs().size(), 1, "to_iso8601 expects exactly 1 input"); + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto packed = asView(inputColumns[0]); + auto parts = localAndOffset(packed, stream, mr); + auto dateStr = cudf::strings::from_timestamps( + parts.localMillis->view(), + "%Y-%m-%dT%H:%M:%S.%3f", + cudf::strings_column_view{}, + stream, + mr); + // ISO8601 uses "+HH:MM" but renders a zero offset as "Z". + auto offsetStr = formatOffsetStrings( + parts.offsetSeconds->view(), + /*includeColon=*/true, + std::string("Z"), + stream, + mr); + return cudf::strings::concatenate( + cudf::table_view{{dateStr->view(), offsetStr->view()}}, + cudf::string_scalar("", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + } +}; + +// format_datetime(timestamp with time zone, varchar) -> varchar. +class FormatDatetimeFunction : public CudfFunction { + public: + explicit FormatDatetimeFunction( + const std::shared_ptr& expr) { + VELOX_CHECK_EQ( + expr->inputs().size(), 2, "format_datetime expects exactly 2 inputs"); + strftime_ = jodaToStrftime(constStringArg(expr, 1), trailing_); + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto packed = asView(inputColumns[0]); + auto parts = localAndOffset(packed, stream, mr); + auto dateStr = cudf::strings::from_timestamps( + parts.localMillis->view(), + strftime_, + cudf::strings_column_view{}, + stream, + mr); + if (trailing_ == TrailingZone::kNone) { + return dateStr; + } + + std::unique_ptr zoneStr; + switch (trailing_) { + case TrailingZone::kOffsetNoColon: + zoneStr = formatOffsetStrings( + parts.offsetSeconds->view(), + /*includeColon=*/false, + std::nullopt, + stream, + mr); + break; + case TrailingZone::kOffsetColon: + zoneStr = formatOffsetStrings( + parts.offsetSeconds->view(), + /*includeColon=*/true, + std::nullopt, + stream, + mr); + break; + case TrailingZone::kZoneId: { + // Each row renders its own zone name; the column may mix zones. + zoneStr = perRowZoneName(packed, stream, mr); + break; + } + case TrailingZone::kZoneName: + // The zone abbreviation/name ('z') is DST- and instant-dependent; cuDF + // cannot render it on device. + VELOX_NYI( + "format_datetime zone-name token 'z' is not supported on GPU"); + case TrailingZone::kNone: + VELOX_UNREACHABLE(); + } + return cudf::strings::concatenate( + cudf::table_view{{dateStr->view(), zoneStr->view()}}, + cudf::string_scalar("", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + } + + private: + std::string strftime_; + TrailingZone trailing_{TrailingZone::kNone}; +}; + +// Mirrors the CPU pack() range check: throws if any non-null millis value falls +// outside [kMinMillisUtc, kMaxMillisUtc]. Without this, from_unixtime would +// shift an out-of-range instant into the zone-key bits and silently corrupt the +// packed value instead of rejecting it as CPU does. +void checkMillisInRange( + const cudf::column_view& millis, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (millis.size() == 0 || millis.null_count() == millis.size()) { + return; + } + auto minScalar = cudf::reduce( + millis, + *cudf::make_min_aggregation(), + int64Type(), + stream, + mr); + auto maxScalar = cudf::reduce( + millis, + *cudf::make_max_aggregation(), + int64Type(), + stream, + mr); + const auto lo = static_cast*>(minScalar.get()) + ->value(stream); + const auto hi = static_cast*>(maxScalar.get()) + ->value(stream); + VELOX_USER_CHECK( + lo >= kMinMillisUtc && hi <= kMaxMillisUtc, + "TimestampWithTimeZone overflow: [{}, {}] ms", + lo, + hi); +} + +// Mirrors the CPU offset bound: from_iso8601_timestamp normalizes its parsed +// offset through tz::getTimeZoneID, which rejects magnitudes beyond +/-14h (840 +// minutes). magnitudeMinutes holds the absolute offset minutes (always +// non-negative), so an upper bound covers both signs. Without this, an offset +// like "+99:00" (5940 minutes) maps to a zone key that overflows the 12-bit +// zone field and corrupts the packed millis. +void checkOffsetMagnitudeInRange( + const cudf::column_view& magnitudeMinutes, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (magnitudeMinutes.size() == 0 || + magnitudeMinutes.null_count() == magnitudeMinutes.size()) { + return; + } + auto maxScalar = cudf::reduce( + magnitudeMinutes, + *cudf::make_max_aggregation(), + int64Type(), + stream, + mr); + const auto hi = static_cast*>(maxScalar.get()) + ->value(stream); + VELOX_USER_CHECK_LE( + hi, 840, "Invalid timezone offset in from_iso8601_timestamp (minutes)"); +} + +// True if any row of the boolean mask is set. An empty or all-null mask -> +// false (so a batch of only SQL-NULL rows raises no error). +bool anyRowTrue( + const cudf::column_view& mask, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + if (mask.size() == 0) { + return false; + } + auto reduced = cudf::reduce( + mask, + *cudf::make_any_aggregation(), + cudf::data_type{kBool8}, + stream, + mr); + auto& scalar = static_cast&>(*reduced); + return scalar.is_valid(stream) && scalar.value(stream); +} + +// Maps captured offset groups (sign character, hours digits, minutes digits) to +// a signed offset in whole minutes as an INT64 column. A null in the hours or +// minutes group is treated as absent (contributes 0), and a null sign defaults +// to '+', so a missing offset (Z or no suffix) yields 0 (GMT). The sign is read +// from the sign character so "-00:30" stays negative. Rejects magnitudes beyond +// +/-840 minutes via checkOffsetMagnitudeInRange, matching CPU's +// tz::getTimeZoneID bound. +std::unique_ptr signedOffsetMinutes( + const cudf::column_view& signChar, + const cudf::column_view& hoursDigits, + const cudf::column_view& minutesDigits, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto offsetHours = cudf::replace_nulls( + cudf::strings::to_integers( + cudf::strings_column_view(hoursDigits), int64Type(), stream, mr) + ->view(), + int64Scalar(0, stream), + stream, + mr); + auto offsetMins = cudf::replace_nulls( + cudf::strings::to_integers( + cudf::strings_column_view(minutesDigits), int64Type(), stream, mr) + ->view(), + int64Scalar(0, stream), + stream, + mr); + auto signStr = cudf::replace_nulls( + signChar, cudf::string_scalar("+", true, stream), stream, mr); + auto isNegativeSign = cudf::strings::starts_with( + cudf::strings_column_view(signStr->view()), + cudf::string_scalar("-", true, stream), + stream, + mr); + auto hourMinutes = binaryOp( + offsetHours->view(), + int64Scalar(60, stream), + cudf::binary_operator::MUL, + int64Type(), + stream, + mr); + auto magnitude = cudf::binary_operation( + hourMinutes->view(), + offsetMins->view(), + cudf::binary_operator::ADD, + int64Type(), + stream, + mr); + // Reject offsets beyond +/-14h before they pack into a zone key that + // overflows the 12-bit zone field, matching CPU's tz::getTimeZoneID bound. + checkOffsetMagnitudeInRange(magnitude->view(), stream, mr); + auto negativeMagnitude = cudf::binary_operation( + int64Scalar(0, stream), + magnitude->view(), + cudf::binary_operator::SUB, + int64Type(), + stream, + mr); + return cudf::copy_if_else( + negativeMagnitude->view(), + magnitude->view(), + isNegativeSign->view(), + stream, + mr); +} + +// Maps a signed offset-minutes INT64 column to packed fixed-offset zone keys, +// mirroring Velox's TimeZoneMap ordering: 0 -> 0 (GMT); <0 -> offset+841; +// >0 -> offset+840. +std::unique_ptr zoneKeyFromOffsetMinutes( + const cudf::column_view& offsetMinutes, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto idPositive = binaryOp( + offsetMinutes, + int64Scalar(840, stream), + cudf::binary_operator::ADD, + int64Type(), + stream, + mr); + auto idNegative = binaryOp( + offsetMinutes, + int64Scalar(841, stream), + cudf::binary_operator::ADD, + int64Type(), + stream, + mr); + auto isNegativeOffset = binaryOp( + offsetMinutes, + int64Scalar(0, stream), + cudf::binary_operator::LESS, + cudf::data_type{kBool8}, + stream, + mr); + auto idNonZero = cudf::copy_if_else( + idNegative->view(), + idPositive->view(), + isNegativeOffset->view(), + stream, + mr); + auto isZeroOffset = binaryOp( + offsetMinutes, + int64Scalar(0, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + return cudf::copy_if_else( + int64Scalar(0, stream), + idNonZero->view(), + isZeroOffset->view(), + stream, + mr); +} + +// Selects the millisecond rounding for from_unixtime, which differs between the +// two CPU overloads. from_unixtime(double, varchar) rounds the whole value with +// llround(x*1000); from_unixtime(double, hours, minutes) floors the seconds and +// rounds the fractional millisecond separately. The two agree except on +// negative-fractional input, where they can differ by 1 ms (e.g. -0.0005 s -> +// -1 ms for kWhole, 0 ms for kFloorThenFraction). +enum class FromUnixtimeRounding { + kWhole, + kFloorThenFraction, +}; + +// from_unixtime(double, ...) -> timestamp with time zone. The zone id is fixed +// at construction (from a zone name or an hour/minute offset). +class FromUnixtimeWithZoneFunction : public CudfFunction { + public: + FromUnixtimeWithZoneFunction(int16_t zoneId, FromUnixtimeRounding rounding) + : zoneId_(zoneId), rounding_(rounding) {} + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto seconds = asView(inputColumns[0]); + const auto doubleType = cudf::data_type{cudf::type_id::FLOAT64}; + + // Emulate std::llround on a FLOAT64 column: add +/-0.5 then truncate toward + // zero via the FLOAT64->INT64 cast. + auto llroundEmu = [&](const cudf::column_view& value) { + auto isNegative = cudf::binary_operation( + value, + cudf::numeric_scalar(0.0, true, stream), + cudf::binary_operator::LESS, + cudf::data_type{kBool8}, + stream, + mr); + auto half = cudf::copy_if_else( + cudf::numeric_scalar(-0.5, true, stream), + cudf::numeric_scalar(0.5, true, stream), + isNegative->view(), + stream, + mr); + auto adjusted = cudf::binary_operation( + value, + half->view(), + cudf::binary_operator::ADD, + doubleType, + stream, + mr); + return cudf::cast(adjusted->view(), int64Type(), stream, mr); + }; + + std::unique_ptr millis; + if (rounding_ == FromUnixtimeRounding::kWhole) { + auto millisDouble = cudf::binary_operation( + seconds, + cudf::numeric_scalar(1000.0, true, stream), + cudf::binary_operator::MUL, + doubleType, + stream, + mr); + millis = llroundEmu(millisDouble->view()); + } else { + // floor(x) whole seconds, plus the fractional second rounded on its own + // (matching CPU's no-zone fromUnixtime). The fraction is in [0, 1), so + // its round is the non-negative case; a fraction that rounds up to 1000 + // ms carries naturally through the addition. + auto secondsFloor = cudf::unary_operation( + seconds, cudf::unary_operator::FLOOR, stream, mr); + auto fraction = cudf::binary_operation( + seconds, + secondsFloor->view(), + cudf::binary_operator::SUB, + doubleType, + stream, + mr); + auto fractionMillisDouble = cudf::binary_operation( + fraction->view(), + cudf::numeric_scalar(1000.0, true, stream), + cudf::binary_operator::MUL, + doubleType, + stream, + mr); + auto fractionMillis = llroundEmu(fractionMillisDouble->view()); + auto secondsInt = + cudf::cast(secondsFloor->view(), int64Type(), stream, mr); + auto secondsMillis = binaryOp( + secondsInt->view(), + int64Scalar(1000, stream), + cudf::binary_operator::MUL, + int64Type(), + stream, + mr); + millis = cudf::binary_operation( + secondsMillis->view(), + fractionMillis->view(), + cudf::binary_operator::ADD, + int64Type(), + stream, + mr); + } + + // Match CPU's non-finite handling, which a FLOAT64->INT64 cast does not + // give on its own: NaN maps to pack(0), and +/-Inf saturates out of range + // so pack (here checkMillisInRange) rejects it. A null input stays null + // because a null comparison yields a null mask element, which copy_if_else + // treats as false and so keeps the (null) computed millis. + const auto infinity = std::numeric_limits::infinity(); + auto isNan = cudf::binary_operation( + seconds, + seconds, + cudf::binary_operator::NOT_EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + auto isPositiveInf = cudf::binary_operation( + seconds, + cudf::numeric_scalar(infinity, true, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + auto isNegativeInf = cudf::binary_operation( + seconds, + cudf::numeric_scalar(-infinity, true, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + auto isInf = cudf::binary_operation( + isPositiveInf->view(), + isNegativeInf->view(), + cudf::binary_operator::LOGICAL_OR, + cudf::data_type{kBool8}, + stream, + mr); + millis = cudf::copy_if_else( + int64Scalar(0, stream), millis->view(), isNan->view(), stream, mr); + millis = cudf::copy_if_else( + int64Scalar(kMaxMillisUtc + 1, stream), + millis->view(), + isInf->view(), + stream, + mr); + + checkMillisInRange(millis->view(), stream, mr); + auto shifted = binaryOp( + millis->view(), + int64Scalar(kMillisShift, stream), + cudf::binary_operator::SHIFT_LEFT, + int64Type(), + stream, + mr); + return binaryOp( + shifted->view(), + int64Scalar(zoneId_ & kTimezoneMask, stream), + cudf::binary_operator::BITWISE_OR, + int64Type(), + stream, + mr); + } + + private: + int16_t zoneId_; + FromUnixtimeRounding rounding_; +}; + +// now() / current_timestamp -> timestamp with time zone. Emits a constant +// column packing the session start time with the session zone, matching CPU's +// CurrentTimestampFunction (the value is not compared against a live CPU now(), +// which is non-deterministic). Rejects like CPU when the session zone is +// unusable: getTimeZoneFromConfig returns null when +// adjust_timestamp_to_session_timezone is off or the session timezone is empty, +// and CPU then throws "Timezone cannot be null". +class NowFunction : public CudfFunction { + public: + ColumnOrView eval( + [[maybe_unused]] std::vector& inputColumns, + cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + VELOX_USER_CHECK( + context_.adjustTimestampToTimezone && !context_.sessionTimezone.empty(), + "Timezone cannot be null"); + const auto zoneId = tz::getTimeZoneID(context_.sessionTimezone); + const int64_t packed = pack(context_.sessionStartTimeMs, zoneId); + auto scalar = int64Scalar(packed, stream); + return cudf::make_column_from_scalar(scalar, numRows, stream, mr); + } +}; + +// parse_datetime(varchar, varchar) -> timestamp with time zone. +class ParseDatetimeFunction : public CudfFunction { + public: + explicit ParseDatetimeFunction( + const std::shared_ptr& expr) { + VELOX_CHECK_EQ( + expr->inputs().size(), 2, "parse_datetime expects exactly 2 inputs"); + TrailingZone trailing = TrailingZone::kNone; + strptime_ = jodaToStrftime(constStringArg(expr, 1), trailing); + if (trailing == TrailingZone::kOffsetNoColon || + trailing == TrailingZone::kOffsetColon) { + // to_timestamps folds the %z offset into the UTC instant. The parsed + // offset is recovered per-row in eval so the packed zone key reflects it + // instead of GMT. + strptime_ += "%z"; + hasOffset_ = true; + // Trailing signed offset with an optional colon, matching both "-09:00" + // and "-0900". Groups: 0 sign, 1 hours, 2 minutes. + offsetProgram_ = + cudf::strings::regex_program::create("([+-])([0-9]{2}):?([0-9]{2})$"); + } else if (trailing != TrailingZone::kNone) { + VELOX_NYI("parse_datetime zone-name token is not supported on GPU"); + } + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto input = asView(inputColumns[0]); + // cuDF parses the wall clock as UTC. With no embedded zone the result is + // interpreted in the session timezone (GMT when unset), so the parsed + // value equals the UTC instant in the GMT case the tests exercise. + if (!context_.sessionTimezone.empty()) { + VELOX_NYI( + "parse_datetime on GPU with a non-UTC session timezone is not yet " + "supported"); + } + auto parsed = cudf::strings::to_timestamps( + cudf::strings_column_view(input), + cudf::data_type{cudf::type_id::TIMESTAMP_MILLISECONDS}, + strptime_, + stream, + mr); + auto millis = bitcastColumn(parsed->view(), kInt64); + auto shifted = binaryOp( + millis, + int64Scalar(kMillisShift, stream), + cudf::binary_operator::SHIFT_LEFT, + int64Type(), + stream, + mr); + if (!hasOffset_) { + // pack(millis, GMT) == millis << 12. + return shifted; + } + // Recover the per-row offset that to_timestamps folded into the instant and + // pack the matching fixed-offset zone key, so timezone_hour/to_iso8601 + // reflect the parsed offset instead of GMT. + auto groups = cudf::strings::extract( + cudf::strings_column_view(input), *offsetProgram_, stream, mr); + auto g = groups->view(); + auto offsetMinutes = + signedOffsetMinutes(g.column(0), g.column(1), g.column(2), stream, mr); + auto zoneId = zoneKeyFromOffsetMinutes(offsetMinutes->view(), stream, mr); + return cudf::binary_operation( + shifted->view(), + zoneId->view(), + cudf::binary_operator::BITWISE_OR, + int64Type(), + stream, + mr); + } + + private: + std::string strptime_; + // True when the Joda format carries a numeric offset token (%z appended); + // gates the per-row offset recovery in eval. + bool hasOffset_{false}; + // Compiled trailing-offset extraction program, built once when hasOffset_. + std::unique_ptr offsetProgram_; +}; + +// from_iso8601_timestamp(varchar) -> timestamp with time zone. +class FromIso8601Function : public CudfFunction { + public: + explicit FromIso8601Function(const std::shared_ptr& expr) { + VELOX_CHECK_EQ( + expr->inputs().size(), + 1, + "from_iso8601_timestamp expects exactly 1 input"); + // Permissive ISO8601 (date-anchored). The year is required; month, day, the + // time fields, the fractional seconds and the zone suffix are all optional, + // and missing components default to the start of the period (matching CPU). + // A 'T' may appear with no time after it ("2021-01-01T", "2021T+14:00"); + // the date/time separator is a literal 'T' only, since CPU rejects a space. + // Time-only inputs ("T11:38") carry no date; eval prefixes the epoch date + // "1970-01-01" to them before this program runs, so the single + // date-anchored program still covers them. The whole zone suffix is + // captured (group 7) to tell an absent suffix from an explicit "Z"; the + // sign is captured on its own (group 8) so a sub-hour offset like "-00:30" + // keeps it. Groups: 0 year, 1 month, 2 day, 3 hour, 4 minute, 5 second, 6 + // fraction, 7 zone suffix, 8 sign, 9 offset hours, 10 offset minutes. + // Batch-independent, so build once. + isoProgram_ = cudf::strings::regex_program::create( + "^([0-9]{4})(?:-([0-9]{2}))?(?:-([0-9]{2}))?" + "(?:T([0-9]{2})?(?::([0-9]{2}))?(?::([0-9]{2}))?)?" + "(?:[.,]([0-9]+))?" + "(Z|([+-])([0-9]{2})(?::?([0-9]{2}))?)?$"); + // Identifies a leading time-only form ("Thh...") so eval can prefix the + // epoch date "1970-01-01" and reuse the date-anchored program. A bare "T" + // (no digits) does not match, so it stays unprefixed and is later rejected + // as malformed, like CPU. + timeOnlyProgram_ = cudf::strings::regex_program::create("^T[0-9]{2}"); + // Matches an otherwise-valid ISO8601 string whose year is signed or has 5+ + // digits -- the CPU-valid extreme years cudf::strings::to_timestamps (int16 + // %Y) cannot represent. Same tail as isoProgram_ so only the year token + // differs; used only as a match test (captures are ignored). + extremeProgram_ = cudf::strings::regex_program::create( + "^(?:[+-][0-9]{4,}|[0-9]{5,})(?:-([0-9]{2}))?(?:-([0-9]{2}))?" + "(?:T([0-9]{2})?(?::([0-9]{2}))?(?::([0-9]{2}))?)?" + "(?:[.,]([0-9]+))?" + "(Z|([+-])([0-9]{2})(?::?([0-9]{2}))?)?$"); + } + + ColumnOrView eval( + std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const override { + auto input = asView(inputColumns[0]); + // Time-only inputs carry no date; CPU defaults them to 1970-01-01. Prefix + // that date to leading-'T' rows so the single date-anchored program (built + // in the ctor) handles them; every other row is passed through unchanged. + auto isTimeOnly = cudf::replace_nulls( + cudf::strings::matches_re( + cudf::strings_column_view(input), *timeOnlyProgram_, stream, mr) + ->view(), + cudf::numeric_scalar(false, true, stream), + stream, + mr); + auto epochDate = cudf::make_column_from_scalar( + cudf::string_scalar("1970-01-01", true, stream), + input.size(), + stream, + mr); + auto prefixed = cudf::strings::concatenate( + cudf::table_view{{epochDate->view(), input}}, + cudf::string_scalar("", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + auto work = cudf::copy_if_else( + prefixed->view(), input, isTimeOnly->view(), stream, mr); + auto workView = cudf::strings_column_view(work->view()); + // Extract the ISO8601 fields with the program built in the constructor; see + // there for the field layout, and the group-column map just below. + auto groups = cudf::strings::extract(workView, *isoProgram_, stream, mr); + auto g = groups->view(); + // Columns: 0 year, 1 month, 2 day, 3 hour, 4 minute, 5 second, 6 fraction, + // 7 zone suffix, 8 sign, 9 offset hours, 10 offset minutes. + + auto orDefault = [&](int index, const char* value) { + return cudf::replace_nulls( + g.column(index), + cudf::string_scalar(value, true, stream), + stream, + mr); + }; + // cudf's extract yields an empty string (not a null) for an optional group + // that did not participate in an otherwise-matching row, so replace_nulls + // alone leaves an absent month or day empty. to_timestamps then reads the + // empty numeric field as 0 and underflows, e.g. "2021" (no month/day) + // parses as 2020-11-30 instead of 2021-01-01. Month and day must default to + // "01", so replace an empty (or null) capture explicitly. The time fields + // default to 0, which an empty string already yields, so they keep + // replace_nulls. + auto orFirstOfPeriod = [&](int index) { + auto filled = cudf::replace_nulls( + g.column(index), cudf::string_scalar("01", true, stream), stream, mr); + auto length = cudf::strings::count_characters( + cudf::strings_column_view(filled->view()), stream, mr); + auto isEmpty = cudf::binary_operation( + length->view(), + cudf::numeric_scalar(0, true, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + return cudf::copy_if_else( + cudf::string_scalar("01", true, stream), + filled->view(), + isEmpty->view(), + stream, + mr); + }; + auto month = orFirstOfPeriod(1); + auto day = orFirstOfPeriod(2); + auto hour = orDefault(3, "00"); + auto minute = orDefault(4, "00"); + auto second = orDefault(5, "00"); + + // Build "YYYY-MM-DD" first; the throw block below and the canonical + // timestamp string both reuse it. A non-matching row leaves the year null, + // so separator_on_nulls yields a null that parses to null. + auto ymd = cudf::strings::concatenate( + cudf::table_view{{g.column(0), month->view(), day->view()}}, + cudf::string_scalar("-", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + + // Match CPU exactly for every non-null row: parse it, or throw. Genuine + // SQL-NULL rows are excluded via is_valid, so they keep propagating as + // NULL. A non-null row falls into one of three buckets: + // - matches the in-range program (isoProgram_) and names a real calendar + // date -> parsed normally below; + // - matches neither program, or matches isoProgram_ but names a + // nonexistent date (month/day out of range) -> malformed, exactly as + // CPU's fromTimestampWithTimezoneString / isValidDate -> + // VELOX_USER_FAIL; + // - matches only the extreme-year program -> CPU-valid but beyond what + // to_timestamps (int16 %Y) can represent -> VELOX_NYI. + // Malformed is checked first so a batch mixing malformed + extreme rows + // reports the parse error, as CPU would. + { + const auto falseScalar = cudf::numeric_scalar(false, true, stream); + const auto trueScalar = cudf::numeric_scalar(true, true, stream); + auto nonNull = cudf::is_valid(input, stream, mr); + auto tier1 = cudf::replace_nulls( + cudf::strings::matches_re(workView, *isoProgram_, stream, mr)->view(), + falseScalar, + stream, + mr); + auto extreme = cudf::replace_nulls( + cudf::strings::matches_re(workView, *extremeProgram_, stream, mr) + ->view(), + falseScalar, + stream, + mr); + auto known = cudf::binary_operation( + tier1->view(), + extreme->view(), + cudf::binary_operator::LOGICAL_OR, + cudf::data_type{kBool8}, + stream, + mr); + auto unknown = cudf::unary_operation( + known->view(), cudf::unary_operator::NOT, stream, mr); + auto malformedShape = cudf::binary_operation( + nonNull->view(), + unknown->view(), + cudf::binary_operator::LOGICAL_AND, + cudf::data_type{kBool8}, + stream, + mr); + + // cudf::strings::to_timestamps normalizes an out-of-range month or day + // (month 13 -> next year, day 30 in Feb -> March) instead of failing, so + // "2021-13-45" matches isoProgram_ yet is not a real date. Parse the + // date, read month and day back, and require they equal the parsed input. + // Any normalization moves the value into a different month (a day + // underflow or overflow always crosses a month boundary), so comparing + // month and day catches every invalid combination, including a non-leap + // Feb 29. The 4-digit year is regex-bounded to [0000, 9999], so it always + // round-trips. + const auto int16Type = cudf::data_type{cudf::type_id::INT16}; + auto dateTs = cudf::strings::to_timestamps( + cudf::strings_column_view(ymd->view()), + cudf::data_type{cudf::type_id::TIMESTAMP_MILLISECONDS}, + "%Y-%m-%d", + stream, + mr); + auto backMonth = cudf::datetime::extract_datetime_component( + dateTs->view(), + cudf::datetime::datetime_component::MONTH, + stream, + mr); + auto backDay = cudf::datetime::extract_datetime_component( + dateTs->view(), cudf::datetime::datetime_component::DAY, stream, mr); + auto inMonth = cudf::strings::to_integers( + cudf::strings_column_view(month->view()), int16Type, stream, mr); + auto inDay = cudf::strings::to_integers( + cudf::strings_column_view(day->view()), int16Type, stream, mr); + auto monthOk = cudf::binary_operation( + backMonth->view(), + inMonth->view(), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + auto dayOk = cudf::binary_operation( + backDay->view(), + inDay->view(), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + auto dateOk = cudf::binary_operation( + monthOk->view(), + dayOk->view(), + cudf::binary_operator::LOGICAL_AND, + cudf::data_type{kBool8}, + stream, + mr); + // A non-tier1 row has a null date, hence a null dateOk; treat null as ok + // so only tier1 rows can be flagged here (also AND'd with tier1 below). + auto dateOkFilled = + cudf::replace_nulls(dateOk->view(), trueScalar, stream, mr); + auto dateInvalid = cudf::unary_operation( + dateOkFilled->view(), cudf::unary_operator::NOT, stream, mr); + auto tier1NonNull = cudf::binary_operation( + nonNull->view(), + tier1->view(), + cudf::binary_operator::LOGICAL_AND, + cudf::data_type{kBool8}, + stream, + mr); + auto calendarBad = cudf::binary_operation( + tier1NonNull->view(), + dateInvalid->view(), + cudf::binary_operator::LOGICAL_AND, + cudf::data_type{kBool8}, + stream, + mr); + auto malformed = cudf::binary_operation( + malformedShape->view(), + calendarBad->view(), + cudf::binary_operator::LOGICAL_OR, + cudf::data_type{kBool8}, + stream, + mr); + VELOX_USER_CHECK( + !anyRowTrue(malformed->view(), stream, mr), + "Unable to parse timestamp value in from_iso8601_timestamp"); + auto extremeNonNull = cudf::binary_operation( + nonNull->view(), + extreme->view(), + cudf::binary_operator::LOGICAL_AND, + cudf::data_type{kBool8}, + stream, + mr); + if (anyRowTrue(extremeNonNull->view(), stream, mr)) { + VELOX_NYI( + "from_iso8601_timestamp does not support years outside [0000, 9999] on GPU"); + } + } + + auto hms = cudf::strings::concatenate( + cudf::table_view{{hour->view(), minute->view(), second->view()}}, + cudf::string_scalar(":", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + auto canonical = cudf::strings::concatenate( + cudf::table_view{{ymd->view(), hms->view()}}, + cudf::string_scalar("T", true, stream), + cudf::string_scalar("", false, stream), + cudf::strings::separator_on_nulls::YES, + stream, + mr); + auto wallTs = cudf::strings::to_timestamps( + cudf::strings_column_view(canonical->view()), + cudf::data_type{cudf::type_id::TIMESTAMP_MILLISECONDS}, + "%Y-%m-%dT%H:%M:%S", + stream, + mr); + auto wallMillisBase = bitcastColumn(wallTs->view(), kInt64); + + // Fractional seconds -> milliseconds: first 3 digits, right-padded to 3 + // (".1" -> 100, ".12" -> 120, ".123456" -> 123). Missing -> 0. + auto frac3 = cudf::strings::slice_strings( + cudf::strings_column_view(g.column(6)), + cudf::numeric_scalar(0, true, stream), + cudf::numeric_scalar(3, true, stream), + cudf::numeric_scalar(1, true, stream), + stream, + mr); + auto fracPadded = cudf::strings::pad( + cudf::strings_column_view(frac3->view()), + 3, + cudf::strings::side_type::RIGHT, + "0", + stream, + mr); + auto fracInts = cudf::strings::to_integers( + cudf::strings_column_view(fracPadded->view()), int64Type(), stream, mr); + auto fracMillis = cudf::replace_nulls( + fracInts->view(), int64Scalar(0, stream), stream, mr); + auto wallMillis = cudf::binary_operation( + wallMillisBase, + fracMillis->view(), + cudf::binary_operator::ADD, + int64Type(), + stream, + mr); + + // Signed offset minutes from the captured sign + HH(:MM); a missing offset + // (Z or no suffix) yields 0 (GMT). The sign is read from the sign character + // so "-00:30" stays negative. + auto offsetMinutes = + signedOffsetMinutes(g.column(8), g.column(9), g.column(10), stream, mr); + + // utcMillis = wallMillis - offsetMinutes * 60'000. + auto offsetMillis = binaryOp( + offsetMinutes->view(), + int64Scalar(60'000, stream), + cudf::binary_operator::MUL, + int64Type(), + stream, + mr); + auto utcMillis = cudf::binary_operation( + wallMillis->view(), + offsetMillis->view(), + cudf::binary_operator::SUB, + int64Type(), + stream, + mr); + + // zoneId from offset minutes: 0 -> 0; <0 -> off+841; >0 -> off+840. + auto zoneId = zoneKeyFromOffsetMinutes(offsetMinutes->view(), stream, mr); + + // An offset-less input is interpreted in the session timezone, not GMT, + // when one is set -- matching CPU's FromIso8601Timestamp, which reads the + // wall clock as that zone's local time (via Timestamp::toGMT) and packs the + // session zone key. Rows carrying an explicit "Z" or numeric offset keep + // the result computed above; the captured zone suffix (group 7) + // distinguishes them from an absent suffix. toUtcTimestamp does the exact + // DST-aware local->UTC conversion: it fails on a nonexistent local time + // (spring-forward gap) and resolves an ambiguous one (fall-back overlap) + // to the earliest instant, like CPU. Only the offset-less rows are + // converted; the rest are nulled out so their wall clock is never flagged + // as a gap. + std::unique_ptr selectedMillis; + std::unique_ptr selectedZone; + cudf::column_view finalMillis = utcMillis->view(); + cudf::column_view finalZone = zoneId->view(); + if (!context_.sessionTimezone.empty()) { + auto zoneSuffix = cudf::replace_nulls( + g.column(7), cudf::string_scalar("", true, stream), stream, mr); + auto suffixLength = cudf::strings::count_characters( + cudf::strings_column_view(zoneSuffix->view()), stream, mr); + auto hasExplicitZone = cudf::binary_operation( + suffixLength->view(), + cudf::numeric_scalar(0, true, stream), + cudf::binary_operator::GREATER, + cudf::data_type{kBool8}, + stream, + mr); + auto offsetless = cudf::binary_operation( + suffixLength->view(), + cudf::numeric_scalar(0, true, stream), + cudf::binary_operator::EQUAL, + cudf::data_type{kBool8}, + stream, + mr); + auto wallTimestamp = bitcastColumn( + wallMillis->view(), cudf::type_id::TIMESTAMP_MILLISECONDS); + // Null the explicit-zone rows so only the offset-less rows reach the gap + // check inside toUtcTimestamp; their result is discarded below anyway. + auto nullWall = cudf::make_default_constructed_scalar( + cudf::data_type{cudf::type_id::TIMESTAMP_MILLISECONDS}, stream, mr); + auto sessionWall = cudf::copy_if_else( + wallTimestamp, *nullWall, offsetless->view(), stream, mr); + auto sessionUtcTimestamp = toUtcTimestamp( + sessionWall->view(), context_.sessionTimezone, stream, mr); + auto sessionUtcMillis = + bitcastColumn(sessionUtcTimestamp->view(), kInt64); + const auto sessionZoneKey = tz::getTimeZoneID(context_.sessionTimezone); + selectedMillis = cudf::copy_if_else( + utcMillis->view(), + sessionUtcMillis, + hasExplicitZone->view(), + stream, + mr); + selectedZone = cudf::copy_if_else( + zoneId->view(), + int64Scalar(sessionZoneKey & kTimezoneMask, stream), + hasExplicitZone->view(), + stream, + mr); + finalMillis = selectedMillis->view(); + finalZone = selectedZone->view(); + } + + // pack(finalMillis, finalZone). + auto shifted = binaryOp( + finalMillis, + int64Scalar(kMillisShift, stream), + cudf::binary_operator::SHIFT_LEFT, + int64Type(), + stream, + mr); + return cudf::binary_operation( + shifted->view(), + finalZone, + cudf::binary_operator::BITWISE_OR, + int64Type(), + stream, + mr); + } + + private: + // Compiled ISO8601 field-extraction program. Batch-independent, so it is + // built once in the constructor and reused across eval calls. + std::unique_ptr isoProgram_; + // Recognizes a leading time-only form ("Thh...") so eval can prefix the epoch + // date and reuse isoProgram_. + std::unique_ptr timeOnlyProgram_; + // Matches an ISO8601 string whose year is signed or 5+ digits -- CPU-valid + // but unrepresentable by to_timestamps; eval raises VELOX_NYI for these. + std::unique_ptr extremeProgram_; +}; + +exec::FunctionSignaturePtr twtzArgSignature(const std::string& returnType) { + return exec::FunctionSignatureBuilder() + .returnType(returnType) + .argumentType("timestamp with time zone") + .build(); +} + +} // namespace + +void registerTimezoneFunctions(const std::string& prefix) { + using exec::FunctionSignatureBuilder; + + // The signatures below reference the TIMESTAMP WITH TIME ZONE custom type, + // which the worker has not registered yet when it registers cuDF (cuDF is + // registered before the CPU prestosql functions). Register it here instead of + // depending on registration order; registerCustomType is idempotent. + registerTimestampWithTimeZoneType(); + + registerCudfFunction( + prefix + "to_unixtime", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr); + }, + {twtzArgSignature("double")}); + + registerCudfFunction( + prefix + "at_timezone", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr); + }, + {FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .argumentType("timestamp with time zone") + .constantArgumentType("varchar") + .build()}); + + registerCudfFunction( + prefix + "timezone_hour", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr, /*minute=*/false); + }, + {twtzArgSignature("bigint")}); + + registerCudfFunction( + prefix + "timezone_minute", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr, /*minute=*/true); + }, + {twtzArgSignature("bigint")}); + + registerCudfFunction( + prefix + "to_iso8601", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr); + }, + {twtzArgSignature("varchar")}); + + registerCudfFunction( + prefix + "format_datetime", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr); + }, + {FunctionSignatureBuilder() + .returnType("varchar") + .argumentType("timestamp with time zone") + .constantArgumentType("varchar") + .build()}); + + registerCudfFunction( + prefix + "from_unixtime", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared( + tz::getTimeZoneID(constStringArg(expr, 1)), + FromUnixtimeRounding::kWhole); + }, + {FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .argumentType("double") + .constantArgumentType("varchar") + .build()}); + + registerCudfFunction( + prefix + "from_unixtime", + [](const std::string&, const std::shared_ptr& expr) { + // Compute hours*60 + minutes in int64 with overflow checks, mirroring + // CPU FromUnixtimeFunction; tz::getTimeZoneID then bounds the result to + // +/-840 minutes. Guards against a large hours value overflowing the + // product and truncating into a bogus in-range offset. + const auto offsetMinutes = checkedPlus( + checkedMultiply(constIntArg(expr, 1), 60), + constIntArg(expr, 2)); + return std::make_shared( + tz::getTimeZoneID(static_cast(offsetMinutes)), + FromUnixtimeRounding::kFloorThenFraction); + }, + {FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .argumentType("double") + .constantArgumentType("bigint") + .constantArgumentType("bigint") + .build()}); + + registerCudfFunction( + prefix + "parse_datetime", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr); + }, + {FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .argumentType("varchar") + .constantArgumentType("varchar") + .build()}); + + registerCudfFunction( + prefix + "from_iso8601_timestamp", + [](const std::string&, const std::shared_ptr& expr) { + return std::make_shared(expr); + }, + {FunctionSignatureBuilder() + .returnType("timestamp with time zone") + .argumentType("varchar") + .build()}); + + // now() / current_timestamp take no arguments; an empty signature list always + // matches by name. + registerCudfFunctions( + {prefix + "now", prefix + "current_timestamp"}, + [](const std::string&, const std::shared_ptr&) { + return std::make_shared(); + }, + {}); +} + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/prestosql/TimezoneFunctions.h b/velox/experimental/cudf/expression/prestosql/TimezoneFunctions.h new file mode 100644 index 00000000000..875bc3e4a22 --- /dev/null +++ b/velox/experimental/cudf/expression/prestosql/TimezoneFunctions.h @@ -0,0 +1,30 @@ +/* + * 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. + */ + +#pragma once + +#include + +namespace facebook::velox::cudf_velox { + +/// Registers GPU implementations of the Presto TIMESTAMP WITH TIME ZONE +/// function family: from_unixtime (with zone name or hour/minute offset), +/// to_unixtime, at_timezone, timezone_hour, timezone_minute, to_iso8601, +/// format_datetime, parse_datetime, from_iso8601_timestamp, and +/// now/current_timestamp. Names are prefixed with `prefix`. +void registerTimezoneFunctions(const std::string& prefix); + +} // namespace facebook::velox::cudf_velox diff --git a/velox/experimental/cudf/expression/sparksql/DateAddFunction.cpp b/velox/experimental/cudf/expression/sparksql/DateAddFunction.cpp index ad009bc3382..21ad1dce7ee 100644 --- a/velox/experimental/cudf/expression/sparksql/DateAddFunction.cpp +++ b/velox/experimental/cudf/expression/sparksql/DateAddFunction.cpp @@ -38,6 +38,7 @@ DateAddFunction::DateAddFunction( ColumnOrView DateAddFunction::eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const { auto inputCol = asView(inputColumns[0]); diff --git a/velox/experimental/cudf/expression/sparksql/DateAddFunction.h b/velox/experimental/cudf/expression/sparksql/DateAddFunction.h index 9af4e405613..ff692f61a40 100644 --- a/velox/experimental/cudf/expression/sparksql/DateAddFunction.h +++ b/velox/experimental/cudf/expression/sparksql/DateAddFunction.h @@ -32,6 +32,7 @@ class DateAddFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override; diff --git a/velox/experimental/cudf/expression/sparksql/HashFunction.cpp b/velox/experimental/cudf/expression/sparksql/HashFunction.cpp index 4fcb6e9a1c0..f7f43ea8f5a 100644 --- a/velox/experimental/cudf/expression/sparksql/HashFunction.cpp +++ b/velox/experimental/cudf/expression/sparksql/HashFunction.cpp @@ -63,6 +63,7 @@ HashFunction::HashFunction(const std::shared_ptr& expr) { ColumnOrView HashFunction::eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const { VELOX_CHECK(!inputColumns.empty()); diff --git a/velox/experimental/cudf/expression/sparksql/HashFunction.h b/velox/experimental/cudf/expression/sparksql/HashFunction.h index 948e1f27889..f685c10178d 100644 --- a/velox/experimental/cudf/expression/sparksql/HashFunction.h +++ b/velox/experimental/cudf/expression/sparksql/HashFunction.h @@ -29,6 +29,7 @@ class HashFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override; diff --git a/velox/experimental/cudf/expression/sparksql/SubStringFunction.cpp b/velox/experimental/cudf/expression/sparksql/SubStringFunction.cpp index 08b766edac7..8fe32e493ae 100644 --- a/velox/experimental/cudf/expression/sparksql/SubStringFunction.cpp +++ b/velox/experimental/cudf/expression/sparksql/SubStringFunction.cpp @@ -92,6 +92,7 @@ class SubStringFunction : public CudfFunction { ColumnOrView eval( std::vector& inputColumns, + [[maybe_unused]] cudf::size_type numRows, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const override { VELOX_CHECK( diff --git a/velox/experimental/cudf/tests/CMakeLists.txt b/velox/experimental/cudf/tests/CMakeLists.txt index 3b3d02f6b16..2885f13b519 100644 --- a/velox/experimental/cudf/tests/CMakeLists.txt +++ b/velox/experimental/cudf/tests/CMakeLists.txt @@ -151,6 +151,18 @@ velox_add_cudf_test( LIBS velox_cudf_expression velox_cudf_exec velox_exec velox_exec_test_lib velox_test_util ) +velox_add_cudf_test( + NAME velox_cudf_timezone_extraction_test + SOURCES Main.cpp TimezoneExtractionTest.cpp + LIBS ${CUDF_TEST_DEFAULT_LIBS} velox_functions_test_lib +) + +velox_add_cudf_test( + NAME velox_cudf_timezone_function_test + SOURCES Main.cpp TimezoneFunctionTest.cpp + LIBS ${CUDF_TEST_DEFAULT_LIBS} velox_functions_test_lib +) + velox_add_cudf_test( NAME velox_cudf_group_id_test SOURCES Main.cpp GroupIdTest.cpp diff --git a/velox/experimental/cudf/tests/CudfFunctionBaseTest.h b/velox/experimental/cudf/tests/CudfFunctionBaseTest.h index 7ec3a2053e0..82143f172cf 100644 --- a/velox/experimental/cudf/tests/CudfFunctionBaseTest.h +++ b/velox/experimental/cudf/tests/CudfFunctionBaseTest.h @@ -48,8 +48,13 @@ class CudfFunctionBaseTest : public velox::functions::test::FunctionBaseTest { auto stream = cudf::get_default_stream(); auto cudfTable = velox::cudf_velox::with_arrow::toCudfTable( input, pool_.get(), stream, cudf::get_current_device_resource_ref()); - auto filterEvaluator = - createCudfExpression({exprSet.exprs()[0]}, input->rowType()); + // Build the evaluation context from the query config exactly as + // CudfFilterProject does, so a test can exercise timezone-aware functions + // under a session timezone via queryCtx_->testingOverrideConfigUnsafe. + const auto exprContext = + contextFromConfig(execCtx_.queryCtx()->queryConfig()); + auto filterEvaluator = createCudfExpression( + {exprSet.exprs()[0]}, input->rowType(), exprContext); auto ownedColumns = cudfTable->release(); std::vector inputViews; inputViews.reserve(ownedColumns.size()); diff --git a/velox/experimental/cudf/tests/FilterProjectTest.cpp b/velox/experimental/cudf/tests/FilterProjectTest.cpp index 9419960ab85..ca4081ac44d 100644 --- a/velox/experimental/cudf/tests/FilterProjectTest.cpp +++ b/velox/experimental/cudf/tests/FilterProjectTest.cpp @@ -20,14 +20,17 @@ #include "velox/common/base/tests/GTestUtils.h" #include "velox/core/Expressions.h" +#include "velox/core/QueryConfig.h" #include "velox/dwio/common/tests/utils/BatchMaker.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/OperatorTestBase.h" #include "velox/exec/tests/utils/PlanBuilder.h" #include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h" #include "velox/functions/prestosql/registration/RegistrationFunctions.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/parse/TypeResolver.h" #include "velox/type/Time.h" +#include "velox/type/tz/TimeZoneMap.h" #include @@ -630,6 +633,27 @@ class CudfFilterProjectTest : public OperatorTestBase { assertPlanMatchesVelox(plan); } + // Runs a projection plan under a session timezone (adjust_timestamp enabled) + // on GPU and CPU and asserts equal results. Used by the timezone-aware + // date_add(timestamp) tests. + void assertProjectMatchesVeloxWithTimezone( + const std::vector& input, + const std::vector& projections, + const std::string& sessionTimezone) { + auto plan = PlanBuilder().values(input).project(projections).planNode(); + auto run = [&]() { + return AssertQueryBuilder(plan) + .config(core::QueryConfig::kSessionTimezone, sessionTimezone) + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()); + }; + auto cudfResult = run(); + cudf_velox::unregisterCudf(); + auto veloxResult = run(); + cudf_velox::registerCudf(); + facebook::velox::test::assertEqualVectors(veloxResult, cudfResult); + } + void runTest(core::PlanNodePtr planNode, const std::string& duckDbSql) { SCOPED_TRACE("run without spilling"); assertQuery(planNode, duckDbSql); @@ -1346,6 +1370,216 @@ TEST_F(CudfFilterProjectTest, dateAddDateScaledOverflowMatchesVelox) { assertProjectMatchesVelox(vectors, projections); } +// date_add(unit, value, timestamp): under a session timezone, addToTimestamp +// adds on the local wall clock for every unit; matches CPU. +TEST_F(CudfFilterProjectTest, dateAddTimestampSessionTimezone) { + auto input = makeRowVector({makeNullableFlatVector( + {Timestamp(1'736'971'261, 123'000'000), // 2025-01-15 20:01:01.123 UTC + Timestamp(1'709'251'200, 0), + std::nullopt}, + TIMESTAMP())}); + std::vector projections; + for (const auto* unit : + {"second", + "minute", + "hour", + "day", + "week", + "month", + "quarter", + "year"}) { + projections.push_back( + std::string("date_add('") + unit + "', 3, c0) AS add_" + unit); + projections.push_back( + std::string("date_add('") + unit + "', -5, c0) AS sub_" + unit); + } + assertProjectMatchesVeloxWithTimezone( + {input}, projections, "America/Los_Angeles"); +} + +// A fractional-offset zone (+05:30). +TEST_F(CudfFilterProjectTest, dateAddTimestampFractionalOffset) { + auto input = makeRowVector({makeFlatVector( + {Timestamp(1'736'971'261, 0), Timestamp(0, 0)}, TIMESTAMP())}); + assertProjectMatchesVeloxWithTimezone( + {input}, + {"date_add('hour', 5, c0) AS h", + "date_add('day', 2, c0) AS d", + "date_add('month', 1, c0) AS mo"}, + "Asia/Kolkata"); +} + +// With no session timezone the add is on the raw UTC instant. +TEST_F(CudfFilterProjectTest, dateAddTimestampNoSessionTimezone) { + auto input = makeRowVector({makeFlatVector( + {Timestamp(1'736'971'261, 0), Timestamp(0, 0)}, TIMESTAMP())}); + assertProjectMatchesVelox( + {input}, + {"date_add('hour', 5, c0) AS h", + "date_add('day', 2, c0) AS d", + "date_add('month', -3, c0) AS mo"}); +} + +// A per-row value column. +TEST_F(CudfFilterProjectTest, dateAddTimestampColumnValue) { + auto input = makeRowVector( + {"amt", "c0"}, + {makeFlatVector({1, -2, 100}), + makeFlatVector( + {Timestamp(1'736'971'261, 0), + Timestamp(0, 0), + Timestamp(1'709'251'200, 0)}, + TIMESTAMP())}); + assertProjectMatchesVeloxWithTimezone( + {input}, + {"date_add('month', amt, c0) AS mo", "date_add('hour', amt, c0) AS h"}, + "Asia/Kolkata"); +} + +// A day add whose result lands in a spring-forward gap must throw, matching +// CPU toGMT. 2024-03-09 02:30 America/Los_Angeles + 1 day == 2024-03-10 02:30, +// which does not exist. +TEST_F(CudfFilterProjectTest, dateAddTimestampSpringForwardThrows) { + auto input = makeRowVector( + {makeFlatVector({Timestamp(1'709'980'200, 0)}, TIMESTAMP())}); + auto plan = PlanBuilder() + .values({input}) + .project({"date_add('day', 1, c0) AS result"}) + .planNode(); + VELOX_ASSERT_THROW( + AssertQueryBuilder(plan) + .config(core::QueryConfig::kSessionTimezone, "America/Los_Angeles") + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()), + "does not exist in the time zone"); +} + +// A value outside int32 range throws, matching CPU checkValueInInt32Range. +TEST_F(CudfFilterProjectTest, dateAddTimestampValueOutOfRange) { + auto input = makeRowVector( + {makeFlatVector({Timestamp(0, 0)}, TIMESTAMP())}); + auto plan = PlanBuilder() + .values({input}) + .project({"date_add('day', 3000000000, c0) AS result"}) + .planNode(); + VELOX_ASSERT_THROW( + AssertQueryBuilder(plan).copyResults(pool()), + "date_add value is out of range"); +} + +// date_add(timestamp with time zone) adds on each row's embedded zone, +// independent of the session zone, matching CPU addToTimestampWithTimezone. +// Sub-day units add on the UTC instant; day-and-above units add on the local +// wall clock and convert back to UTC. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneUnits) { + // 2025-01-15 20:01:01.123 UTC == 12:01:01.123 America/Los_Angeles. + auto input = makeRowVector({makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + std::vector projections; + for (const auto* unit : + {"second", + "minute", + "hour", + "day", + "week", + "month", + "quarter", + "year"}) { + projections.push_back( + std::string("date_add('") + unit + "', 3, c0) AS add_" + unit); + projections.push_back( + std::string("date_add('") + unit + "', -5, c0) AS sub_" + unit); + } + assertProjectMatchesVelox({input}, projections); +} + +// A fractional-offset zone (+05:30) exercises both the sub-day UTC add and the +// day+ local->UTC round-trip on a non-whole-hour offset. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneFractionalOffset) { + auto input = makeRowVector({makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox( + {input}, + {"date_add('hour', 5, c0) AS h", + "date_add('day', 2, c0) AS d", + "date_add('month', 1, c0) AS mo"}); +} + +// Each row adds on its own embedded zone. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneMixedZones) { + auto input = makeRowVector({makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles")), + pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox( + {input}, + {"date_add('day', 10, c0) AS d", "date_add('hour', 2, c0) AS h"}); +} + +// A per-row value column. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneColumnValue) { + auto input = makeRowVector( + {"amt", "c0"}, + {makeFlatVector({1, -2, 100}), + makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles")), + pack(0, tz::getTimeZoneID("Asia/Kolkata")), + pack(1'709'251'200'000, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox( + {input}, + {"date_add('month', amt, c0) AS mo", "date_add('hour', amt, c0) AS h"}); +} + +// A null row stays null through the TSWTZ path. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneNull) { + auto input = makeRowVector({makeNullableFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata")), + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox( + {input}, {"date_add('day', 3, c0) AS d", "date_add('hour', 3, c0) AS h"}); +} + +// Unlike date_add(timestamp) under a session zone, a TSWTZ day-add whose result +// lands in a spring-forward gap does NOT throw: it resolves the nonexistent +// local forward, matching CPU addToTimestampWithTimezone. 2024-03-09 02:30 +// America/Los_Angeles + 1 day == 2024-03-10 02:30, which does not exist. +TEST_F( + CudfFilterProjectTest, + dateAddTimestampWithTimeZoneSpringForwardNoThrow) { + auto input = makeRowVector({makeFlatVector( + {pack(1'709'980'200'000, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox({input}, {"date_add('day', 1, c0) AS d"}); +} + +// A day add whose result lands in a fall-back overlap resolves to the earliest +// instant, matching CPU to_sys(kEarliest). 2024-11-02 01:30 America/Los_Angeles +// (PDT) + 1 day == 2024-11-03 01:30, ambiguous during the fall-back overlap. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneFallBack) { + auto input = makeRowVector({makeFlatVector( + {pack(1'730'536'200'000, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox({input}, {"date_add('day', 1, c0) AS d"}); +} + +// A value outside int32 range throws, matching CPU checkValueInInt32Range. +TEST_F(CudfFilterProjectTest, dateAddTimestampWithTimeZoneValueOutOfRange) { + auto input = makeRowVector({makeFlatVector( + {pack(0, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + auto plan = PlanBuilder() + .values({input}) + .project({"date_add('day', 3000000000, c0) AS result"}) + .planNode(); + VELOX_ASSERT_THROW( + AssertQueryBuilder(plan).copyResults(pool()), + "date_add value is out of range"); +} + TEST_F(CudfFilterProjectTest, dateTruncTimestampUnits) { auto vectors = makeTimestampExtractVectors(); const std::vector projections{ @@ -1361,6 +1595,69 @@ TEST_F(CudfFilterProjectTest, dateTruncTimestampUnits) { assertProjectMatchesVelox(vectors, projections); } +// date_trunc(timestamp with time zone) truncates on each row's embedded zone +// (independent of the session zone), matching CPU +// DateTruncFunction::call(TSWTZ). +TEST_F(CudfFilterProjectTest, dateTruncTimestampWithTimeZoneUnits) { + // 2025-01-15 20:01:01.123 UTC == 12:01:01.123 America/Los_Angeles. + auto input = makeRowVector({makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + const std::vector projections{ + "date_trunc('second', c0) AS s", + "date_trunc('minute', c0) AS mi", + "date_trunc('hour', c0) AS h", + "date_trunc('day', c0) AS d", + "date_trunc('week', c0) AS w", + "date_trunc('month', c0) AS mo", + "date_trunc('quarter', c0) AS q", + "date_trunc('year', c0) AS y"}; + assertProjectMatchesVelox({input}, projections); +} + +// A fractional-offset zone (+05:30) exercises the sub-day delta and the day+ +// local->UTC round-trip on a non-whole-hour offset. +TEST_F(CudfFilterProjectTest, dateTruncTimestampWithTimeZoneFractionalOffset) { + auto input = makeRowVector({makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + const std::vector projections{ + "date_trunc('hour', c0) AS h", + "date_trunc('day', c0) AS d", + "date_trunc('month', c0) AS mo"}; + assertProjectMatchesVelox({input}, projections); +} + +// Each row truncates on its own embedded zone. +TEST_F(CudfFilterProjectTest, dateTruncTimestampWithTimeZoneMixedZones) { + auto input = makeRowVector({makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles")), + pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox( + {input}, {"date_trunc('day', c0) AS d", "date_trunc('hour', c0) AS h"}); +} + +// A null row stays null through the TSWTZ path. +TEST_F(CudfFilterProjectTest, dateTruncTimestampWithTimeZoneNull) { + auto input = makeRowVector({makeNullableFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata")), + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox({input}, {"date_trunc('month', c0) AS mo"}); +} + +// A fall-back overlap: day-trunc crosses the ambiguous local midnight; the +// local->UTC conversion must resolve to the earliest instant (matches toGMT). +TEST_F(CudfFilterProjectTest, dateTruncTimestampWithTimeZoneDstFallBack) { + // 2024-11-03 08:30:00 UTC == 01:30 America/Los_Angeles during the fall-back + // overlap. + auto input = makeRowVector({makeFlatVector( + {pack(1'730'622'600'000, tz::getTimeZoneID("America/Los_Angeles"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertProjectMatchesVelox({input}, {"date_trunc('day', c0) AS d"}); +} + TEST_F(CudfFilterProjectTest, dateTruncDateUnits) { auto vectors = makeTimestampExtractVectors(); const std::vector projections{ diff --git a/velox/experimental/cudf/tests/FunctionRegistryTest.cpp b/velox/experimental/cudf/tests/FunctionRegistryTest.cpp index 29f143354ac..b34d2a68e2a 100644 --- a/velox/experimental/cudf/tests/FunctionRegistryTest.cpp +++ b/velox/experimental/cudf/tests/FunctionRegistryTest.cpp @@ -15,10 +15,12 @@ */ #include "velox/experimental/cudf/expression/ExpressionEvaluator.h" +#include "velox/experimental/cudf/expression/prestosql/TimezoneFunctions.h" #include "velox/expression/Expr.h" #include "velox/expression/FieldReference.h" #include "velox/expression/FunctionSignature.h" +#include "velox/type/Type.h" #include @@ -41,6 +43,7 @@ class TagFunction : public CudfFunction { ColumnOrView eval( std::vector& /*inputColumns*/, + cudf::size_type /*numRows*/, rmm::cuda_stream_view /*stream*/, rmm::device_async_resource_ref /*mr*/) const override { VELOX_UNREACHABLE("TagFunction::eval should not be called in these tests"); @@ -111,6 +114,24 @@ TEST_F(FunctionRegistryTest, singleSignatureDispatch) { EXPECT_EQ(tagOf(fn), "double"); } +// The timezone function signatures reference the TIMESTAMP WITH TIME ZONE +// custom type, which lives in the prestosql type registry. The Presto worker +// registers cuDF before the CPU prestosql functions, so that type is absent +// when these signatures are built; registerTimezoneFunctions must register the +// type itself rather than rely on registration order (otherwise +// FunctionSignatureBuilder aborts with "Type doesn't exist: 'TIMESTAMP WITH +// TIME ZONE'" at worker startup). +TEST_F(FunctionRegistryTest, timezoneFunctionsRegisterTheirCustomType) { + // Reproduce a clean worker startup where the type has not been registered + // yet, independent of registration order within this binary. + unregisterCustomType("TIMESTAMP WITH TIME ZONE"); + ASSERT_FALSE(hasType("TIMESTAMP WITH TIME ZONE")); + + EXPECT_NO_THROW(registerTimezoneFunctions("regtest_tz_")); + + EXPECT_TRUE(hasType("TIMESTAMP WITH TIME ZONE")); +} + // Two function registrations, same name collision with different signatures. // Test that both can coexist. TEST_F(FunctionRegistryTest, multipleSignaturesDispatchByInputTypes) { diff --git a/velox/experimental/cudf/tests/HashJoinTest.cpp b/velox/experimental/cudf/tests/HashJoinTest.cpp index 462fa1c9c60..20f2afb8d2c 100644 --- a/velox/experimental/cudf/tests/HashJoinTest.cpp +++ b/velox/experimental/cudf/tests/HashJoinTest.cpp @@ -436,6 +436,54 @@ TEST_P(MultiThreadedHashJoinTest, filter) { .run(); } +// A timezone-sensitive join filter must read the session-local wall clock, the +// same way the CPU evaluates it. Both probe rows join their build key, so only +// the filter hour(t_ts) = 18 decides which row survives. Under +// America/Los_Angeles (UTC-8) the two instants land on opposite sides of the +// predicate from their UTC values: 2021-01-01 02:00 UTC is the local 18:00 and +// passes, while 2021-01-01 18:00 UTC is the local 10:00 and fails. Running the +// same plan on GPU (cuDF registered) and CPU (cuDF unregistered) under the same +// session timezone must therefore select the same probe row. Until the join +// filter honors the session timezone the GPU keeps the UTC-18:00 row instead. +TEST_F(HashJoinTest, joinFilterHonorsSessionTimezone) { + // 2021-01-01 18:00:00 UTC == 2021-01-01 10:00:00 America/Los_Angeles. + constexpr int64_t kUtc1800 = 1'609'524'000; + // 2021-01-01 02:00:00 UTC == 2020-12-31 18:00:00 America/Los_Angeles. + constexpr int64_t kLocal1800 = 1'609'466'400; + + auto probe = makeRowVector( + {"t_k0", "t_ts"}, + {makeFlatVector({1, 2}), + makeFlatVector( + {Timestamp(kUtc1800, 0), Timestamp(kLocal1800, 0)}, TIMESTAMP())}); + auto build = makeRowVector({"u_k0"}, {makeFlatVector({1, 2})}); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = + PlanBuilder(planNodeIdGenerator) + .values({probe}) + .hashJoin( + {"t_k0"}, + {"u_k0"}, + PlanBuilder(planNodeIdGenerator).values({build}).planNode(), + "hour(t_ts) = 18", + {"t_k0"}) + .planNode(); + + auto run = [&]() { + return AssertQueryBuilder(plan) + .config(core::QueryConfig::kSessionTimezone, "America/Los_Angeles") + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()); + }; + + auto gpu = run(); + cudf_velox::unregisterCudf(); + auto cpu = run(); + cudf_velox::registerCudf(); + facebook::velox::test::assertEqualVectors(cpu, gpu); +} + DEBUG_ONLY_TEST_P(MultiThreadedHashJoinTest, filterSpillOnFirstProbeInput) { auto spillDirectory = TempDirectoryPath::create(); std::atomic_bool injectProbeSpillOnce{true}; diff --git a/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp b/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp index d4e455c1c9c..c2b6d78bb6a 100644 --- a/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp +++ b/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp @@ -14,12 +14,16 @@ * limitations under the License. */ +#include "velox/experimental/cudf/CudfConfig.h" #include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/core/QueryConfig.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/HiveConnectorTestBase.h" #include "velox/exec/tests/utils/PlanBuilder.h" +#include + using namespace facebook::velox; using namespace facebook::velox::exec; using namespace facebook::velox::exec::test; @@ -184,6 +188,67 @@ TEST_F(CudfNestedLoopJoinTest, innerJoinWithFilter) { assertQuery(plan, "SELECT t.c0, u.c0 FROM t INNER JOIN u ON t.c0 < u.c0"); } +// A timezone-sensitive join condition must read the session-local wall clock, +// the same way the CPU evaluates it. The single build row makes the cross +// product equal to the probe rows, so only hour(l_ts) = 18 decides which row +// survives. Under America/Los_Angeles (UTC-8) the two instants land on opposite +// sides of the predicate from their UTC values: 2021-01-01 02:00 UTC is the +// local 18:00 and passes, while 2021-01-01 18:00 UTC is the local 10:00 and +// fails. Running the same plan on GPU (cuDF registered) and CPU (cuDF +// unregistered) under the same session timezone must therefore select the same +// probe row. Until the condition honors the session timezone the GPU keeps the +// UTC-18:00 row instead. +TEST_F(CudfNestedLoopJoinTest, joinConditionHonorsSessionTimezone) { + // 2021-01-01 18:00:00 UTC == 2021-01-01 10:00:00 America/Los_Angeles. + constexpr int64_t kUtc1800 = 1'609'524'000; + // 2021-01-01 02:00:00 UTC == 2020-12-31 18:00:00 America/Los_Angeles. + constexpr int64_t kLocal1800 = 1'609'466'400; + + auto left = makeRowVector( + {"l_id", "l_ts"}, + {makeFlatVector({1, 2}), + makeFlatVector( + {Timestamp(kUtc1800, 0), Timestamp(kLocal1800, 0)}, TIMESTAMP())}); + auto right = makeRowVector({"r_id"}, {makeFlatVector({0})}); + + auto planNodeIdGenerator = std::make_shared(); + auto plan = + PlanBuilder(planNodeIdGenerator) + .values({left}) + .nestedLoopJoin( + PlanBuilder(planNodeIdGenerator).values({right}).planNode(), + "hour(l_ts) = 18", + {"l_id"}, + core::JoinType::kInner) + .planNode(); + + auto run = [&]() { + return AssertQueryBuilder(plan) + .config(core::QueryConfig::kSessionTimezone, "America/Los_Angeles") + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()); + }; + + RowVectorPtr gpu; + { + // Disable CPU fallback for the GPU run so a missing GPU path throws instead + // of silently running on CPU and hiding the timezone gap. The CPU run below + // does not need this guard because it unregisters cuDF entirely. + auto& cudfConfig = cudf_velox::CudfConfig::getInstance(); + const bool savedAllowCpuFallback = cudfConfig.allowCpuFallback; + cudfConfig.allowCpuFallback = false; + SCOPE_EXIT { + cudfConfig.allowCpuFallback = savedAllowCpuFallback; + }; + gpu = run(); + } + + cudf_velox::unregisterCudf(); + auto cpu = run(); + cudf_velox::registerCudf(); + facebook::velox::test::assertEqualVectors(cpu, gpu); +} + // Test 6: Multiple batches (tests streaming behavior) TEST_F(CudfNestedLoopJoinTest, multipleBatches) { std::vector probeVectors; diff --git a/velox/experimental/cudf/tests/TableScanTest.cpp b/velox/experimental/cudf/tests/TableScanTest.cpp index 6b3a366896c..c8d41687ff7 100644 --- a/velox/experimental/cudf/tests/TableScanTest.cpp +++ b/velox/experimental/cudf/tests/TableScanTest.cpp @@ -818,3 +818,52 @@ TEST_F(TableScanTest, decimalRemainingFilter) { {filePath}, "SELECT c0, c1 FROM tmp WHERE c0 = CAST('-5.00' AS DECIMAL(5, 2))"); } + +// hour(ts) selects different rows under UTC versus a non-UTC session timezone, +// so the remaining filter on the cuDF hive scan must evaluate hour() in the +// session timezone. Two instants straddle that boundary: +// kUtc1800 2021-01-01 18:00:00 UTC == 2021-01-01 10:00:00 Los_Angeles +// kLocal1800 2021-01-01 02:00:00 UTC == 2020-12-31 18:00:00 Los_Angeles +// Under UTC "hour(ts) = 18" keeps id 1; under America/Los_Angeles it keeps +// id 2. The hive path always evaluates the remaining filter on the GPU and the +// fixture keeps the cuDF hive connector registered, so a registry-toggle CPU +// oracle is not available; the expected ids are built by hand instead. +TEST_F(TableScanTest, remainingFilterHonorsSessionTimezone) { + constexpr int64_t kUtc1800 = 1'609'524'000; + constexpr int64_t kLocal1800 = 1'609'466'400; + auto rowType = ROW({"id", "ts"}, {BIGINT(), TIMESTAMP()}); + auto vector = makeRowVector( + {"id", "ts"}, + {makeFlatVector({1, 2}), + makeFlatVector( + {Timestamp(kUtc1800, 0), Timestamp(kLocal1800, 0)}, TIMESTAMP())}); + + auto filePath = TempFilePath::create(); + writeToFile(filePath->getPath(), {vector}); + + auto assignments = + facebook::velox::exec::test::HiveConnectorTestBase::allRegularColumns( + rowType); + + auto plan = PlanBuilder(pool_.get()) + .startTableScan() + .connectorId(kCudfHiveConnectorId) + .outputType(rowType) + .dataColumns(rowType) + .assignments(assignments) + .remainingFilter("hour(ts) = 18") + .endTableScan() + .planNode(); + + auto result = + AssertQueryBuilder(plan) + .split(makeCudfHiveSplit(filePath->getPath())) + .config(core::QueryConfig::kSessionTimezone, "America/Los_Angeles") + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()); + + // Assert only on the id column to stay independent of how the timestamp + // round-trips through the parquet writer and the cuDF reader. + auto expectedIds = makeFlatVector({2}); + facebook::velox::test::assertEqualVectors(expectedIds, result->childAt(0)); +} diff --git a/velox/experimental/cudf/tests/TimezoneExtractionTest.cpp b/velox/experimental/cudf/tests/TimezoneExtractionTest.cpp new file mode 100644 index 00000000000..4b9fc9846bc --- /dev/null +++ b/velox/experimental/cudf/tests/TimezoneExtractionTest.cpp @@ -0,0 +1,400 @@ +/* + * 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. + */ + +// Class A correctness tests: GPU (cuDF) date/time extraction functions must +// honor the session timezone exactly as CPU does. These reproduce the missing +// timezone support by failing while it is absent. +// +// On CPU, year/month/day/hour/... read the session timezone (via +// getTimeZoneFromConfig + getDateTime in velox/functions/lib/TimeUtils.h) and +// return the component of the *local* wall-clock time whenever +// adjust_timestamp_to_session_timezone is set. On GPU, +// ExtractComponentFunction::eval (velox/experimental/cudf/expression/ +// ExpressionEvaluator.cpp) currently calls +// cudf::datetime::extract_datetime_component directly on the raw epoch column, +// so it always returns the UTC component and ignores the session timezone. +// +// Each test runs the same projection twice under a non-UTC session timezone -- +// once with cuDF registered (GPU) and once without (CPU) -- and asserts the two +// results match. The timestamps are chosen so the local value lands in a +// different calendar field than the UTC value, so the assertion holds only once +// the GPU path applies the session-timezone offset before extracting the field. +// A matching control under UTC proves a failure is specifically timezone-driven +// and not a pre-existing extraction bug. +// +// IMPORTANT: these are test-driven-development tests for the *target* behavior, +// so they FAIL today and pass once the gap is closed. Each test below fails +// with the local-vs-UTC mismatch noted in its comment until the GPU path +// applies the session-timezone offset before extraction. A red test here means +// the timezone gap is still open; do not "fix" a failure by weakening the +// assertion. The UTC controls and the sub-minute fields already pass and guard +// the harness. +// +// The plan/operator path is used here rather than +// CudfFunctionBaseTest::assertExpressionMatchesCpu because that lightweight +// harness evaluates the expression with finalize=false and cannot relabel a +// narrow cuDF result (e.g. extract_datetime_component returns SMALLINT) to the +// Velox BIGINT result type. The operator path applies the finalizing cast, the +// same way a real query does, and is already exercised by +// FilterProjectTest.extractTimestampComponents. +// +// These tests require a GPU and are labeled cuda_driver; they will not run in a +// CPU-only environment. + +#include "velox/experimental/cudf/CudfConfig.h" +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/expression/ExpressionEvaluator.h" +#include "velox/experimental/cudf/expression/PrestoFunctions.h" + +#include "velox/common/file/FileSystems.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" +#include "velox/exec/tests/utils/OperatorTestBase.h" +#include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h" +#include "velox/functions/prestosql/registration/RegistrationFunctions.h" +#include "velox/parse/TypeResolver.h" + +using namespace facebook::velox; +using namespace facebook::velox::exec; +using namespace facebook::velox::exec::test; + +namespace { + +class TimezoneExtractionTest : public OperatorTestBase { + protected: + void SetUp() override { + OperatorTestBase::SetUp(); + filesystems::registerLocalFileSystem(); + cudf_velox::CudfConfig::getInstance().allowCpuFallback = false; + cudf_velox::registerCudf(); + cudf_velox::registerPrestoFunctions( + cudf_velox::CudfConfig::getInstance().functionNamePrefix); + } + + void TearDown() override { + cudf_velox::unregisterFunctions(); + cudf_velox::unregisterCudf(); + OperatorTestBase::TearDown(); + } + + // Builds a single-row TIMESTAMP input column named ts. + RowVectorPtr timestampInput(int64_t seconds, uint64_t nanos = 0) { + return makeRowVector( + {"ts"}, + {makeFlatVector({Timestamp(seconds, nanos)}, TIMESTAMP())}); + } + + // Evaluates the projection with the given session timezone. cuDF is expected + // to be registered or unregistered by the caller, which selects GPU or CPU + // execution respectively (allowCpuFallback is false, so a registered cuDF + // never falls back to CPU). + RowVectorPtr project( + const RowVectorPtr& input, + const std::string& projection, + std::string_view timezone) { + auto plan = PlanBuilder().values({input}).project({projection}).planNode(); + return AssertQueryBuilder(plan) + .config(core::QueryConfig::kSessionTimezone, std::string(timezone)) + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()); + } + + // Runs the projection on GPU (cuDF registered) and CPU (cuDF unregistered) + // under the same session timezone and asserts the single output columns are + // equal. This is the target behavior for every extraction function: the GPU + // result must equal the CPU result regardless of the session timezone. + void assertGpuMatchesCpu( + const RowVectorPtr& input, + const std::string& projection, + std::string_view timezone) { + auto gpu = project(input, projection, timezone); + cudf_velox::unregisterCudf(); + auto cpu = project(input, projection, timezone); + cudf_velox::registerCudf(); + SCOPED_TRACE( + projection + " under session timezone " + std::string(timezone)); + facebook::velox::test::assertEqualVectors(cpu->childAt(0), gpu->childAt(0)); + } + + // Filters then projects under the given session timezone. cuDF registration + // by the caller selects GPU vs CPU execution, exactly like project(). + RowVectorPtr filterProject( + const RowVectorPtr& input, + const std::string& filter, + const std::string& projection, + std::string_view timezone) { + auto plan = PlanBuilder() + .values({input}) + .filter(filter) + .project({projection}) + .planNode(); + return AssertQueryBuilder(plan) + .config(core::QueryConfig::kSessionTimezone, std::string(timezone)) + .config(core::QueryConfig::kAdjustTimestampToTimezone, "true") + .copyResults(pool()); + } + + // Runs a filter+project on GPU (cuDF registered) and CPU (cuDF unregistered) + // under the same session timezone and asserts the projected columns are + // equal. The filter is the work this exercises: a non-AST function inside the + // predicate (e.g. hour(ts)) is precomputed as a column, so the precompute + // must honor the session timezone or the GPU selects the wrong rows. + void assertFilterGpuMatchesCpu( + const RowVectorPtr& input, + const std::string& filter, + const std::string& projection, + std::string_view timezone) { + auto gpu = filterProject(input, filter, projection, timezone); + cudf_velox::unregisterCudf(); + auto cpu = filterProject(input, filter, projection, timezone); + cudf_velox::registerCudf(); + SCOPED_TRACE( + "filter " + filter + ", project " + projection + + " under session timezone " + std::string(timezone)); + facebook::velox::test::assertEqualVectors(cpu->childAt(0), gpu->childAt(0)); + } +}; + +// America/Los_Angeles is UTC-8 in January. This instant is 2021-01-01 02:00:00 +// UTC, which is 2020-12-31 18:00:00 local, so the local year/month/day/quarter/ +// hour/day_of_week/day_of_year all land in the previous day, month, quarter and +// year. +constexpr int64_t kJan2021At0200Utc = 1'609'466'400; + +// 2021-01-04 02:00:00 UTC is a Monday (ISO week 1 of 2021); 2021-01-03 18:00:00 +// America/Los_Angeles is the preceding Sunday, which still belongs to ISO week +// 53 of week-year 2020. Used for the week / year_of_week fields. +constexpr int64_t kJan2021MondayUtc = 1'609'725'600; + +// 2021-01-01 00:00:00 UTC. Asia/Kolkata is UTC+5:30, a half-hour offset, so the +// local minute differs from the UTC minute. Used for the minute field, which a +// whole-hour offset zone like America/Los_Angeles cannot exercise. Under +// Asia/Kolkata the local hour is 5 (UTC hour 0). +constexpr int64_t kJan2021MidnightUtc = 1'609'459'200; + +// 2021-01-01 05:00:00 UTC. Under Asia/Kolkata (+5:30) the local hour is 10 +// while the UTC hour is 5. Paired with kJan2021MidnightUtc so a filter on +// hour = 5 selects different rows in UTC (this one) and in Asia/Kolkata (the +// midnight row), exposing a filter that evaluates hour() in the wrong zone. +constexpr int64_t kJan2021At0500Utc = 1'609'477'200; + +constexpr std::string_view kLosAngeles = "America/Los_Angeles"; +constexpr std::string_view kKolkata = "Asia/Kolkata"; + +TEST_F(TimezoneExtractionTest, yearHonorsSessionTimezone) { + // Expect local year 2020; GPU currently returns UTC year 2021. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "year(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, monthHonorsSessionTimezone) { + // Expect local month 12; GPU currently returns UTC month 1. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "month(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, dayHonorsSessionTimezone) { + // Expect local day 31; GPU currently returns UTC day 1. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "day(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, quarterHonorsSessionTimezone) { + // Expect local quarter 4; GPU currently returns UTC quarter 1. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "quarter(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, hourHonorsSessionTimezone) { + // Expect local hour 18; GPU currently returns UTC hour 2. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "hour(ts)", kLosAngeles); +} + +// A filter "hour(ts) = 5" routes through the createAstTree precompute path: the +// comparison is an AST node, but hour(ts) is a non-AST function precomputed as +// a column (an AstContext PrecomputeInstruction built via createCudfExpression +// at AstExpressionUtils.h). That precompute must receive the session-timezone +// context, or the GPU evaluates hour() in UTC and the predicate selects the +// wrong rows -- a silent wrong result, not an error. The projection tests above +// cover the projection path; this covers the filter/AST-precompute path. The +// two-row input selects different rows under UTC and Asia/Kolkata, so a GPU +// that ignores the session timezone returns the UTC-hour-5 row while CPU +// returns the local-hour-5 row. +TEST_F(TimezoneExtractionTest, filterPrecomputeHonorsSessionTimezone) { + auto input = makeRowVector( + {"ts"}, + {makeFlatVector( + {Timestamp(kJan2021MidnightUtc, 0), Timestamp(kJan2021At0500Utc, 0)}, + TIMESTAMP())}); + assertFilterGpuMatchesCpu(input, "hour(ts) = 5", "ts", kKolkata); +} + +TEST_F(TimezoneExtractionTest, dayOfWeekHonorsSessionTimezone) { + // Expect local 2020-12-31 (Thursday); GPU currently returns UTC 2021-01-01 + // (Friday). + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "day_of_week(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, dowHonorsSessionTimezone) { + // dow is an alias of day_of_week. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "dow(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, dayOfYearHonorsSessionTimezone) { + // Expect local 2020-12-31 (day 366 of leap year 2020); GPU currently returns + // UTC 2021-01-01 (day 1). + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "day_of_year(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, doyHonorsSessionTimezone) { + // doy is an alias of day_of_year. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "doy(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, weekHonorsSessionTimezone) { + // Expect local 2021-01-03 (Sunday), ISO week 53; GPU currently returns UTC + // 2021-01-04 (Monday), ISO week 1. + assertGpuMatchesCpu( + timestampInput(kJan2021MondayUtc), "week(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, weekOfYearHonorsSessionTimezone) { + // week_of_year is an alias of week. + assertGpuMatchesCpu( + timestampInput(kJan2021MondayUtc), "week_of_year(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, yearOfWeekHonorsSessionTimezone) { + // Expect local 2021-01-03 to belong to week-year 2020; GPU currently returns + // UTC 2021-01-04, week-year 2021. + assertGpuMatchesCpu( + timestampInput(kJan2021MondayUtc), "year_of_week(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, yowHonorsSessionTimezone) { + // yow is an alias of year_of_week. + assertGpuMatchesCpu( + timestampInput(kJan2021MondayUtc), "yow(ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, minuteHonorsHalfHourOffsetZone) { + // Asia/Kolkata is UTC+5:30. Expect local minute 30; GPU currently returns UTC + // minute 0. + assertGpuMatchesCpu( + timestampInput(kJan2021MidnightUtc), "minute(ts)", kKolkata); +} + +// second and millisecond cannot diverge: every IANA timezone offset is a whole +// number of minutes, so sub-minute fields are identical in UTC and in any +// session timezone. On CPU, second/millisecond are additionally computed +// without applying the session timezone at all (getDateTime(timestamp, +// nullptr)). Assert that GPU still matches CPU under a non-UTC timezone to +// document the boundary of the gap. +TEST_F(TimezoneExtractionTest, secondUnaffectedByTimezone) { + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc, 123'000'000), + "second(ts)", + kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, millisecondUnaffectedByTimezone) { + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc, 123'000'000), + "millisecond(ts)", + kLosAngeles); +} + +// Control: under UTC the GPU (which always computes in UTC) already matches the +// CPU for every extraction function. This passes today and guards the harness: +// a failure here would point to an extraction bug unrelated to the session +// timezone, isolating it from the timezone-driven failures above. +TEST_F(TimezoneExtractionTest, allComponentsMatchUnderUtc) { + auto boundary = timestampInput(kJan2021At0200Utc, 123'000'000); + auto monday = timestampInput(kJan2021MondayUtc); + + for (const auto& projection : + {"year(ts)", + "month(ts)", + "day(ts)", + "quarter(ts)", + "hour(ts)", + "minute(ts)", + "second(ts)", + "millisecond(ts)", + "day_of_week(ts)", + "dow(ts)", + "day_of_year(ts)", + "doy(ts)"}) { + SCOPED_TRACE(projection); + assertGpuMatchesCpu(boundary, projection, "UTC"); + } + for (const auto& projection : + {"week(ts)", "week_of_year(ts)", "year_of_week(ts)", "yow(ts)"}) { + SCOPED_TRACE(projection); + assertGpuMatchesCpu(monday, projection, "UTC"); + } +} + +// date_trunc(timestamp) must truncate on the session-local wall clock, matching +// the CPU truncateTimestamp, rather than on the raw UTC epoch. day-and-above +// and hour differ from UTC under a non-UTC session: these fail while the GPU +// truncates in UTC and pass once DateTruncFunction converts to local, +// truncates, then converts back. +TEST_F(TimezoneExtractionTest, dateTruncDayHonorsSessionTimezone) { + // 2021-01-01 02:00 UTC is 2020-12-31 18:00 in America/Los_Angeles, so the + // truncated day is the previous local day, not the UTC day. + assertGpuMatchesCpu( + timestampInput(kJan2021At0200Utc), "date_trunc('day', ts)", kLosAngeles); +} + +TEST_F( + TimezoneExtractionTest, + dateTruncWeekMonthQuarterYearHonorSessionTimezone) { + auto boundary = timestampInput(kJan2021At0200Utc); + for (const auto& projection : + {"date_trunc('month', ts)", + "date_trunc('quarter', ts)", + "date_trunc('year', ts)"}) { + SCOPED_TRACE(projection); + assertGpuMatchesCpu(boundary, projection, kLosAngeles); + } + // 2021-01-04 02:00 UTC is a Monday; 2021-01-03 18:00 in America/Los_Angeles + // is the preceding Sunday, so the local week starts the prior Monday. + assertGpuMatchesCpu( + timestampInput(kJan2021MondayUtc), "date_trunc('week', ts)", kLosAngeles); +} + +TEST_F(TimezoneExtractionTest, dateTruncHourHonorsHalfHourOffsetZone) { + // Asia/Kolkata is +05:30, so the local hour boundary is offset by 30 minutes + // from the UTC hour boundary; the DST-safe UTC delta must reproduce it. + assertGpuMatchesCpu( + timestampInput(kJan2021MidnightUtc), "date_trunc('hour', ts)", kKolkata); +} + +TEST_F(TimezoneExtractionTest, dateTruncSecondMinuteUnaffectedByTimezone) { + // second/minute truncate the UTC epoch directly (offsets are whole minutes), + // so they match CPU under any session timezone. + auto input = timestampInput(kJan2021At0200Utc, 123'000'000); + assertGpuMatchesCpu(input, "date_trunc('second', ts)", kKolkata); + assertGpuMatchesCpu(input, "date_trunc('minute', ts)", kKolkata); +} + +} // namespace diff --git a/velox/experimental/cudf/tests/TimezoneFunctionTest.cpp b/velox/experimental/cudf/tests/TimezoneFunctionTest.cpp new file mode 100644 index 00000000000..6d952e00c06 --- /dev/null +++ b/velox/experimental/cudf/tests/TimezoneFunctionTest.cpp @@ -0,0 +1,861 @@ +/* + * 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. + */ + +// Class B correctness tests: GPU (cuDF) evaluation of the TIMESTAMP WITH TIME +// ZONE function family must match CPU. These reproduce the missing timezone +// support by failing while it is absent. +// +// None of these functions are implemented on the GPU path today. Each works on +// CPU (it is registered by registerAllScalarFunctions), but forcing GPU +// evaluation currently throws from FunctionExpression::eval: +// +// "Unsupported expression for recursive evaluation: " +// (velox/experimental/cudf/expression/ExpressionEvaluator.cpp) +// +// This is true both for functions that produce a TIMESTAMP WITH TIME ZONE from +// plain double/varchar inputs and for functions that consume a TIMESTAMP WITH +// TIME ZONE column. The input conversion (Velox -> cuDF) is *not* the failure +// site: cuDF has no TIMESTAMP WITH TIME ZONE type, but a TIMESTAMP WITH TIME +// ZONE column is carried as its physical BIGINT and round-trips without error +// (see timestampWithTimeZoneColumnPreservedThroughGpu). The failure surfaces +// only when an unsupported function is evaluated. +// +// CudfFunctionBaseTest::evaluate forces GPU execution (it does not consult +// allowCpuFallback), and assertExpressionMatchesCpu compares that GPU result to +// the CPU result -- so until a function is implemented the GPU evaluation +// throws and the test fails. +// +// IMPORTANT: these are test-driven-development tests for the *target* behavior, +// so they FAIL today and pass once the gap is closed. Each asserts the GPU +// result equals CPU; until the function (and a GPU TIMESTAMP WITH TIME ZONE +// representation) is implemented, GPU evaluation throws and the test is red. A +// red test here means the function is still unsupported on GPU; do not "fix" a +// failure by weakening the assertion. The one exception is +// timestampWithTimeZoneColumnPreservedThroughGpu, which passes today: a plain +// passthrough never touches timezone semantics and the physical BIGINT +// round-trips losslessly, so it serves as a baseline proving the function tests +// fail in the function and not in the column conversion. +// +// These tests require a GPU and are labeled cuda_driver; they will not run in a +// CPU-only environment. +// +// Note: Presto's with_timezone is intentionally not covered here. Velox does +// not register a with_timezone scalar function, so it cannot be compiled on CPU +// and is out of scope for a GPU-vs-CPU gap. + +#include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/experimental/cudf/tests/CudfFunctionBaseTest.h" + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h" +#include "velox/functions/prestosql/registration/RegistrationFunctions.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" +#include "velox/parse/TypeResolver.h" +#include "velox/type/tz/TimeZoneMap.h" + +#include + +using namespace facebook::velox; +using namespace facebook::velox::cudf_velox; + +namespace { + +class TimezoneFunctionTest : public cudf_velox::CudfFunctionBaseTest { + protected: + static void SetUpTestCase() { + parse::registerTypeResolver(); + functions::prestosql::registerAllScalarFunctions(); + aggregate::prestosql::registerAllAggregateFunctions(); + memory::MemoryManager::testingSetInstance(memory::MemoryManager::Options{}); + cudf_velox::registerCudf(); + } + + static void TearDownTestCase() { + cudf_velox::unregisterCudf(); + } + + // Builds a single-row TIMESTAMP WITH TIME ZONE input column named c0, packing + // the UTC millis with the given zone's key (same layout as + // TimestampWithTimeZoneType: upper 52 bits millis, lower 12 bits zone key). + RowVectorPtr timestampWithTimeZoneInput(int64_t millisUtc, const char* zone) { + auto zoneId = tz::getTimeZoneID(zone); + return makeRowVector({makeFlatVector( + {pack(millisUtc, zoneId)}, TIMESTAMP_WITH_TIME_ZONE())}); + } + + // Builds a two-row TIMESTAMP WITH TIME ZONE column [value, NULL] in the given + // zone, to check that a NULL row propagates as NULL through the GPU path. + RowVectorPtr timestampWithTimeZoneAndNullInput( + int64_t millisUtc, + const char* zone) { + auto zoneId = tz::getTimeZoneID(zone); + return makeRowVector({makeNullableFlatVector( + {pack(millisUtc, zoneId), std::nullopt}, TIMESTAMP_WITH_TIME_ZONE())}); + } + + // Builds a two-row, entirely-NULL TIMESTAMP WITH TIME ZONE column. + RowVectorPtr allNullTimestampWithTimeZoneInput() { + return makeRowVector({makeNullableFlatVector( + {std::nullopt, std::nullopt}, TIMESTAMP_WITH_TIME_ZONE())}); + } + + // Builds a two-row TIMESTAMP WITH TIME ZONE column whose rows carry different + // zone keys, to exercise per-row (non-uniform) zone handling. + RowVectorPtr twoZoneTimestampWithTimeZoneInput( + int64_t millisUtcA, + const char* zoneA, + int64_t millisUtcB, + const char* zoneB) { + return makeRowVector({makeFlatVector( + {pack(millisUtcA, tz::getTimeZoneID(zoneA)), + pack(millisUtcB, tz::getTimeZoneID(zoneB))}, + TIMESTAMP_WITH_TIME_ZONE())}); + } + + // Builds a single-row double input column named c0. + RowVectorPtr doubleInput(double value) { + return makeRowVector({makeFlatVector({value})}); + } + + // Builds a single-row varchar input column named c0. + RowVectorPtr varcharInput(const std::string& value) { + return makeRowVector({makeFlatVector({value})}); + } + + // Asserts the expression evaluates to the same result on GPU (forced by + // CudfFunctionBaseTest::evaluate) and CPU. The input's own type is the + // projection's row type. + void assertMatchesCpu(const std::string& expr, const RowVectorPtr& input) { + assertExpressionMatchesCpu(expr, input, asRowType(input->type())); + } + + // Sets the session timezone for subsequent evaluate() calls, mirroring + // DateTimeFunctionsTest::setQueryTimeZone, so a test can exercise the + // session-timezone path the harness otherwise runs with an empty session. + void setSessionTimezone(const std::string& zone) { + queryCtx_->testingOverrideConfigUnsafe({ + {core::QueryConfig::kSessionTimezone, zone}, + {core::QueryConfig::kAdjustTimestampToTimezone, "true"}, + }); + } + + // Sets the session start time (consumed by now()/current_timestamp) and the + // session timezone, with adjust-to-session-timezone on, for subsequent + // evaluate() calls. testingOverrideConfigUnsafe replaces the whole config, so + // all three keys are set together. + void setSessionStartTimeAndTimeZone( + int64_t startTimeMs, + const std::string& zone) { + queryCtx_->testingOverrideConfigUnsafe({ + {core::QueryConfig::kSessionStartTime, std::to_string(startTimeMs)}, + {core::QueryConfig::kSessionTimezone, zone}, + {core::QueryConfig::kAdjustTimestampToTimezone, "true"}, + }); + } +}; + +// A TIMESTAMP WITH TIME ZONE column projected unchanged must round-trip through +// the GPU as TIMESTAMP WITH TIME ZONE and match CPU. cuDF has no native +// TIMESTAMP WITH TIME ZONE type and carries the column as its physical BIGINT; +// this asserts that representation preserves the packed millis+zone so the +// passthrough result is indistinguishable from CPU. It isolates the column +// conversion from the function evaluation exercised below. +TEST_F(TimezoneFunctionTest, timestampWithTimeZoneColumnPreservedThroughGpu) { + auto input = + timestampWithTimeZoneInput(1'609'466'400'000, "America/Los_Angeles"); + assertMatchesCpu("c0", input); +} + +// Functions that consume a TIMESTAMP WITH TIME ZONE column. The column converts +// to cuDF (as its physical BIGINT) without error; the function is the work the +// GPU must learn to do. + +TEST_F(TimezoneFunctionTest, toUnixtimeFromTimestampWithTimeZone) { + // to_unixtime(timestamp with time zone) -> double. + auto input = + timestampWithTimeZoneInput(1'609'466'400'000, "America/Los_Angeles"); + assertMatchesCpu("to_unixtime(c0)", input); +} + +// Coverage: a pre-1970 (negative-millis) instant exercises the arithmetic right +// shift in unpackMillis, which differs from a logical shift only for negative +// packed values -- every other test here uses a positive 2021 instant. +// to_unixtime recovers the seconds and to_iso8601 unpacks then renders, so both +// must match CPU for the negative instant 1938-04-24T17:33:20 UTC. +TEST_F(TimezoneFunctionTest, toUnixtimePre1970Instant) { + auto input = + timestampWithTimeZoneInput(-1'000'000'000'000, "America/Los_Angeles"); + assertMatchesCpu("to_unixtime(c0)", input); +} + +TEST_F(TimezoneFunctionTest, toIso8601Pre1970Instant) { + auto input = + timestampWithTimeZoneInput(-1'000'000'000'000, "America/Los_Angeles"); + assertMatchesCpu("to_iso8601(c0)", input); +} + +TEST_F(TimezoneFunctionTest, atTimezone) { + // at_timezone(timestamp with time zone, varchar) -> timestamp with time zone. + auto input = + timestampWithTimeZoneInput(1'609'466'400'000, "America/Los_Angeles"); + assertMatchesCpu("at_timezone(c0, 'America/New_York')", input); +} + +TEST_F(TimezoneFunctionTest, timezoneHour) { + // timezone_hour(timestamp with time zone) -> bigint. + auto input = + timestampWithTimeZoneInput(1'609'466'400'000, "America/Los_Angeles"); + assertMatchesCpu("timezone_hour(c0)", input); +} + +TEST_F(TimezoneFunctionTest, timezoneMinute) { + // timezone_minute(timestamp with time zone) -> bigint. + auto input = timestampWithTimeZoneInput(1'609'466'400'000, "Asia/Kolkata"); + assertMatchesCpu("timezone_minute(c0)", input); +} + +// Reproducers: timezone_hour/timezone_minute must return NULL for a NULL row, +// matching CPU (a plain call() -> NULL for NULL). The GPU offset primitive +// (utcOffsetSeconds) builds an all-valid column via make_column_from_scalar / +// gather and never re-applies the input mask (TimezoneConversion.h documents +// the all-valid contract), so the field functions' scalar DIV/MOD yield 0 +// instead of NULL. Red until the input validity is carried onto the offset +// column. The single-row tests above use non-null inputs and so never exercise +// this. +TEST_F(TimezoneFunctionTest, timezoneHourPropagatesNull) { + auto input = + timestampWithTimeZoneAndNullInput(1'609'466'400'000, "Asia/Kolkata"); + assertMatchesCpu("timezone_hour(c0)", input); +} + +TEST_F(TimezoneFunctionTest, timezoneMinutePropagatesNull) { + auto input = + timestampWithTimeZoneAndNullInput(1'609'466'400'000, "Asia/Kolkata"); + assertMatchesCpu("timezone_minute(c0)", input); +} + +// Reproducers: a TSWTZ column mixing zone keys must be handled per row (CPU +// unpacks each row's own key). The GPU's uniformZoneKey VELOX_USER_CHECK-fails +// on mixed zones (the "one zone per column" limitation). Red until the per-row +// offset path lands. +TEST_F(TimezoneFunctionTest, timezoneHourMixedZones) { + auto input = twoZoneTimestampWithTimeZoneInput( + 1'609'466'400'000, + "America/Los_Angeles", + 1'609'466'400'000, + "Asia/Kolkata"); + assertMatchesCpu("timezone_hour(c0)", input); +} + +TEST_F(TimezoneFunctionTest, timezoneMinuteMixedZones) { + auto input = twoZoneTimestampWithTimeZoneInput( + 1'609'466'400'000, + "America/Los_Angeles", + 1'609'466'400'000, + "Asia/Kolkata"); + assertMatchesCpu("timezone_minute(c0)", input); +} + +// Mixed zones plus a null row: the null must stay null through the per-row path. +TEST_F(TimezoneFunctionTest, timezoneHourMixedZonesWithNull) { + auto input = makeRowVector({makeNullableFlatVector( + {pack(1'609'466'400'000, tz::getTimeZoneID("America/Los_Angeles")), + pack(1'609'466'400'000, tz::getTimeZoneID("Asia/Kolkata")), + std::nullopt}, + TIMESTAMP_WITH_TIME_ZONE())}); + assertMatchesCpu("timezone_hour(c0)", input); +} + +// Mixed zones at a DST-varying instant: the per-row offset must be computed for +// each row's own instant, not a uniform one. 2021-07-01T02:00:00Z puts +// America/Los_Angeles in PDT (-07:00, not the -08:00 PST the January cases use) +// while Asia/Kolkata is fixed at +05:30. The existing single-zone timezone_hour +// test uses a January (PST) instant, so this is the only DST-active per-row +// case. +TEST_F(TimezoneFunctionTest, timezoneHourMixedZonesDst) { + auto input = twoZoneTimestampWithTimeZoneInput( + 1'625'104'800'000, + "America/Los_Angeles", + 1'625'104'800'000, + "Asia/Kolkata"); + assertMatchesCpu("timezone_hour(c0)", input); +} + +TEST_F(TimezoneFunctionTest, toIso8601FromTimestampWithTimeZone) { + // to_iso8601(timestamp with time zone) -> varchar. + auto input = + timestampWithTimeZoneInput(1'609'466'400'000, "America/Los_Angeles"); + assertMatchesCpu("to_iso8601(c0)", input); +} + +// Reproducer for the zero-offset divergence: to_iso8601 of a UTC/GMT instant +// must render a trailing 'Z', matching CPU (ToISO8601Function passes +// zeroOffsetText="Z"). The GPU's formatOffsetStrings has no zero-offset branch +// and emits '+00:00'. Red until the 'Z' branch is added. (The only other +// to_iso8601 test uses a non-zero offset, so it does not exercise this.) +TEST_F(TimezoneFunctionTest, toIso8601RendersZForZeroOffset) { + auto input = timestampWithTimeZoneInput(1'609'466'400'000, "UTC"); + assertMatchesCpu("to_iso8601(c0)", input); +} + +// to_iso8601 over mixed zones: each row renders its own offset (LA -08:00 vs +// Kolkata +05:30 on the same UTC instant). Red until the per-row offset lands. +TEST_F(TimezoneFunctionTest, toIso8601MixedZones) { + auto input = twoZoneTimestampWithTimeZoneInput( + 1'609'466'400'000, + "America/Los_Angeles", + 1'609'466'400'000, + "Asia/Kolkata"); + assertMatchesCpu("to_iso8601(c0)", input); +} + +// Contract/regression test: an entirely-NULL TIMESTAMP WITH TIME ZONE column +// must yield an all-NULL result like CPU. uniformZoneKey reduces min/max over +// the (all-null) zone-key column; reduce excludes nulls, so its scalars come +// back invalid and value() would be a meaningless device read (UB) before +// VELOX_USER_CHECK_EQ(lo, hi). uniformZoneKey guards null_count() == size() and +// defaults to GMT (key 0), as the empty-column path does. This is not a +// differential RED for the UB -- the bad read happens to yield 0/GMT in this +// environment, so the output is already correct -- so it instead pins the +// all-null -> all-null contract and guards against the guard's removal. +TEST_F(TimezoneFunctionTest, toIso8601AllNullColumn) { + assertMatchesCpu("to_iso8601(c0)", allNullTimestampWithTimeZoneInput()); +} + +TEST_F(TimezoneFunctionTest, formatDatetimeOfTimestampWithTimeZone) { + // format_datetime(timestamp with time zone, varchar) -> varchar. + auto input = + timestampWithTimeZoneInput(1'609'466'400'000, "America/Los_Angeles"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss ZZ')", input); +} + +// format_datetime over mixed zones: the local wall clock and the numeric offset +// token ('ZZ' -> "+HH:MM") are both per-row. This exercises localAndOffset +// through the per-row offset path (LA -08:00 vs Kolkata +05:30 on the same UTC +// instant give different local times and different rendered offsets). +TEST_F(TimezoneFunctionTest, formatDatetimeMixedZones) { + auto input = twoZoneTimestampWithTimeZoneInput( + 1'609'466'400'000, + "America/Los_Angeles", + 1'609'466'400'000, + "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss ZZ')", input); +} + +// Reproducers for the Joda zone-token divergences. CPU (DateTimeFormatter) +// distinguishes the run length and letter; the GPU collapses Z/z into one flag +// and always emits '+HH:MM'. Only the (correct) ZZ case is covered above. Each +// is red until jodaToStrftime threads the run length and letter. + +// Single 'Z' renders the offset WITHOUT a colon (e.g. +0530) on CPU +// (appendTimezoneOffset, includeColon=false); the GPU emits +05:30. +TEST_F(TimezoneFunctionTest, formatDatetimeSingleZNoColon) { + auto input = timestampWithTimeZoneInput(1'609'466'400'000, "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss Z')", input); +} + +// 'ZZZ' (3+ repeats) renders the zone id (Asia/Kolkata) on CPU; the GPU emits +// the numeric offset. +TEST_F(TimezoneFunctionTest, formatDatetimeZoneIdToken) { + auto input = timestampWithTimeZoneInput(1'609'466'400'000, "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss ZZZ')", input); +} + +// format_datetime zone-id token ('ZZZ' -> zone name) over mixed zones: each row +// renders its own zone name (America/Los_Angeles vs Asia/Kolkata) via +// perRowZoneName. formatDatetimeZoneIdToken above covers only a single zone; +// this pins the per-row name path the owner scoped into this PR. Red until +// perRowZoneName replaces the uniformZoneKey single-name render. +TEST_F(TimezoneFunctionTest, formatDatetimeZoneIdMixedZones) { + auto input = twoZoneTimestampWithTimeZoneInput( + 1'609'466'400'000, + "America/Los_Angeles", + 1'609'466'400'000, + "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss ZZZ')", input); +} + +// Lowercase 'z' is a distinct Joda specifier (zone abbreviation/name, e.g. +// IST). It is DST- and instant-dependent, so the GPU cannot render it on +// device; it rejects the token with VELOX_NYI rather than silently emit a wrong +// (numeric offset) result. Asserting the guard pins the scoped limitation. +TEST_F(TimezoneFunctionTest, formatDatetimeZoneNameTokenUnsupportedOnGpu) { + auto input = timestampWithTimeZoneInput(1'609'466'400'000, "Asia/Kolkata"); + auto exprSet = compileExpression( + "format_datetime(c0, 'yyyy-MM-dd HH:mm:ss z')", asRowType(input->type())); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Reproducers for the Joda fractional-second run length. CPU +// (formatFractionOfSecond) renders exactly digits: a single 'S' +// is 1 digit, 'SSSSSS' is 6. The GPU's jodaToStrftime maps any 'S' run to +// "%3f" (3 digits), so single-'S' and 6-'S' diverge while 'SSS' happens to +// match. Red until the run length feeds the "%f" width. The 123 ms +// sub-second instant makes the fractional digits observable. +TEST_F(TimezoneFunctionTest, formatDatetimeFractionSingleDigit) { + auto input = timestampWithTimeZoneInput(1'609'466'400'123, "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss.S')", input); +} + +// 'SSS' -> 3 digits; matches the GPU's current %3f (control case, stays green). +TEST_F(TimezoneFunctionTest, formatDatetimeFractionMillis) { + auto input = timestampWithTimeZoneInput(1'609'466'400'123, "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss.SSS')", input); +} + +// 'SSSSSS' -> 6 digits; the millisecond value is right-padded with zeros. +TEST_F(TimezoneFunctionTest, formatDatetimeFractionMicros) { + auto input = timestampWithTimeZoneInput(1'609'466'400'123, "Asia/Kolkata"); + assertMatchesCpu("format_datetime(c0, 'yyyy-MM-dd HH:mm:ss.SSSSSS')", input); +} + +// Functions that produce a TIMESTAMP WITH TIME ZONE from plain inputs. The +// inputs convert to cuDF fine; the function is the work the GPU must learn. + +TEST_F(TimezoneFunctionTest, fromUnixtimeWithZoneName) { + // from_unixtime(double, varchar) -> timestamp with time zone. + assertMatchesCpu( + "from_unixtime(c0, 'America/Los_Angeles')", doubleInput(1'609'466'400.0)); +} + +TEST_F(TimezoneFunctionTest, fromUnixtimeWithHoursMinutes) { + // from_unixtime(double, bigint, bigint) -> timestamp with time zone. + assertMatchesCpu("from_unixtime(c0, 7, 30)", doubleInput(1'609'466'400.0)); +} + +// Reproducer: from_unixtime(double, bigint, bigint) computes the fixed offset as +// hours*60 + minutes. INT64_MAX hours overflows that int64 product. CPU +// (FromUnixtimeFunction) uses checkedMultiply/checkedPlus and throws; the GPU +// registration multiplies unchecked, then casts to int32 -- on this platform the +// UB wraps to -60, an in-range offset tz::getTimeZoneID happily accepts. Red +// until the GPU mirrors CPU's checked arithmetic. compileExpression succeeds on +// both (the CPU arithmetic error is a user error captured in initialize() and +// re-thrown at eval); both throws carry "overflow". +TEST_F(TimezoneFunctionTest, fromUnixtimeHoursMinutesOverflowRejectedLikeCpu) { + auto input = doubleInput(0.0); + auto exprSet = compileExpression( + "from_unixtime(c0, 9223372036854775807, 0)", asRowType(input->type())); + VELOX_ASSERT_THROW( + functions::test::FunctionBaseTest::evaluate(*exprSet, input), "overflow"); + VELOX_ASSERT_THROW(evaluate(*exprSet, input), "overflow"); +} + +// Reproducer: from_unixtime of an out-of-range instant must throw to match CPU. +// CPU pack() VELOX_USER_CHECKs the millis range and throws an overflow error; +// the CPU suite asserts from_unixtime(2251799813685.248, 'GMT') throws. The GPU +// shifts millis << 12 with no guard and silently overflows into the zone-key +// bits. Red until the range/NaN check is added. +TEST_F(TimezoneFunctionTest, fromUnixtimeOverflowRejectedLikeCpu) { + auto input = doubleInput(2'251'799'813'685.248); + auto exprSet = + compileExpression("from_unixtime(c0, 'GMT')", asRowType(input->type())); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Reproducer: from_unixtime(NaN) must map to the epoch like CPU, which returns +// pack(0, zone) for a NaN unixtime, rather than reading a meaningless value out +// of a float->int cast of NaN. Red until NaN is mapped to 0 before packing. +TEST_F(TimezoneFunctionTest, fromUnixtimeNanMapsToEpochLikeCpu) { + auto input = doubleInput(std::numeric_limits::quiet_NaN()); + assertMatchesCpu("from_unixtime(c0, 'GMT')", input); +} + +// Reproducer: from_unixtime(+/-Inf) must throw to match CPU. CPU saturates the +// millis to int64 min/max, which pack()'s range check then rejects as overflow. +// The GPU must throw too rather than rely on float->int cast behavior for Inf. +TEST_F(TimezoneFunctionTest, fromUnixtimeInfinityRejectedLikeCpu) { + auto input = doubleInput(std::numeric_limits::infinity()); + auto exprSet = + compileExpression("from_unixtime(c0, 'GMT')", asRowType(input->type())); + EXPECT_ANY_THROW( + functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Reproducer for the two-overload rounding split. The (double, hours, minutes) +// overload rounds via floor-seconds + a separate fractional llround (CPU's +// no-zone fromUnixtime), differing from the varchar overload's llround(x*1000) +// by up to 1 ms on negative-fractional input. For -0.0005 s the hours/minutes +// overload yields 0 ms while the varchar overload yields -1 ms. The GPU uses +// the varchar rounding for both, so the hours/minutes case is red. +TEST_F( + TimezoneFunctionTest, + fromUnixtimeHoursMinutesNegativeFractionalRounding) { + auto input = doubleInput(-0.0005); + assertMatchesCpu("from_unixtime(c0, 0, 0)", input); +} + +// Control: the varchar overload's llround(x*1000) already matches CPU for the +// same negative-fractional input (-0.0005 -> -1 ms), so this stays green. +TEST_F(TimezoneFunctionTest, fromUnixtimeVarcharNegativeFractionalRounding) { + auto input = doubleInput(-0.0005); + assertMatchesCpu("from_unixtime(c0, 'GMT')", input); +} + +TEST_F(TimezoneFunctionTest, parseDatetime) { + // parse_datetime(varchar, varchar) -> timestamp with time zone. + assertMatchesCpu( + "parse_datetime(c0, 'yyyy-MM-dd HH:mm:ss')", + varcharInput("2021-01-01 02:00:00")); +} + +// When the Joda format carries a colon offset token (ZZ), CPU folds the offset +// into the UTC instant AND packs the parsed fixed-offset zone key, so +// timezone_hour reports -9 and to_iso8601 prints -09:00. GPU currently packs +// GMT (timezone_hour = 0, to_iso8601 = Z). Compare through projections that read +// the zone key, since assertMatchesCpu on the TSWTZ value alone ignores it. +TEST_F(TimezoneFunctionTest, parseDatetimePreservesParsedOffset) { + auto input = varcharInput("2021-01-01 02:00:00 -09:00"); + assertMatchesCpu( + "timezone_hour(parse_datetime(c0, 'yyyy-MM-dd HH:mm:ss ZZ'))", input); + assertMatchesCpu( + "to_iso8601(parse_datetime(c0, 'yyyy-MM-dd HH:mm:ss ZZ'))", input); +} + +// Same, for the no-colon offset token (Z) matching -0900. +TEST_F(TimezoneFunctionTest, parseDatetimeNoColonOffset) { + auto input = varcharInput("2021-01-01 02:00:00 -0900"); + assertMatchesCpu( + "timezone_hour(parse_datetime(c0, 'yyyy-MM-dd HH:mm:ss Z'))", input); + assertMatchesCpu( + "to_iso8601(parse_datetime(c0, 'yyyy-MM-dd HH:mm:ss Z'))", input); +} + +TEST_F(TimezoneFunctionTest, fromIso8601Timestamp) { + // from_iso8601_timestamp(varchar) -> timestamp with time zone. + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T02:00:00+05:30")); +} + +// Reproducers: from_iso8601_timestamp must accept the ISO8601 shapes CPU does +// (see DateTimeFunctionsTest.fromIso8601Timestamp). The GPU's rigid regex +// requires a full yyyy-MM-ddTHH:mm:ss with a colon offset, so it rejects short +// forms (-> NULL), discards sub-second digits, rejects hours-only offsets, and +// loses the sign of offsets in (-1h, 0). Inputs without an embedded offset are +// interpreted as GMT under the default session, matching CPU. Each is red until +// the GPU parser matches CPU. + +// Date-only: CPU -> midnight GMT; GPU regex needs a time component -> NULL. +TEST_F(TimezoneFunctionTest, fromIso8601DateOnly) { + assertMatchesCpu("from_iso8601_timestamp(c0)", varcharInput("2021-01-01")); +} + +// Minute precision (no seconds): CPU accepts; GPU regex needs seconds -> NULL. +TEST_F(TimezoneFunctionTest, fromIso8601MinutePrecision) { + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-02T11:38")); +} + +// Sub-second digits: CPU preserves .123; GPU discards them (parses to seconds). +TEST_F(TimezoneFunctionTest, fromIso8601FractionalSeconds) { + assertMatchesCpu( + "from_iso8601_timestamp(c0)", + varcharInput("2021-01-01T02:00:00.123+05:30")); +} + +// Hours-only offset: CPU expands +05 -> +05:00; GPU requires minutes -> NULL. +TEST_F(TimezoneFunctionTest, fromIso8601HoursOnlyOffset) { + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T02:00:00+05")); +} + +// Offset in (-1h, 0): CPU keeps the sign (-00:30); GPU reads -00 as 0 and +// yields +30 -- wrong instant and wrong zone key. +TEST_F(TimezoneFunctionTest, fromIso8601NegativeHalfHourOffset) { + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T02:00:00-00:30")); +} + +// Reproducer: an offset outside +/-14h must be rejected like CPU rather than +// silently corrupt the packed value. CPU normalizes "+99:00" to an unknown zone +// name and throws ("Unknown timezone value"); the +/-840-minute bound is the +// same one tz::getTimeZoneID enforces. The GPU parser has no bound -- +99:00 -> +// 5940 minutes -> zone key 6780, which overflows the 12-bit zone field and +// corrupts the packed millis (the key is not masked with kTimezoneMask). Red +// until the parsed offset magnitude is bounded with a user error. +TEST_F(TimezoneFunctionTest, fromIso8601OffsetOutOfRangeRejectedLikeCpu) { + auto input = varcharInput("2021-01-01T02:00:00+99:00"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + // CPU rejects the out-of-range offset; confirm parity is "both throw". + EXPECT_ANY_THROW( + functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Reproducer: an offset-less ISO string is interpreted in the session timezone +// on CPU (the wall clock is that zone's local time, and the packed zone key is +// the session zone), not GMT. Asia/Kolkata has a fixed +05:30 offset (no DST), +// so the conversion is exact. The GPU treats offset-less input as GMT +// regardless of the session, so it produces both a wrong instant and a wrong +// zone key. Red until the session offset is applied. +TEST_F(TimezoneFunctionTest, fromIso8601OffsetlessUsesSessionZone) { + setSessionTimezone("Asia/Kolkata"); + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T02:00:00")); +} + +// Reproducer: the offset-less session-zone conversion must match CPU even for a +// DST zone whose offset depends on the instant. America/Los_Angeles springs +// forward on 2021-03-14 at 02:00 PST (-08:00) to 03:00 PDT (-07:00), i.e. at +// 10:00:00 UTC. The wall clock 2021-03-14T03:30:00 is a valid post-gap local +// time (PDT), so CPU resolves it to 2021-03-14T10:30:00 UTC (to_unixtime +// 1615717800). The GPU uses the local->UTC approximation, which keys the wall +// clock as if it were UTC: 03:30 UTC precedes the 10:00 UTC transition, so it +// reads the pre-gap offset (-08:00) and yields 2021-03-14T11:30:00 UTC +// (to_unixtime 1615721400) -- one hour late. Red until an inverse (local-keyed) +// transition lookup replaces the approximation. The fixed-offset Kolkata case +// above stays green because its offset does not vary with the instant. +TEST_F(TimezoneFunctionTest, fromIso8601OffsetlessSessionZoneDstTransition) { + setSessionTimezone("America/Los_Angeles"); + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-03-14T03:30:00")); +} + +// Reproducer: a wall clock inside the spring-forward gap is a nonexistent local +// time, so CPU's toGMT throws and from_iso8601_timestamp fails. America/ +// Los_Angeles springs forward on 2021-03-14 from 02:00 PST to 03:00 PDT, so +// local times in [02:00, 03:00) never occur; 02:30:00 is one of them. The GPU +// local->UTC approximation does plain arithmetic and never throws, so it +// silently returns an instant. Asserting both paths throw is red until the +// inverse (local-keyed) transition lookup flags the gap and fails like CPU. +TEST_F(TimezoneFunctionTest, fromIso8601OffsetlessSessionZoneGapThrows) { + setSessionTimezone("America/Los_Angeles"); + auto input = varcharInput("2021-03-14T02:30:00"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW( + functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Reproducer: a wall clock in a fall-back overlap is ambiguous, and CPU's toGMT +// resolves it to the earliest instant (TChoose::kEarliest). Australia/Sydney +// falls back on 2021-04-04 from 03:00 AEDT (+11:00) to 02:00 AEST (+10:00), so +// local times in [02:00, 03:00) occur twice; 02:30:00 is one. CPU keeps the +// earlier AEDT reading -- 2021-04-03T15:30:00 UTC. The GPU approximation keys +// the wall clock as UTC, which lands after the 2021-04-03T16:00 UTC transition +// and reads the later AEST offset, yielding 2021-04-03T16:30:00 UTC -- one hour +// late. Red until the inverse transition lookup keeps the pre-transition offset +// over the overlap, matching kEarliest. A western-hemisphere zone like +// Los_Angeles cannot exercise this: its negative offsets place the overlap +// window before the UTC transition, where the approximation already reads the +// earlier offset. +TEST_F( + TimezoneFunctionTest, + fromIso8601OffsetlessSessionZoneAmbiguousPicksEarliest) { + setSessionTimezone("Australia/Sydney"); + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-04-04T02:30:00")); +} + +// Control: an explicit numeric offset wins over the session zone on both paths, +// so this stays green and guards that the session change does not hijack rows +// that carry their own offset. +TEST_F(TimezoneFunctionTest, fromIso8601ExplicitOffsetIgnoresSessionZone) { + setSessionTimezone("Asia/Kolkata"); + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T02:00:00+09:00")); +} + +// Control: an explicit "Z" designator is GMT on both paths (distinct from an +// absent zone), so this stays green and guards that "Z" is not mistaken for an +// offset-less input and rerouted through the session zone. +TEST_F(TimezoneFunctionTest, fromIso8601ZuluIgnoresSessionZone) { + setSessionTimezone("Asia/Kolkata"); + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T02:00:00Z")); +} + +// Trailing 'T' with no time component: CPU treats it as the date at midnight; +// the current regex needs 2 digits after T -> NULL. (Oracle: DateTimeFunctions +// fromIso8601Timestamp accepts "1970-01-01T"/"1970-01T"/"1970T".) +TEST_F(TimezoneFunctionTest, fromIso8601TrailingT) { + assertMatchesCpu("from_iso8601_timestamp(c0)", varcharInput("2021-01-01T")); +} + +// Year-only and year-month: CPU -> start-of-period midnight GMT. +TEST_F(TimezoneFunctionTest, fromIso8601YearOnly) { + assertMatchesCpu("from_iso8601_timestamp(c0)", varcharInput("2021")); +} +TEST_F(TimezoneFunctionTest, fromIso8601YearMonth) { + assertMatchesCpu("from_iso8601_timestamp(c0)", varcharInput("2021-07")); +} + +// No time, explicit offset ("T+01:00", "T+14:00"): CPU applies the +// offset to the start-of-period wall clock. +TEST_F(TimezoneFunctionTest, fromIso8601DateThenOffset) { + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("2021-01-01T+01:00")); +} + +// Time-only ("Thh[:mm[:ss[.fff]]]" [offset]): CPU defaults the date to +// 1970-01-01; the date-anchored regex needs a leading year -> NULL today. +TEST_F(TimezoneFunctionTest, fromIso8601TimeOnly) { + assertMatchesCpu("from_iso8601_timestamp(c0)", varcharInput("T11:38:56")); +} +TEST_F(TimezoneFunctionTest, fromIso8601TimeOnlyHourOnly) { + assertMatchesCpu("from_iso8601_timestamp(c0)", varcharInput("T11")); +} +TEST_F(TimezoneFunctionTest, fromIso8601TimeOnlyWithOffset) { + assertMatchesCpu( + "from_iso8601_timestamp(c0)", varcharInput("T11:38:56.123-14:00")); +} + +// Malformed input: CPU (util::fromTimestampWithTimezoneString) throws; the GPU +// must not silently return NULL. Red until the throw block lands. +TEST_F(TimezoneFunctionTest, fromIso8601MalformedThrowsLikeCpu) { + auto input = varcharInput("not-a-timestamp"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Space separator: CPU rejects "yyyy-MM-dd HH:mm" (only 'T' is legal). GPU now +// rejects it too (regex tightened to 'T'; unmatched non-null row -> throw). +TEST_F(TimezoneFunctionTest, fromIso8601SpaceSeparatorThrowsLikeCpu) { + auto input = varcharInput("2021-01-02 11:38"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Empty string and a bare "T": both malformed on CPU. +TEST_F(TimezoneFunctionTest, fromIso8601EmptyStringThrowsLikeCpu) { + auto input = varcharInput(""); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} +TEST_F(TimezoneFunctionTest, fromIso8601BareTThrowsLikeCpu) { + auto input = varcharInput("T"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Well-formed shape but a nonexistent calendar date: CPU rejects both +// (isValidDate -- month > 12 / day past the month's length, leap year aware). +// The regex alone accepts the two-digit fields, and cudf::to_timestamps would +// silently normalize them (13 -> next year, Feb 30 -> March), so the GPU must +// detect the normalization and throw rather than return a wrong value. +TEST_F(TimezoneFunctionTest, fromIso8601InvalidMonthDayThrowsLikeCpu) { + auto input = varcharInput("2021-13-45"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} +TEST_F(TimezoneFunctionTest, fromIso8601InvalidFebruaryThrowsLikeCpu) { + auto input = varcharInput("2021-02-30"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_ANY_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); +} + +// Extreme-but-valid years (5-digit / signed): CPU parses them; the GPU +// (to_timestamps is int16, <=4-digit %Y) cannot, so it throws VELOX_NYI rather +// than returning NULL or a wrong value -- the query stops (owner decision +// 2026-07-16, no silent NULL). CPU is asserted to succeed to document the +// divergence. Literals are the max/min still in CPU's non-overflow range. +TEST_F(TimezoneFunctionTest, fromIso8601FiveDigitYearNyiOnGpu) { + auto input = varcharInput("73326-09-11T20:14:45.247"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_NO_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + VELOX_ASSERT_THROW(evaluate(*exprSet, input), "does not support years"); +} +TEST_F(TimezoneFunctionTest, fromIso8601NegativeYearNyiOnGpu) { + auto input = varcharInput("-69387-04-22T03:45:14.752"); + auto exprSet = + compileExpression("from_iso8601_timestamp(c0)", asRowType(input->type())); + EXPECT_NO_THROW(functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + VELOX_ASSERT_THROW(evaluate(*exprSet, input), "does not support years"); +} + +// Control: a genuine SQL NULL row stays NULL -- it must trip neither throw. +TEST_F(TimezoneFunctionTest, fromIso8601NullRowStaysNull) { + auto input = makeRowVector( + {makeNullableFlatVector({std::nullopt}, VARCHAR())}); + assertMatchesCpu("from_iso8601_timestamp(c0)", input); +} + +// now()/current_timestamp -> timestamp with time zone. now() is +// non-deterministic -- a live CPU now() and a separate GPU now() observe +// different instants -- so this cannot assert CPU == GPU against a live clock. +// Instead it pins the deterministic contract CPU's CurrentTimestampFunction +// implements: pack(sessionStartTimeMs, sessionZone). The GPU must emit a +// TIMESTAMP WITH TIME ZONE whose UTC millis are the session start time and +// whose zone key is the session zone. A dummy column sizes the batch. +TEST_F(TimezoneFunctionTest, nowUsesSessionStartTimeAndTimezone) { + constexpr int64_t kStartMs = 1'609'466'400'000; // 2021-01-01T02:00:00 UTC. + setSessionStartTimeAndTimeZone(kStartMs, "America/Los_Angeles"); + auto input = doubleInput(0.0); + auto exprSet = compileExpression("now()", asRowType(input->type())); + auto result = evaluate(*exprSet, input); + ASSERT_NE(result, nullptr); + ASSERT_EQ(result->size(), input->size()); + ASSERT_TRUE(isTimestampWithTimeZoneType(result->type())) + << "now() must produce TIMESTAMP WITH TIME ZONE, got " + << result->type()->toString(); + const auto packed = result->as>()->valueAt(0); + EXPECT_EQ(unpackMillisUtc(packed), kStartMs); + EXPECT_EQ(unpackZoneKeyId(packed), tz::getTimeZoneID("America/Los_Angeles")); +} + +// now()/current_timestamp must be rejected exactly when CPU rejects it. CPU's +// CurrentTimestampFunction throws "Timezone cannot be null" when +// getTimeZoneFromConfig returns null -- i.e. when +// adjust_timestamp_to_session_timezone is off, or the session timezone is +// empty. The GPU previously honored neither condition (defaulting the zone key +// to GMT) and silently produced a value where CPU failed. Assert both paths +// throw for the two rejection configs. The exception is captured at +// initialize() and re-thrown at eval time, so compileExpression succeeds and +// the throw surfaces from evaluate(). +TEST_F(TimezoneFunctionTest, nowWithoutAdjustedSessionTimezoneRejectedLikeCpu) { + auto input = doubleInput(0.0); + const auto rowType = asRowType(input->type()); + + // Adjust on but no session timezone -> rejected on both paths. + queryCtx_->testingOverrideConfigUnsafe({ + {core::QueryConfig::kAdjustTimestampToTimezone, "true"}, + }); + { + auto exprSet = compileExpression("now()", rowType); + EXPECT_ANY_THROW( + functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); + } + + // Session timezone set but adjust off -> rejected on both paths. + queryCtx_->testingOverrideConfigUnsafe({ + {core::QueryConfig::kSessionTimezone, "America/Los_Angeles"}, + {core::QueryConfig::kAdjustTimestampToTimezone, "false"}, + }); + { + auto exprSet = compileExpression("now()", rowType); + EXPECT_ANY_THROW( + functions::test::FunctionBaseTest::evaluate(*exprSet, input)); + EXPECT_ANY_THROW(evaluate(*exprSet, input)); + } +} + +} // namespace diff --git a/velox/experimental/cudf/tests/ToCudfSelectionTest.cpp b/velox/experimental/cudf/tests/ToCudfSelectionTest.cpp index f1b94097874..6db921d0632 100644 --- a/velox/experimental/cudf/tests/ToCudfSelectionTest.cpp +++ b/velox/experimental/cudf/tests/ToCudfSelectionTest.cpp @@ -24,7 +24,9 @@ #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/OperatorTestBase.h" #include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/type/Time.h" +#include "velox/type/tz/TimeZoneMap.h" namespace facebook::velox::exec::test { @@ -171,7 +173,10 @@ TEST_F(ToCudfSelectionTest, prestoDateAddVariableUnitFallsBack) { ASSERT_TRUE(wasDefaultFilterProjectUsed(task)); } -TEST_F(ToCudfSelectionTest, prestoDateAddTimestampFallsBack) { +TEST_F(ToCudfSelectionTest, prestoDateAddTimestampUsesCudf) { + // date_add(timestamp) is now evaluated on GPU (session-timezone aware), so + // the plan runs on cuDF rather than falling back -- both plain and under a + // session timezone. auto input = makeRowVector( {"amount", "event_ts"}, {makeFlatVector({1, 2, -1, 13}), @@ -189,12 +194,45 @@ TEST_F(ToCudfSelectionTest, prestoDateAddTimestampFallsBack) { std::shared_ptr task; AssertQueryBuilder(plan).config("cudf.enabled", true).countResults(task); + ASSERT_TRUE(wasCudfFilterProjectUsed(task)); + ASSERT_FALSE(wasDefaultFilterProjectUsed(task)); - ASSERT_FALSE(wasCudfFilterProjectUsed(task)); - ASSERT_TRUE(wasDefaultFilterProjectUsed(task)); + std::shared_ptr tzTask; + AssertQueryBuilder(plan) + .config("cudf.enabled", true) + .config(QueryConfig::kSessionTimezone, "America/Los_Angeles") + .config(QueryConfig::kAdjustTimestampToTimezone, "true") + .countResults(tzTask); + ASSERT_TRUE(wasCudfFilterProjectUsed(tzTask)); + ASSERT_FALSE(wasDefaultFilterProjectUsed(tzTask)); } -TEST_F(ToCudfSelectionTest, prestoDateTruncTimestampAdjustTimezoneFallsBack) { +TEST_F(ToCudfSelectionTest, prestoDateAddTimestampWithTimeZoneUsesCudf) { + // date_add(timestamp with time zone) is evaluated on GPU (per-row embedded + // zone), so the plan runs on cuDF rather than falling back. + auto input = makeRowVector( + {"amount", "c0"}, + {makeFlatVector({1, 2}), + makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles")), + pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + + auto plan = PlanBuilder() + .values({input}) + .project({"date_add('day', amount, c0) AS result"}) + .planNode(); + + std::shared_ptr task; + AssertQueryBuilder(plan).config("cudf.enabled", true).countResults(task); + + ASSERT_TRUE(wasCudfFilterProjectUsed(task)); + ASSERT_FALSE(wasDefaultFilterProjectUsed(task)); +} + +TEST_F(ToCudfSelectionTest, prestoDateTruncTimestampAdjustTimezoneUsesCudf) { + // date_trunc(timestamp) is timezone-aware on GPU, so it runs on cuDF even + // when adjust_timestamp_to_session_timezone is enabled. auto input = makeRowVector( {"event_ts"}, {makeFlatVector( @@ -212,8 +250,30 @@ TEST_F(ToCudfSelectionTest, prestoDateTruncTimestampAdjustTimezoneFallsBack) { .config(QueryConfig::kAdjustTimestampToTimezone, "true") .countResults(task); - ASSERT_FALSE(wasCudfFilterProjectUsed(task)); - ASSERT_TRUE(wasDefaultFilterProjectUsed(task)); + ASSERT_TRUE(wasCudfFilterProjectUsed(task)); + ASSERT_FALSE(wasDefaultFilterProjectUsed(task)); +} + +TEST_F(ToCudfSelectionTest, prestoDateTruncTimestampWithTimeZoneUsesCudf) { + // date_trunc(timestamp with time zone) is evaluated on GPU (per-row embedded + // zone), so the plan runs on cuDF rather than falling back. + auto input = makeRowVector( + {"c0"}, + {makeFlatVector( + {pack(1'736'971'261'123, tz::getTimeZoneID("America/Los_Angeles")), + pack(1'736'971'261'123, tz::getTimeZoneID("Asia/Kolkata"))}, + TIMESTAMP_WITH_TIME_ZONE())}); + + auto plan = PlanBuilder() + .values({input}) + .project({"date_trunc('day', c0) AS result"}) + .planNode(); + + std::shared_ptr task; + AssertQueryBuilder(plan).config("cudf.enabled", true).countResults(task); + + ASSERT_TRUE(wasCudfFilterProjectUsed(task)); + ASSERT_FALSE(wasDefaultFilterProjectUsed(task)); } TEST_F(ToCudfSelectionTest, prestoDateTruncSubHourAdjustTimezoneUsesCudf) { @@ -242,7 +302,7 @@ TEST_F(ToCudfSelectionTest, prestoDateTruncSubHourAdjustTimezoneUsesCudf) { TEST_F( ToCudfSelectionTest, - nestedPrestoDateTruncTimestampAdjustTimezoneFallsBack) { + nestedPrestoDateTruncTimestampAdjustTimezoneUsesCudf) { auto input = makeRowVector( {"event_ts"}, {makeFlatVector( @@ -259,15 +319,17 @@ TEST_F( ASSERT_TRUE(wasCudfFilterProjectUsed(cudfTask)); ASSERT_FALSE(wasDefaultFilterProjectUsed(cudfTask)); - std::shared_ptr fallbackTask; + // A nested timezone-sensitive date_trunc(timestamp) also stays on cuDF under + // adjust_timestamp_to_session_timezone now that it is timezone-aware. + std::shared_ptr adjustTask; AssertQueryBuilder(plan) .config("cudf.enabled", true) .config(QueryConfig::kSessionTimezone, "Asia/Kolkata") .config(QueryConfig::kAdjustTimestampToTimezone, "true") - .countResults(fallbackTask); + .countResults(adjustTask); - ASSERT_FALSE(wasCudfFilterProjectUsed(fallbackTask)); - ASSERT_TRUE(wasDefaultFilterProjectUsed(fallbackTask)); + ASSERT_TRUE(wasCudfFilterProjectUsed(adjustTask)); + ASSERT_FALSE(wasDefaultFilterProjectUsed(adjustTask)); } TEST_F(ToCudfSelectionTest, prestoDateTruncDateAdjustTimezoneUsesCudf) {