feat(cudf): GPU timezone-aware date_trunc and date_add

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_
   }

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<velox::exec::Expr>& expr) {
-  const auto dateTruncName =
-      CudfConfig::getInstance().functionNamePrefix + "date_trunc";
-  return expr->name() == dateTruncName &&
-      DateTruncFunction::isTimezoneSensitive(expr);
-}
-
-bool containsTimezoneSensitiveDateTrunc(
-    const std::shared_ptr<velox::exec::Expr>& 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<exec::ExprSet> 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);
         });
   }
 

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<velox::RowTypePtr> 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);
     }
   }
 }

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 <cudf/ast/expressions.hpp>
@@ -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

velox/experimental/cudf/expression/ArrayAccessFunctions.cpp

@@ -519,6 +519,7 @@ class ArrayAccessFunction : public CudfFunction {
 
   ColumnOrView eval(
       std::vector<ColumnOrView>& 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.

velox/experimental/cudf/expression/AstExpression.cpp

@@ -35,10 +35,11 @@ cudf::ast::expression const& createAstTree(
     cudf::ast::tree& tree,
     std::vector<std::unique_ptr<cudf::scalar>>& scalars,
     const RowTypePtr& inputRowSchema,
-    std::vector<PrecomputeInstruction>& precomputeInstructions) {
-  AstContext context{
-      tree, scalars, {inputRowSchema}, {precomputeInstructions}, expr};
-  return context.pushExprToTree(expr);
+    std::vector<PrecomputeInstruction>& 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<PrecomputeInstruction>& leftPrecomputeInstructions,
-    std::vector<PrecomputeInstruction>& rightPrecomputeInstructions) {
-  AstContext context{
+    std::vector<PrecomputeInstruction>& 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<velox::exec::Expr> 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<velox::exec::Expr> expr) {
         return ASTExpression::canEvaluate(expr);
       },
-      [](std::shared_ptr<velox::exec::Expr> expr, const RowTypePtr& row) {
-        return std::make_shared<ASTExpression>(std::move(expr), row);
+      [](std::shared_ptr<velox::exec::Expr> expr,
+         const RowTypePtr& row,
+         const CudfDateTimeContext& context) {
+        return std::make_shared<ASTExpression>(std::move(expr), row, context);
       },
       /*overwrite=*/false);
 }

velox/experimental/cudf/expression/AstExpression.h

@@ -29,7 +29,8 @@ cudf::ast::expression const& createAstTree(
     cudf::ast::tree& tree,
     std::vector<std::unique_ptr<cudf::scalar>>& scalars,
     const RowTypePtr& inputRowSchema,
-    std::vector<PrecomputeInstruction>& precomputeInstructions);
+    std::vector<PrecomputeInstruction>& precomputeInstructions,
+    const CudfDateTimeContext& context);
 
 cudf::ast::expression const& createAstTree(
     const std::shared_ptr<velox::exec::Expr>& expr,
@@ -38,7 +39,8 @@ cudf::ast::expression const& createAstTree(
     const RowTypePtr& leftRowSchema,
     const RowTypePtr& rightRowSchema,
     std::vector<PrecomputeInstruction>& leftPrecomputeInstructions,
-    std::vector<PrecomputeInstruction>& rightPrecomputeInstructions);
+    std::vector<PrecomputeInstruction>& 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<velox::exec::Expr> expr,
-      const RowTypePtr& inputRowSchema);
+      const RowTypePtr& inputRowSchema,
+      const CudfDateTimeContext& context);
 
   // Evaluates the expression tree for the given input columns
   ColumnOrView eval(

velox/experimental/cudf/expression/AstExpressionUtils.h

@@ -415,6 +415,10 @@ struct AstContext {
       precomputeInstructions;
   const std::shared_ptr<velox::exec::Expr>
       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);

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)