diff --git a/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp b/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp index ab423cda7cf..38bdfc96109 100644 --- a/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp +++ b/velox/experimental/cudf/exec/CudfNestedLoopJoin.cpp @@ -41,6 +41,8 @@ #include #include +#include + 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 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 l(mutex_); + VELOX_CHECK_GE(retainedBuildBytes_, bytes); + retainedBuildBytes_ -= bytes; +} + // ============================================================================ // Build Operator Implementation // ============================================================================ @@ -136,7 +168,10 @@ CudfNestedLoopJoinBuild::CudfNestedLoopJoinBuild( NvtxMethodFlag::kNoMoreInput, std::nullopt, joinNode), - joinNode_(joinNode) {} + joinNode_(joinNode), + maxBuildBytes_(driverCtx->queryConfig().get( + CudfConfig::kCudfNestedLoopJoinMaxBuildBytes, + std::numeric_limits::max())) {} // Accumulates input batches in memory. // All batches are kept as CudfVectors (GPU memory) until join completes. @@ -144,7 +179,25 @@ void CudfNestedLoopJoinBuild::doAddInput(RowVectorPtr input) { if (input->size() > 0) { auto cudfInput = std::dynamic_pointer_cast(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(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,25 +230,47 @@ 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 peerBuilds; + peerBuilds.reserve(peers.size()); + auto mergedInputCount = inputs_.size(); for (auto& peer : peers) { auto op = peer->findOperator(planNodeId()); auto* build = dynamic_cast(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 @@ -203,36 +278,49 @@ void CudfNestedLoopJoinBuild::doNoMoreInput() { // 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(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(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(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(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(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(); } diff --git a/velox/experimental/cudf/exec/CudfNestedLoopJoin.h b/velox/experimental/cudf/exec/CudfNestedLoopJoin.h index 6143667dc98..d0f63f6bdb3 100644 --- a/velox/experimental/cudf/exec/CudfNestedLoopJoin.h +++ b/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: 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 makeCudfNestedLoopJoinBuild( - int32_t operatorId, - exec::DriverCtx* driverCtx, - std::shared_ptr joinNode); - /// Performs nested loop join using cuDF APIs. /// /// Supports inner, left, right, full outer, and left semi project joins. diff --git a/velox/experimental/cudf/expression/AstExpressionUtils.h b/velox/experimental/cudf/expression/AstExpressionUtils.h index 6fddb46e53c..3b2acc82503 100644 --- a/velox/experimental/cudf/expression/AstExpressionUtils.h +++ b/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()->isInputColumn(); + if (timestampInputField) { + // Fall through to kFieldAccess and emit a column_reference. + } else if (expr->isCallKind()) { + const auto name = stripPrefix( + expr->asUnchecked()->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()) { diff --git a/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp b/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp index 64d44220631..7c2478e1cc1 100644 --- a/velox/experimental/cudf/tests/NestedLoopJoinTest.cpp +++ b/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()),