draft: fix(cudf): Restore NLJ build byte limit and timestamp AST compares

velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp

@@ -41,6 +41,8 @@
 #include <cudf/search.hpp>
 #include <cudf/stream_compaction.hpp>
 
+#include <limits>
+
 namespace facebook::velox::cudf_velox {
 
 namespace {
@@ -116,6 +118,36 @@ CudfNestedLoopJoinBridge::getBuildStream() {
   return buildStream_;
 }
 
+void CudfNestedLoopJoinBridge::reserveBuildBytes(
+    uint64_t bytes,
+    uint64_t maxBuildBytes) {
+  std::lock_guard<std::mutex> l(mutex_);
+  VELOX_USER_CHECK_LE(
+      retainedBuildBytes_,
+      maxBuildBytes,
+      "CudfNestedLoopJoin build byte accounting exceeded configured limit: "
+      "retained={} limit={} config={}",
+      retainedBuildBytes_,
+      maxBuildBytes,
+      CudfConfig::kCudfNestedLoopJoinMaxBuildBytes);
+  VELOX_USER_CHECK_LE(
+      bytes,
+      maxBuildBytes - retainedBuildBytes_,
+      "CudfNestedLoopJoin build exceeds configured device-memory limit "
+      "before retaining next batch: retained={} next={} limit={} config={}",
+      retainedBuildBytes_,
+      bytes,
+      maxBuildBytes,
+      CudfConfig::kCudfNestedLoopJoinMaxBuildBytes);
+  retainedBuildBytes_ += bytes;
+}
+
+void CudfNestedLoopJoinBridge::releaseBuildBytes(uint64_t bytes) {
+  std::lock_guard<std::mutex> l(mutex_);
+  VELOX_CHECK_GE(retainedBuildBytes_, bytes);
+  retainedBuildBytes_ -= bytes;
+}
+
 // ============================================================================
 // Build Operator Implementation
 // ============================================================================
@@ -136,15 +168,36 @@ CudfNestedLoopJoinBuild::CudfNestedLoopJoinBuild(
           NvtxMethodFlag::kNoMoreInput,
           std::nullopt,
           joinNode),
-      joinNode_(joinNode) {}
+      joinNode_(joinNode),
+      maxBuildBytes_(driverCtx->queryConfig().get<uint64_t>(
+          CudfConfig::kCudfNestedLoopJoinMaxBuildBytes,
+          std::numeric_limits<uint64_t>::max())) {}
 
 // Accumulates input batches in memory.
 // All batches are kept as CudfVectors (GPU memory) until join completes.
 void CudfNestedLoopJoinBuild::doAddInput(RowVectorPtr input) {
   if (input->size() > 0) {
     auto cudfInput = std::dynamic_pointer_cast<CudfVector>(input);
     VELOX_CHECK_NOT_NULL(cudfInput);
-    inputs_.push_back(std::move(cudfInput)); // Store in GPU memory
+    // Enforce the exact device payload before retaining the next batch. This
+    // is the runtime backstop for planner estimates used by replicated MPP
+    // Cartesian joins.
+    const auto inputBytes = cudfInput->estimateFlatSize();
+    if (!buildBridge_) {
+      auto joinBridge = operatorCtx_->task()->getCustomJoinBridge(
+          operatorCtx_->driverCtx()->splitGroupId, planNodeId());
+      buildBridge_ =
+          std::dynamic_pointer_cast<CudfNestedLoopJoinBridge>(joinBridge);
+      VELOX_CHECK_NOT_NULL(buildBridge_);
+    }
+    buildBridge_->reserveBuildBytes(inputBytes, maxBuildBytes_);
+    try {
+      inputs_.push_back(std::move(cudfInput)); // Store in GPU memory
+      bufferedBuildBytes_ += inputBytes;
+    } catch (...) {
+      buildBridge_->releaseBuildBytes(inputBytes);
+      throw;
+    }
   }
 }
 
@@ -177,62 +230,97 @@ void CudfNestedLoopJoinBuild::doNoMoreInput() {
     return; // Not the last driver - just wait
   }
 
-  // This driver was chosen to collect data from all peers
+  // Always wake the waiting build peers, including when host-vector growth or
+  // build concatenation throws. Leaving this guard after the peer merge would
+  // strand the other drivers on an allocation failure.
+  SCOPE_EXIT {
+    peers.clear();
+    for (auto& promise : promises) {
+      promise.setValue(); // Unblock other build operators
+    }
+  };
+
+  std::vector<CudfNestedLoopJoinBuild*> peerBuilds;
+  peerBuilds.reserve(peers.size());
+  auto mergedInputCount = inputs_.size();
   for (auto& peer : peers) {
     auto op = peer->findOperator(planNodeId());
     auto* build = dynamic_cast<CudfNestedLoopJoinBuild*>(op);
     VELOX_CHECK_NOT_NULL(build);
+    VELOX_CHECK_LE(
+        build->inputs_.size(),
+        inputs_.max_size() - mergedInputCount,
+        "CudfNestedLoopJoin build input vector exceeds host size limit");
+    mergedInputCount += build->inputs_.size();
+    peerBuilds.push_back(build);
+  }
+
+  // Allocate once before moving peer-owned vectors and byte-accounting
+  // ownership. With sufficient capacity, shared_ptr moves below are noexcept.
+  inputs_.reserve(mergedInputCount);
+
+  // This driver was chosen to collect data from all peers
+  for (auto* build : peerBuilds) {
+    // The shared join bridge reserved these bytes before each peer retained
+    // its batches, so transferring ownership does not need another reserve.
     inputs_.insert(
         inputs_.end(),
         std::make_move_iterator(build->inputs_.begin()),
         std::make_move_iterator(build->inputs_.end()));
+    bufferedBuildBytes_ += build->bufferedBuildBytes_;
+    build->bufferedBuildBytes_ = 0;
   }
 
-  // Wake up peer build operators when we finish transferring data
-  SCOPE_EXIT {
-    peers.clear();
-    for (auto& promise : promises) {
-      promise.setValue(); // Unblock other build operators
-    }
-  };
-
   // Concatenate all input batches into a single cuDF table.
   // getConcatenatedTable throws if the total row count exceeds cudf::size_type
   // limits (~2.1B rows). We don't use getConcatenatedTableBatched here because
   // batching the build side does not prevent output overflow for NLJ: a cross
   // join output is probe_rows × build_rows regardless of how the build is
   // split.
   auto stream = cudfGlobalStreamPool().get_stream();
-  auto table = getConcatenatedTable(
-      std::exchange(inputs_, {}),
-      joinNode_->sources()[1]->outputType(),
-      stream,
-      get_output_mr());
-
-  // Record the build-ready event now, immediately after the build table
-  // is materialized on `stream` - not lazily on the probe side - so it
-  // captures this exact completion point before `stream` can be recycled
-  // by cudfGlobalStreamPool() for unrelated work. Every probe operator
-  // instance/batch just waits on this same event before reading buildData_
-  // (see CudfNestedLoopJoinProbe::waitForBuildReady()). `stream` is also
-  // exposed via setBuildStream() below: buildData_'s eventual free is
-  // stream-ordered on `stream` (it was allocated here), so every probe read
-  // makes `stream` wait on its own completion first (see
-  // CudfNestedLoopJoinProbe::recordReadCompletion()), ensuring that free
-  // can't run before all such reads are done.
-  auto buildReadyEvent = std::make_shared<CudaEvent>(cudaEventDisableTiming);
-  buildReadyEvent->recordFrom(stream);
-
-  // Transfer build data to bridge - this will unblock probe operators.
   auto joinBridge = operatorCtx_->task()->getCustomJoinBridge(
       operatorCtx_->driverCtx()->splitGroupId, planNodeId());
   auto bridge = std::dynamic_pointer_cast<CudfNestedLoopJoinBridge>(joinBridge);
+  VELOX_CHECK_NOT_NULL(bridge);
+  buildBridge_ = bridge;
+
+  try {
+    auto table = getConcatenatedTable(
+        std::exchange(inputs_, {}),
+        joinNode_->sources()[1]->outputType(),
+        stream,
+        get_output_mr());
 
-  bridge->setBuildReadyEvent(std::move(buildReadyEvent));
-  bridge->setBuildStream(stream);
-  bridge->setData(
-      std::make_optional(
-          std::shared_ptr<cudf::table>(std::move(table)))); // Wake probes
+    // Record the build-ready event now, immediately after the build table
+    // is materialized on `stream` - not lazily on the probe side - so it
+    // captures this exact completion point before `stream` can be recycled
+    // by cudfGlobalStreamPool() for unrelated work. Every probe operator
+    // instance/batch just waits on this same event before reading buildData_
+    // (see CudfNestedLoopJoinProbe::waitForBuildReady()). `stream` is also
+    // exposed via setBuildStream() below: buildData_'s eventual free is
+    // stream-ordered on `stream` (it was allocated here), so every probe read
+    // makes `stream` wait on its own completion first (see
+    // CudfNestedLoopJoinProbe::recordReadCompletion()), ensuring that free
+    // can't run before all such reads are done.
+    auto buildReadyEvent = std::make_shared<CudaEvent>(cudaEventDisableTiming);
+    buildReadyEvent->recordFrom(stream);
+
+    // Transfer build data to bridge - this will unblock probe operators.
+    bridge->setBuildReadyEvent(std::move(buildReadyEvent));
+    bridge->setBuildStream(stream);
+    bridge->setData(
+        std::make_optional(
+            std::shared_ptr<cudf::table>(std::move(table)))); // Wake probes
+    // The bridge now owns the complete build. Keep its global accounting until
+    // bridge destruction, but disarm this operator's close-time release.
+    bufferedBuildBytes_ = 0;
+  } catch (...) {
+    // std::exchange empties inputs_ before concatenation. Release accounting
+    // here because doClose() can no longer observe those retained vectors.
+    bridge->releaseBuildBytes(bufferedBuildBytes_);
+    bufferedBuildBytes_ = 0;
+    throw;
+  }
 }
 
 exec::BlockingReason CudfNestedLoopJoinBuild::isBlocked(
@@ -249,7 +337,13 @@ bool CudfNestedLoopJoinBuild::isFinished() {
 }
 
 void CudfNestedLoopJoinBuild::doClose() {
+  if (bufferedBuildBytes_ > 0 && !inputs_.empty()) {
+    VELOX_CHECK_NOT_NULL(buildBridge_);
+    buildBridge_->releaseBuildBytes(bufferedBuildBytes_);
+  }
   inputs_.clear();
+  bufferedBuildBytes_ = 0;
+  buildBridge_.reset();
   Operator::close();
 }
 

velox/experimental/cudf/exec/CudfNestedLoopJoin.h

@@ -133,16 +133,6 @@ class CudfNestedLoopJoinBuild : public CudfOperatorBase {
   ContinueFuture future_{ContinueFuture::makeEmpty()};
 };
 
-/// Constructs the build operator in the same translation unit that defines
-/// CudfNestedLoopJoinBuild.  Keep external adapters on this factory instead of
-/// inlining make_unique<CudfNestedLoopJoinBuild>: doing so prevents a stale
-/// caller object from baking in an older concrete class size after the build
-/// operator gains or removes private state.
-std::unique_ptr<exec::Operator> makeCudfNestedLoopJoinBuild(
-    int32_t operatorId,
-    exec::DriverCtx* driverCtx,
-    std::shared_ptr<const core::NestedLoopJoinNode> joinNode);
-
 /// Performs nested loop join using cuDF APIs.
 ///
 /// Supports inner, left, right, full outer, and left semi project joins.

velox/experimental/cudf/expression/AstExpressionUtils.h

@@ -536,7 +536,33 @@ cudf::ast::expression const& AstContext::pushExprToTree(
   };
 
   if (!detail::isAstExprSupported(expr)) {
-    return compileSubExpression();
+    // TIMESTAMP is gated from AST/JIT evaluation (see
+    // containsAstUnsupportedType), so a two-sided comparison such as
+    // gt(p_ts, b_ts) would otherwise try to precompute the whole predicate
+    // and fail with sideIdx == -2. cuDF AST comparison operators still
+    // accept timestamp column_references, so compile the comparison from
+    // native field refs instead of precomputing across both join sides.
+    const bool timestampInputField = expr->isFieldAccessKind() &&
+        expr->type() && expr->type()->isTimestamp() &&
+        expr->asUnchecked<core::FieldAccessTypedExpr>()->isInputColumn();
+    if (timestampInputField) {
+      // Fall through to kFieldAccess and emit a column_reference.
+    } else if (expr->isCallKind()) {
+      const auto name = stripPrefix(
+          expr->asUnchecked<core::CallTypedExpr>()->name(),
+          CudfConfig::getInstance().functionNamePrefix);
+      auto it = binaryOps.find(name);
+      if (it != binaryOps.end() && len == 2 &&
+          (expr->inputs()[0]->type()->isTimestamp() ||
+           expr->inputs()[1]->type()->isTimestamp())) {
+        auto const& op1 = pushExprToTree(expr->inputs()[0]);
+        auto const& op2 = pushExprToTree(expr->inputs()[1]);
+        return tree.push(Operation{it->second, op1, op2});
+      }
+      return compileSubExpression();
+    } else {
+      return compileSubExpression();
+    }
   }
 
   switch (expr->kind()) {

velox/experimental/cudf/tests/NestedLoopJoinTest.cpp

@@ -15,6 +15,7 @@
  */
 
 #include "velox/experimental/cudf/CudfConfig.h"
+#include "velox/experimental/cudf/exec/CudfConversion.h"
 #include "velox/experimental/cudf/exec/ToCudf.h"
 
 #include "velox/common/base/tests/GTestUtils.h"
@@ -99,6 +100,7 @@ TEST_F(CudfNestedLoopJoinTest, buildByteLimitAndCleanup) {
 
   VELOX_ASSERT_THROW(
       AssertQueryBuilder(plan)
+          .config(cudf_velox::CudfFromVelox::kGpuBatchSizeRows, "1")
           .config(
               cudf_velox::CudfConfig::kCudfNestedLoopJoinMaxBuildBytes, "16")
           .copyResults(pool()),