diff --git a/velox/experimental/cudf/exec/CudfOrderBy.cpp b/velox/experimental/cudf/exec/CudfOrderBy.cpp index 5d984127190..dbeadc00ad2 100644 --- a/velox/experimental/cudf/exec/CudfOrderBy.cpp +++ b/velox/experimental/cudf/exec/CudfOrderBy.cpp @@ -14,8 +14,8 @@ * limitations under the License. */ -#include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/CudfConfig.h" +#include "velox/experimental/cudf/CudfNoDefaults.h" #include "velox/experimental/cudf/exec/CudfOrderBy.h" #include "velox/experimental/cudf/exec/GpuResources.h" #include "velox/experimental/cudf/exec/NvtxHelper.h" @@ -107,8 +107,10 @@ bool isSpillSafeType(const TypePtr& type) { } } return true; - case TypeKind::VARBINARY: case TypeKind::MAP: + return type->size() == 2 && isSpillSafeType(type->childAt(0)) && + isSpillSafeType(type->childAt(1)); + case TypeKind::VARBINARY: case TypeKind::UNKNOWN: case TypeKind::FUNCTION: case TypeKind::OPAQUE: @@ -125,13 +127,13 @@ bool isSupportedSortKeyType(const TypePtr& type) { // Nested cuDF sort semantics have not been validated against Velox. Nested // values may still be carried as payload when every leaf is spill-safe. - return type->kind() != TypeKind::ARRAY && type->kind() != TypeKind::ROW; + return type->kind() != TypeKind::ARRAY && type->kind() != TypeKind::MAP && + type->kind() != TypeKind::ROW; } void updateAtomicMax(std::atomic& target, uint64_t value) { auto current = target.load(); - while (current < value && - !target.compare_exchange_weak(current, value)) { + while (current < value && !target.compare_exchange_weak(current, value)) { } } @@ -406,27 +408,30 @@ void CudfOrderBy::spillSortedRun() { spilled_ = true; } - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=sortRun.concatenate.begin " - "bufferedBytes={} bufferedInputs={} existingRuns={} " - "sortedRunBytes={} mergeFanIn={}", - orderByNode_->id(), - bufferedBytes_, - inputs_.size(), - sortedRuns_.size(), - sortedRunBytes_, - mergeFanIn_)); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=sortRun.concatenate.begin " + "bufferedBytes={} bufferedInputs={} existingRuns={} " + "sortedRunBytes={} mergeFanIn={}", + orderByNode_->id(), + bufferedBytes_, + inputs_.size(), + sortedRuns_.size(), + sortedRunBytes_, + mergeFanIn_)); auto input = getConcatenatedTable( std::exchange(inputs_, {}), outputType_, stateStream_, get_output_mr()); bufferedBytes_ = 0; - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=sortRun.concatenate.end rows={}", - orderByNode_->id(), - input->num_rows())); - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=sortRun.sort.begin rows={}", - orderByNode_->id(), - input->num_rows())); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=sortRun.concatenate.end rows={}", + orderByNode_->id(), + input->num_rows())); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=sortRun.sort.begin rows={}", + orderByNode_->id(), + input->num_rows())); auto sorted = cudf::sort_by_key( input->view(), input->view().select(sortKeys_), @@ -434,10 +439,11 @@ void CudfOrderBy::spillSortedRun() { nullOrder_, stateStream_, get_output_mr()); - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=sortRun.sort.end rows={}", - orderByNode_->id(), - sorted->num_rows())); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=sortRun.sort.end rows={}", + orderByNode_->id(), + sorted->num_rows())); auto path = fmt::format( "{}/run-{:06}.parquet", spillDirectory_, spillFileSequence_++); @@ -445,25 +451,27 @@ void CudfOrderBy::spillSortedRun() { cudf::io::sink_info{path}, sorted->view()) .row_group_size_bytes(kSpillRowGroupBytes) .build(); - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=sortRun.write.begin rows={} " - "existingRuns={} path={}", - orderByNode_->id(), - sorted->num_rows(), - sortedRuns_.size(), - path)); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=sortRun.write.begin rows={} " + "existingRuns={} path={}", + orderByNode_->id(), + sorted->num_rows(), + sortedRuns_.size(), + path)); cudf::io::write_parquet(options, stateStream_); SortedRun run; run.path = std::move(path); sortedRuns_.push_back(std::move(run)); - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=sortRun.write.end rows={} runs={} " - "path={}", - orderByNode_->id(), - sorted->num_rows(), - sortedRuns_.size(), - sortedRuns_.back().path)); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=sortRun.write.end rows={} runs={} " + "path={}", + orderByNode_->id(), + sorted->num_rows(), + sortedRuns_.size(), + sortedRuns_.back().path)); ::malloc_trim(0); } @@ -531,9 +539,7 @@ std::unique_ptr CudfOrderBy::mergeNextPausedBatch( continue; } auto slices = cudf::slice( - run->chunk->view(), - {run->chunkOffset, run->chunk->num_rows()}, - stream); + run->chunk->view(), {run->chunkOffset, run->chunk->num_rows()}, stream); VELOX_CHECK_EQ(slices.size(), 1); activeRuns.push_back(run); remainingViews.push_back(slices.front()); @@ -544,10 +550,8 @@ std::unique_ptr CudfOrderBy::mergeNextPausedBatch( stats.maxActiveRuns = std::max(stats.maxActiveRuns, activeRuns.size()); - stats.maxResidentRows = - std::max(stats.maxResidentRows, residentRows); - stats.maxResidentBytes = - std::max(stats.maxResidentBytes, residentBytes); + stats.maxResidentRows = std::max(stats.maxResidentRows, residentRows); + stats.maxResidentBytes = std::max(stats.maxResidentBytes, residentBytes); updateAtomicMax(observedMaxActiveRuns, activeRuns.size()); VELOX_CHECK_LE( activeRuns.size(), @@ -571,8 +575,8 @@ std::unique_ptr CudfOrderBy::mergeNextPausedBatch( std::vector boundaryRows; boundaryRows.reserve(remainingViews.size()); for (const auto& view : remainingViews) { - auto last = cudf::slice( - view, {view.num_rows() - 1, view.num_rows()}, stream); + auto last = + cudf::slice(view, {view.num_rows() - 1, view.num_rows()}, stream); boundaryRows.push_back(last.front()); } auto boundaryCandidates = cudf::concatenate(boundaryRows, stream, mr); @@ -606,8 +610,7 @@ std::unique_ptr CudfOrderBy::mergeNextPausedBatch( VELOX_CHECK(!safeViews.empty(), "Paused OrderBy merge made no progress"); auto output = safeViews.size() == 1 ? std::make_unique(safeViews.front(), stream, mr) - : cudf::merge( - safeViews, sortKeys_, columnOrder_, nullOrder_, stream, mr); + : cudf::merge(safeViews, sortKeys_, columnOrder_, nullOrder_, stream, mr); for (size_t index = 0; index < activeRuns.size(); ++index) { activeRuns[index]->chunkOffset += consumed[index]; @@ -668,20 +671,15 @@ void CudfOrderBy::compactSortedRunsForMerge() { const auto outputPath = fmt::format( "{}/merge-{:06}.parquet", spillDirectory_, spillFileSequence_++); - auto writerOptions = - cudf::io::chunked_parquet_writer_options::builder( - cudf::io::sink_info{outputPath}) - .row_group_size_bytes(kSpillRowGroupBytes) - .build(); + auto writerOptions = cudf::io::chunked_parquet_writer_options::builder( + cudf::io::sink_info{outputPath}) + .row_group_size_bytes(kSpillRowGroupBytes) + .build(); cudf::io::chunked_parquet_writer writer(writerOptions, stateStream_); bool groupFinished{false}; while (!groupFinished) { auto merged = mergeNextPausedBatch( - runs, - stateStream_, - get_output_mr(), - groupFinished, - levelStats); + runs, stateStream_, get_output_mr(), groupFinished, levelStats); if (merged && merged->num_rows() > 0) { writer.write(merged->view()); } @@ -719,41 +717,43 @@ void CudfOrderBy::compactSortedRunsForMerge() { compactionStats.outputBatches += levelStats.outputBatches; compactionStats.outputRows += levelStats.outputRows; compactionStats.outputBytes += levelStats.outputBytes; - compactionStats.maxResidentRows = std::max( - compactionStats.maxResidentRows, levelStats.maxResidentRows); - compactionStats.maxResidentBytes = std::max( - compactionStats.maxResidentBytes, levelStats.maxResidentBytes); - compactionStats.maxOutputBytes = std::max( - compactionStats.maxOutputBytes, levelStats.maxOutputBytes); - compactionStats.maxActiveRuns = std::max( - compactionStats.maxActiveRuns, levelStats.maxActiveRuns); - - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=compaction.level.end " - "inputRuns={} outputRuns={} sourceChunks={} outputBatches={} " - "maxResidentBytes={} maxOutputBytes={} maxActiveRuns={}", - orderByNode_->id(), - inputRunCount, - sortedRuns_.size(), - levelStats.sourceChunks, - levelStats.outputBatches, - levelStats.maxResidentBytes, - levelStats.maxOutputBytes, - levelStats.maxActiveRuns)); + compactionStats.maxResidentRows = + std::max(compactionStats.maxResidentRows, levelStats.maxResidentRows); + compactionStats.maxResidentBytes = + std::max(compactionStats.maxResidentBytes, levelStats.maxResidentBytes); + compactionStats.maxOutputBytes = + std::max(compactionStats.maxOutputBytes, levelStats.maxOutputBytes); + compactionStats.maxActiveRuns = + std::max(compactionStats.maxActiveRuns, levelStats.maxActiveRuns); + + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=compaction.level.end " + "inputRuns={} outputRuns={} sourceChunks={} outputBatches={} " + "maxResidentBytes={} maxOutputBytes={} maxActiveRuns={}", + orderByNode_->id(), + inputRunCount, + sortedRuns_.size(), + levelStats.sourceChunks, + levelStats.outputBatches, + levelStats.maxResidentBytes, + levelStats.maxOutputBytes, + levelStats.maxActiveRuns)); } stateStream_.synchronize(); - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=compaction.end runs={} " - "sourceChunks={} outputBatches={} maxResidentBytes={} " - "maxOutputBytes={} maxActiveRuns={}", - orderByNode_->id(), - sortedRuns_.size(), - compactionStats.sourceChunks, - compactionStats.outputBatches, - compactionStats.maxResidentBytes, - compactionStats.maxOutputBytes, - compactionStats.maxActiveRuns)); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=compaction.end runs={} " + "sourceChunks={} outputBatches={} maxResidentBytes={} " + "maxOutputBytes={} maxActiveRuns={}", + orderByNode_->id(), + sortedRuns_.size(), + compactionStats.sourceChunks, + compactionStats.outputBatches, + compactionStats.maxResidentBytes, + compactionStats.maxOutputBytes, + compactionStats.maxActiveRuns)); } void CudfOrderBy::initializeSortedRunReaders() { @@ -776,13 +776,14 @@ void CudfOrderBy::initializeSortedRunReaders() { run.chunkBytes = 0; } readersInitialized_ = true; - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=output.merge.begin runs={} " - "chunkReadLimit={} passReadLimit={}", - orderByNode_->id(), - sortedRuns_.size(), - mergeChunkBytes.load(), - kMergePassBytes)); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=output.merge.begin runs={} " + "chunkReadLimit={} passReadLimit={}", + orderByNode_->id(), + sortedRuns_.size(), + mergeChunkBytes.load(), + kMergePassBytes)); } void CudfOrderBy::prepareSpilledOutput() { @@ -803,29 +804,26 @@ std::unique_ptr CudfOrderBy::mergeNextSortedBatch() { runs.push_back(&run); } auto result = mergeNextPausedBatch( - runs, - stateStream_, - get_output_mr(), - mergeFinished_, - outputMergeStats_); + runs, stateStream_, get_output_mr(), mergeFinished_, outputMergeStats_); if (mergeFinished_) { - logDeviceMemorySnapshot(fmt::format( - "operator=CudfOrderBy node={} state=output.merge.end runs={} " - "sourceChunks={} sourceRows={} sourceBytes={} outputBatches={} " - "outputRows={} outputBytes={} maxResidentRows={} " - "maxResidentBytes={} maxOutputBytes={} maxActiveRuns={}", - orderByNode_->id(), - sortedRuns_.size(), - outputMergeStats_.sourceChunks, - outputMergeStats_.sourceRows, - outputMergeStats_.sourceBytes, - outputMergeStats_.outputBatches, - outputMergeStats_.outputRows, - outputMergeStats_.outputBytes, - outputMergeStats_.maxResidentRows, - outputMergeStats_.maxResidentBytes, - outputMergeStats_.maxOutputBytes, - outputMergeStats_.maxActiveRuns)); + logDeviceMemorySnapshot( + fmt::format( + "operator=CudfOrderBy node={} state=output.merge.end runs={} " + "sourceChunks={} sourceRows={} sourceBytes={} outputBatches={} " + "outputRows={} outputBytes={} maxResidentRows={} " + "maxResidentBytes={} maxOutputBytes={} maxActiveRuns={}", + orderByNode_->id(), + sortedRuns_.size(), + outputMergeStats_.sourceChunks, + outputMergeStats_.sourceRows, + outputMergeStats_.sourceBytes, + outputMergeStats_.outputBatches, + outputMergeStats_.outputRows, + outputMergeStats_.outputBytes, + outputMergeStats_.maxResidentRows, + outputMergeStats_.maxResidentBytes, + outputMergeStats_.maxOutputBytes, + outputMergeStats_.maxActiveRuns)); } return result; } @@ -849,11 +847,10 @@ CudfVectorPtr CudfOrderBy::takePendingOutput() { VELOX_CHECK_LT(pendingOutputOffset_, totalRows); const auto remainingRows = totalRows - pendingOutputOffset_; const auto byteLimit = outputChunkBytes.load(); - cudf::size_type targetRows = - std::min(remainingRows, maxOutputRows.load()); + cudf::size_type targetRows = std::min(remainingRows, maxOutputRows.load()); if (pendingOutputBytes_ > byteLimit) { - const auto proportionalRows = static_cast( - std::max( + const auto proportionalRows = + static_cast(std::max( 1, static_cast(totalRows) * byteLimit / pendingOutputBytes_)); @@ -886,12 +883,10 @@ CudfVectorPtr CudfOrderBy::takePendingOutput() { return output; } - const auto proportionalRows = static_cast( - std::max( - 1, - static_cast(targetRows) * byteLimit / actualBytes)); - targetRows = - std::min(targetRows - 1, proportionalRows); + const auto proportionalRows = + static_cast(std::max( + 1, static_cast(targetRows) * byteLimit / actualBytes)); + targetRows = std::min(targetRows - 1, proportionalRows); } } diff --git a/velox/experimental/cudf/exec/CudfWindow.cpp b/velox/experimental/cudf/exec/CudfWindow.cpp index 7aa2c251635..3c3ee959b4d 100644 --- a/velox/experimental/cudf/exec/CudfWindow.cpp +++ b/velox/experimental/cudf/exec/CudfWindow.cpp @@ -27,11 +27,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -42,6 +44,7 @@ #include #include #include +#include namespace facebook::velox::cudf_velox { namespace { @@ -51,7 +54,12 @@ namespace { // avoiding the thousands of tiny spill files produced by hash bucketing. constexpr uint64_t kWindowSortedRunBytes = 3ULL << 30; constexpr uint64_t kWindowMergeChunkBytes = 256ULL << 20; +constexpr uint64_t kWindowStreamingActiveRowsBytes = 256ULL << 20; std::atomic windowSpillDirectorySequence{0}; +std::atomic testingStreamingActiveRowsBytes{0}; +std::atomic testingStreamingReplayChunkBytes{0}; +std::atomic observedStreamingSpillWrites{0}; +std::atomic observedStreamingSpillCleanups{0}; bool isSupportedScalarWindowType(const TypePtr& type) { switch (type->kind()) { @@ -116,6 +124,32 @@ bool isSupportedSumInputType(const TypePtr& type) { } } +bool isSupportedFullPartitionCountFunction( + const core::WindowNode::Function& function) { + if (function.functionCall->name() != "count" || function.ignoreNulls || + function.functionCall->type()->kind() != TypeKind::BIGINT || + function.functionCall->inputs().size() != 1 || + !isFullPartitionRowsFrame(function.frame)) { + return false; + } + const auto constant = + std::dynamic_pointer_cast( + function.functionCall->inputs().front()); + return constant != nullptr && !constant->isNull(); +} + +bool isSupportedRangeRunningSumFunction( + const core::WindowNode::Function& function) { + if (function.functionCall->name() != "sum" || function.ignoreNulls || + function.functionCall->inputs().size() != 1 || + function.frame.type != core::WindowNode::WindowType::kRange || + !isUnboundedPrecedingToCurrentRowFrame(function.frame)) { + return false; + } + const auto& input = function.functionCall->inputs()[0]; + return isFieldAccessExpr(input) && isSupportedSumInputType(input->type()); +} + bool isSupportedFullPartitionSumFunction( const core::WindowNode::Function& function) { if (function.functionCall->name() != "sum" || function.ignoreNulls || @@ -189,6 +223,75 @@ cudf::size_type firstSearchPosition( return result; } +std::unique_ptr makeSizeTypeColumn( + const std::vector& values, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto column = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT32}, + static_cast(values.size()), + cudf::mask_state::UNALLOCATED, + stream, + mr); + if (!values.empty()) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + column->mutable_view().data(), + values.data(), + values.size() * sizeof(cudf::size_type), + cudaMemcpyHostToDevice, + stream.value())); + } + return column; +} + +std::unique_ptr makeInt64Column( + const std::vector& values, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto column = cudf::make_fixed_width_column( + cudf::data_type{cudf::type_id::INT64}, + static_cast(values.size()), + cudf::mask_state::UNALLOCATED, + stream, + mr); + if (!values.empty()) { + CUDF_CUDA_TRY(cudaMemcpyAsync( + column->mutable_view().data(), + values.data(), + values.size() * sizeof(int64_t), + cudaMemcpyHostToDevice, + stream.value())); + } + return column; +} + +std::unique_ptr copyScalar( + const cudf::scalar& value, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + auto column = cudf::make_column_from_scalar(value, 1, stream, mr); + return cudf::get_element(column->view(), 0, stream, mr); +} + +std::unique_ptr addValidScalars( + const cudf::scalar& left, + const cudf::scalar& right, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + VELOX_CHECK(left.is_valid(stream)); + VELOX_CHECK(right.is_valid(stream)); + VELOX_CHECK(left.type() == right.type()); + auto leftColumn = cudf::make_column_from_scalar(left, 1, stream, mr); + auto result = cudf::binary_operation( + leftColumn->view(), + right, + cudf::binary_operator::ADD, + left.type(), + stream, + mr); + return cudf::get_element(result->view(), 0, stream, mr); +} + } // namespace bool isSupportedCudfWindowNode( @@ -197,17 +300,23 @@ bool isSupportedCudfWindowNode( return false; } - bool hasRowNumber = false; + bool hasRankLike = false; + bool hasFullPartitionCount = false; bool hasFullPartitionSum = false; bool hasRunningSum = false; + bool hasRangeRunningSum = false; bool hasFirstValue = false; for (const auto& function : node->windowFunctions()) { if (isSupportedRankLikeFunction(function)) { - hasRowNumber = true; + hasRankLike = true; + } else if (isSupportedFullPartitionCountFunction(function)) { + hasFullPartitionCount = true; } else if (isSupportedFullPartitionSumFunction(function)) { hasFullPartitionSum = true; } else if (isSupportedRunningSumFunction(function)) { hasRunningSum = true; + } else if (isSupportedRangeRunningSumFunction(function)) { + hasRangeRunningSum = true; } else if (isSupportedFirstValueFunction(function)) { hasFirstValue = true; } else { @@ -215,16 +324,30 @@ bool isSupportedCudfWindowNode( } } - if ((hasRowNumber || hasFirstValue || hasRunningSum) && + if ((hasRankLike || hasFirstValue || hasRunningSum || hasRangeRunningSum) && node->sortingKeys().empty()) { return false; } - if (hasFirstValue && node->partitionKeys().empty()) { + if ((hasFirstValue || hasFullPartitionCount || hasRangeRunningSum) && + node->partitionKeys().empty()) { + return false; + } + + if ((hasFullPartitionCount || hasRangeRunningSum) && !node->inputsSorted()) { return false; } - if ((hasRowNumber || hasFirstValue || hasRunningSum) && hasFullPartitionSum) { + if (hasFullPartitionCount && !node->sortingKeys().empty()) { + return false; + } + + const auto functionKinds = static_cast(hasRankLike) + + static_cast(hasFullPartitionCount) + + static_cast(hasFullPartitionSum) + static_cast(hasRunningSum) + + static_cast(hasRangeRunningSum) + static_cast(hasFirstValue); + if (functionKinds > 1 && + (hasFullPartitionCount || hasFullPartitionSum || hasRangeRunningSum)) { return false; } @@ -243,6 +366,30 @@ bool isSupportedCudfWindowNode( return true; } +void CudfWindow::testingSetStreamingMemoryLimits( + uint64_t activeRowsBytes, + uint64_t replayChunkBytes) { + VELOX_CHECK_GT(activeRowsBytes, 0); + VELOX_CHECK_GT(replayChunkBytes, 0); + testingStreamingActiveRowsBytes.store(activeRowsBytes); + testingStreamingReplayChunkBytes.store(replayChunkBytes); + observedStreamingSpillWrites.store(0); + observedStreamingSpillCleanups.store(0); +} + +void CudfWindow::testingResetStreamingMemoryLimits() { + testingStreamingActiveRowsBytes.store(0); + testingStreamingReplayChunkBytes.store(0); +} + +uint64_t CudfWindow::testingStreamingSpillWrites() { + return observedStreamingSpillWrites.load(); +} + +uint64_t CudfWindow::testingStreamingSpillCleanups() { + return observedStreamingSpillCleanups.load(); +} + CudfWindow::CudfWindow( int32_t operatorId, exec::DriverCtx* driverCtx, @@ -260,14 +407,36 @@ CudfWindow::CudfWindow( windowNode_(windowNode), inputType_(windowNode->sources()[0]->outputType()), rankLikeStreaming_( - windowNode->inputsSorted() && - !windowNode->partitionKeys().empty() && + windowNode->inputsSorted() && !windowNode->partitionKeys().empty() && !windowNode->sortingKeys().empty() && std::all_of( windowNode->windowFunctions().begin(), windowNode->windowFunctions().end(), isSupportedRankLikeFunction)), - stateStream_(cudfGlobalStreamPool().get_stream()) { + fullPartitionCountStreaming_( + windowNode->inputsSorted() && !windowNode->partitionKeys().empty() && + windowNode->sortingKeys().empty() && + std::all_of( + windowNode->windowFunctions().begin(), + windowNode->windowFunctions().end(), + isSupportedFullPartitionCountFunction)), + rangeSumStreaming_( + windowNode->inputsSorted() && !windowNode->partitionKeys().empty() && + !windowNode->sortingKeys().empty() && + std::all_of( + windowNode->windowFunctions().begin(), + windowNode->windowFunctions().end(), + isSupportedRangeRunningSumFunction)), + boundedStreaming_(fullPartitionCountStreaming_ || rangeSumStreaming_), + stateStream_(cudfGlobalStreamPool().get_stream()), + activeRowsLimit_( + testingStreamingActiveRowsBytes.load() > 0 + ? testingStreamingActiveRowsBytes.load() + : kWindowStreamingActiveRowsBytes), + replayChunkLimit_( + testingStreamingReplayChunkBytes.load() > 0 + ? testingStreamingReplayChunkBytes.load() + : kWindowMergeChunkBytes) { for (const auto& key : windowNode->partitionKeys()) { const auto channel = exec::exprToChannel(key.get(), inputType_); VELOX_CHECK( @@ -300,6 +469,28 @@ void CudfWindow::doAddInput(RowVectorPtr input) { auto cudfInput = std::dynamic_pointer_cast(input); VELOX_CHECK_NOT_NULL(cudfInput, "Expected CudfVector input"); + if (boundedStreaming_) { + VELOX_CHECK_NULL( + pendingOutput_, + "CudfWindow received sorted input before prior output was drained"); + VELOX_CHECK_NULL( + deferredInput_, + "CudfWindow received sorted input while deferred input remained"); + VELOX_CHECK_NULL( + streamingReplay_, + "CudfWindow received sorted input while replay was active"); + const auto inputStream = cudfInput->stream(); + if (inputStream.value() != stateStream_.value()) { + std::vector inputStreams{inputStream}; + cudf::detail::join_streams(inputStreams, stateStream_); + } + VELOX_CHECK( + cudfInput->rebindStream(stateStream_), + "CudfWindow cannot rebind sorted input to its state stream"); + deferredInput_ = cudfInput->release(); + advanceBoundedStreaming(); + return; + } if (rankLikeStreaming_) { VELOX_CHECK_NULL( @@ -357,6 +548,19 @@ void CudfWindow::doAddInput(RowVectorPtr input) { void CudfWindow::doNoMoreInput() { Operator::noMoreInput(); + if (boundedStreaming_) { + advanceBoundedStreaming(); + if (activeKey_) { + finalizeActiveGroup(); + advanceBoundedStreaming(); + } + if (pendingOutput_ == nullptr && streamingReplay_ == nullptr && + deferredInput_ == nullptr) { + finished_ = true; + } + return; + } + if (rankLikeStreaming_) { finished_ = true; return; @@ -374,8 +578,29 @@ void CudfWindow::doNoMoreInput() { } RowVectorPtr CudfWindow::doGetOutput() { + auto takePending = [&]() -> RowVectorPtr { + auto output = std::exchange(pendingOutput_, nullptr); + if (boundedStreaming_ && noMoreInput_ && streamingReplay_ == nullptr && + deferredInput_ == nullptr && activeKey_ == nullptr) { + finished_ = true; + } + return output; + }; + if (pendingOutput_ != nullptr) { - return std::exchange(pendingOutput_, nullptr); + return takePending(); + } + + if (boundedStreaming_) { + advanceBoundedStreaming(); + if (pendingOutput_ != nullptr) { + return takePending(); + } + if (noMoreInput_ && streamingReplay_ == nullptr && + deferredInput_ == nullptr && activeKey_ == nullptr) { + finished_ = true; + } + return nullptr; } if (finished_ || !noMoreInput_) { @@ -413,6 +638,590 @@ RowVectorPtr CudfWindow::doGetOutput() { pool(), outputType_, output->num_rows(), std::move(output), stream); } +std::vector CudfWindow::streamingGroupChannels() const { + auto channels = partitionKeyChannels_; + if (rangeSumStreaming_) { + channels.insert( + channels.end(), sortKeyChannels_.begin(), sortKeyChannels_.end()); + } + return channels; +} + +std::vector CudfWindow::streamingGroupOrders() const { + std::vector orders( + partitionKeyChannels_.size(), cudf::order::ASCENDING); + if (rangeSumStreaming_) { + orders.insert(orders.end(), sortOrders_.begin(), sortOrders_.end()); + } + return orders; +} + +std::vector CudfWindow::streamingGroupNullOrders() const { + std::vector nullOrders( + partitionKeyChannels_.size(), cudf::null_order::BEFORE); + if (rangeSumStreaming_) { + nullOrders.insert( + nullOrders.end(), sortNullOrders_.begin(), sortNullOrders_.end()); + } + return nullOrders; +} + +cudf::size_type CudfWindow::trailingGroupStart( + cudf::table_view input, + const std::vector& channels, + const std::vector& orders, + const std::vector& nullOrders, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + VELOX_CHECK_GT(input.num_rows(), 0); + auto keys = cudf::table_view(selectColumns(input, channels)); + auto last = + cudf::slice(keys, {input.num_rows() - 1, input.num_rows()}, stream); + auto positions = + cudf::lower_bound(keys, last.front(), orders, nullOrders, stream, mr); + return firstSearchPosition(positions->view(), stream); +} + +uint64_t CudfWindow::measureTableBytes(std::unique_ptr& table) { + VELOX_CHECK_NOT_NULL(table); + auto vector = std::make_shared( + pool(), inputType_, table->num_rows(), std::move(table), stateStream_); + const auto bytes = vector->estimateFlatSize(); + table = vector->release(); + return bytes; +} + +void CudfWindow::setPendingOutput(std::unique_ptr output) { + VELOX_CHECK_NOT_NULL(output); + VELOX_CHECK_NULL(pendingOutput_); + if (output->num_rows() == 0) { + return; + } + pendingOutput_ = std::make_shared( + pool(), outputType_, output->num_rows(), std::move(output), stateStream_); +} + +std::unique_ptr CudfWindow::appendConstantResults( + std::unique_ptr rows, + const std::vector>& results) { + VELOX_CHECK_NOT_NULL(rows); + VELOX_CHECK_EQ(results.size(), windowNode_->windowFunctions().size()); + const auto numRows = rows->num_rows(); + auto columns = rows->release(); + columns.reserve(columns.size() + results.size()); + for (const auto& result : results) { + VELOX_CHECK_NOT_NULL(result); + columns.push_back( + cudf::make_column_from_scalar( + *result, numRows, stateStream_, get_output_mr())); + } + return std::make_unique(std::move(columns)); +} + +std::unique_ptr CudfWindow::computeFullPartitionCountOutput( + std::unique_ptr input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const { + VELOX_CHECK_NOT_NULL(input); + auto partitionColumns = selectColumns(input->view(), partitionKeyChannels_); + cudf::groupby::groupby grouper( + cudf::table_view(partitionColumns), + cudf::null_policy::INCLUDE, + cudf::sorted::YES, + std::vector( + partitionKeyChannels_.size(), cudf::order::ASCENDING), + std::vector( + partitionKeyChannels_.size(), cudf::null_order::BEFORE)); + auto groups = grouper.get_groups({}, stream, mr); + VELOX_CHECK_GE(groups.offsets.size(), 2); + + std::vector counts; + std::vector repeats; + counts.reserve(groups.offsets.size() - 1); + repeats.reserve(groups.offsets.size() - 1); + for (size_t i = 1; i < groups.offsets.size(); ++i) { + const auto size = groups.offsets[i] - groups.offsets[i - 1]; + counts.push_back(size); + repeats.push_back(size); + } + auto countsColumn = makeInt64Column(counts, stream, mr); + auto repeatsColumn = makeSizeTypeColumn(repeats, stream, mr); + + auto columns = input->release(); + columns.reserve(columns.size() + windowNode_->windowFunctions().size()); + for (size_t i = 0; i < windowNode_->windowFunctions().size(); ++i) { + auto repeated = cudf::repeat( + cudf::table_view{{countsColumn->view()}}, + repeatsColumn->view(), + stream, + mr); + auto resultColumns = repeated->release(); + VELOX_CHECK_EQ(resultColumns.size(), 1); + columns.push_back(std::move(resultColumns.front())); + } + return std::make_unique(std::move(columns)); +} + +std::unique_ptr CudfWindow::computeRangeRunningSumOutput( + std::unique_ptr input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { + VELOX_CHECK_NOT_NULL(input); + VELOX_CHECK_GT(input->num_rows(), 0); + const auto inputView = input->view(); + const auto numFunctions = windowNode_->windowFunctions().size(); + const auto inputSize = inputType_->size(); + + std::vector> runningColumns; + std::vector runningViews; + runningColumns.reserve(numFunctions); + runningViews.reserve(numFunctions); + for (size_t i = 0; i < numFunctions; ++i) { + runningColumns.push_back(computeRunningPartitionSumColumn( + inputView, + windowNode_->windowFunctions()[i], + outputType_->childAt(inputSize + i), + stream, + mr)); + runningViews.push_back(runningColumns.back()->view()); + } + + auto partitionColumns = selectColumns(inputView, partitionKeyChannels_); + const std::vector partitionOrders( + partitionKeyChannels_.size(), cudf::order::ASCENDING); + const std::vector partitionNullOrders( + partitionKeyChannels_.size(), cudf::null_order::BEFORE); + cudf::groupby::groupby partitionGrouper( + cudf::table_view(partitionColumns), + cudf::null_policy::INCLUDE, + cudf::sorted::YES, + partitionOrders, + partitionNullOrders); + const std::vector preceding( + numFunctions, cudf::replace_policy::PRECEDING); + auto filledRunning = partitionGrouper.replace_nulls( + cudf::table_view(runningViews), preceding, stream, mr); + runningColumns = filledRunning.second->release(); + + const auto samePartitionEnd = continuingPrefixSize( + inputView, + partitionKeyChannels_, + cumulativePartitionKey_, + partitionOrders, + partitionNullOrders, + stream, + mr); + if (samePartitionEnd > 0) { + VELOX_CHECK_EQ(cumulativeSums_.size(), numFunctions); + VELOX_CHECK_EQ(cumulativeValidCounts_.size(), numFunctions); + for (size_t i = 0; i < numFunctions; ++i) { + if (cumulativeValidCounts_[i] == 0) { + continue; + } + auto prefix = + cudf::slice(runningColumns[i]->view(), {0, samePartitionEnd}, stream); + auto added = cudf::binary_operation( + prefix.front(), + *cumulativeSums_[i], + cudf::binary_operator::ADD, + runningColumns[i]->type(), + stream, + mr); + auto fixed = + cudf::replace_nulls(added->view(), *cumulativeSums_[i], stream, mr); + if (samePartitionEnd == input->num_rows()) { + runningColumns[i] = std::move(fixed); + } else { + auto suffix = cudf::slice( + runningColumns[i]->view(), + {samePartitionEnd, input->num_rows()}, + stream); + const std::vector pieces{ + fixed->view(), suffix.front()}; + runningColumns[i] = cudf::concatenate(pieces, stream, mr); + } + } + } + + auto peerColumns = selectColumns(inputView, streamingGroupChannels()); + cudf::groupby::groupby peerGrouper( + cudf::table_view(peerColumns), + cudf::null_policy::INCLUDE, + cudf::sorted::YES, + streamingGroupOrders(), + streamingGroupNullOrders()); + auto groups = peerGrouper.get_groups({}, stream, mr); + VELOX_CHECK_GE(groups.offsets.size(), 2); + + std::vector peerEnds; + std::vector repeats; + peerEnds.reserve(groups.offsets.size() - 1); + repeats.reserve(groups.offsets.size() - 1); + for (size_t i = 1; i < groups.offsets.size(); ++i) { + peerEnds.push_back(groups.offsets[i] - 1); + repeats.push_back(groups.offsets[i] - groups.offsets[i - 1]); + } + auto peerEndsColumn = makeSizeTypeColumn(peerEnds, stream, mr); + auto repeatsColumn = makeSizeTypeColumn(repeats, stream, mr); + + runningViews.clear(); + for (const auto& column : runningColumns) { + runningViews.push_back(column->view()); + } + auto peerResults = cudf::gather( + cudf::table_view(runningViews), + peerEndsColumn->view(), + cudf::out_of_bounds_policy::DONT_CHECK, + stream, + mr); + auto repeatedResults = + cudf::repeat(peerResults->view(), repeatsColumn->view(), stream, mr); + auto resultColumns = repeatedResults->release(); + + auto finalPartition = cudf::slice( + cudf::table_view(partitionColumns), + {input->num_rows() - 1, input->num_rows()}, + stream); + cumulativePartitionKey_ = + std::make_unique(finalPartition.front(), stream, mr); + cumulativeSums_.clear(); + cumulativeValidCounts_.clear(); + cumulativeSums_.reserve(numFunctions); + cumulativeValidCounts_.reserve(numFunctions); + for (const auto& column : resultColumns) { + auto value = + cudf::get_element(column->view(), input->num_rows() - 1, stream, mr); + cumulativeValidCounts_.push_back(value->is_valid(stream) ? 1 : 0); + cumulativeSums_.push_back(std::move(value)); + } + + auto columns = input->release(); + columns.reserve(columns.size() + resultColumns.size()); + for (auto& column : resultColumns) { + columns.push_back(std::move(column)); + } + return std::make_unique(std::move(columns)); +} + +void CudfWindow::updateActiveRangeSums(cudf::table_view rows) { + VELOX_CHECK(rangeSumStreaming_); + const auto numFunctions = windowNode_->windowFunctions().size(); + if (activeSums_.empty()) { + activeSums_.resize(numFunctions); + activeValidCounts_.assign(numFunctions, 0); + } + VELOX_CHECK_EQ(activeSums_.size(), numFunctions); + VELOX_CHECK_EQ(activeValidCounts_.size(), numFunctions); + + for (size_t i = 0; i < numFunctions; ++i) { + const auto& function = windowNode_->windowFunctions()[i]; + const auto valueChannel = exec::exprToChannel( + function.functionCall->inputs()[0].get(), inputType_); + VELOX_CHECK_NE(valueChannel, kConstantChannel); + const auto values = rows.column(valueChannel); + const auto validCount = rows.num_rows() - values.null_count(); + auto aggregation = cudf::make_sum_aggregation(); + auto reduced = cudf::reduce( + values, + *aggregation, + veloxToCudfDataType(outputType_->childAt(inputType_->size() + i)), + stateStream_, + get_output_mr()); + if (activeValidCounts_[i] > 0 && validCount > 0) { + activeSums_[i] = addValidScalars( + *activeSums_[i], *reduced, stateStream_, get_output_mr()); + } else if (!activeSums_[i] || validCount > 0) { + activeSums_[i] = std::move(reduced); + } + activeValidCounts_[i] += validCount; + } +} + +void CudfWindow::spillActiveRows() { + if (!activeRows_ || activeRows_->num_rows() == 0) { + return; + } + namespace fs = std::filesystem; + if (spillDirectory_.empty()) { + const auto sequence = windowSpillDirectorySequence.fetch_add(1); + spillDirectory_ = (fs::temp_directory_path() / + fmt::format( + "velox-cudf-window-spill-{}-{}", + static_cast(::getpid()), + sequence)) + .string(); + fs::create_directories(spillDirectory_); + } + + auto path = fmt::format( + "{}/active-{:06}.parquet", spillDirectory_, spillFileSequence_++); + auto options = cudf::io::parquet_writer_options::builder( + cudf::io::sink_info{path}, activeRows_->view()) + .build(); + cudf::io::write_parquet(options, stateStream_); + activeSpillPaths_.push_back(std::move(path)); + activeRows_.reset(); + activeRowsBytes_ = 0; + streamingSpilled_ = true; + observedStreamingSpillWrites.fetch_add(1); + ::malloc_trim(0); +} + +void CudfWindow::appendActiveGroup(std::unique_ptr rows) { + VELOX_CHECK_NOT_NULL(rows); + VELOX_CHECK_GT(rows->num_rows(), 0); + activeRowCount_ += rows->num_rows(); + if (rangeSumStreaming_) { + updateActiveRangeSums(rows->view()); + } + if (activeRows_) { + const std::vector pieces{ + activeRows_->view(), rows->view()}; + activeRows_ = cudf::concatenate(pieces, stateStream_, get_output_mr()); + } else { + activeRows_ = std::move(rows); + } + activeRowsBytes_ = measureTableBytes(activeRows_); + if (activeRowsBytes_ >= activeRowsLimit_) { + spillActiveRows(); + } +} + +void CudfWindow::startActiveGroup(std::unique_ptr rows) { + VELOX_CHECK_NOT_NULL(rows); + VELOX_CHECK_GT(rows->num_rows(), 0); + VELOX_CHECK_NULL(activeKey_); + VELOX_CHECK_EQ(activeRowCount_, 0); + VELOX_CHECK(activeSpillPaths_.empty()); + VELOX_CHECK_NULL(activeRows_); + + auto groupColumns = selectColumns(rows->view(), streamingGroupChannels()); + auto firstGroup = + cudf::slice(cudf::table_view(groupColumns), {0, 1}, stateStream_); + activeKey_ = std::make_unique( + firstGroup.front(), stateStream_, get_output_mr()); + + if (rangeSumStreaming_) { + const std::vector partitionOrders( + partitionKeyChannels_.size(), cudf::order::ASCENDING); + const std::vector partitionNullOrders( + partitionKeyChannels_.size(), cudf::null_order::BEFORE); + const auto samePartition = continuingPrefixSize( + rows->view(), + partitionKeyChannels_, + cumulativePartitionKey_, + partitionOrders, + partitionNullOrders, + stateStream_, + get_output_mr()); + if (samePartition == 0) { + cumulativePartitionKey_.reset(); + cumulativeSums_.clear(); + cumulativeValidCounts_.clear(); + } + } + appendActiveGroup(std::move(rows)); +} + +void CudfWindow::finalizeActiveGroup() { + VELOX_CHECK_NOT_NULL(activeKey_); + VELOX_CHECK_GT(activeRowCount_, 0); + VELOX_CHECK( + activeRows_ != nullptr || !activeSpillPaths_.empty(), + "Active Window group has no retained rows"); + VELOX_CHECK_NULL(streamingReplay_); + + std::vector> results; + const auto numFunctions = windowNode_->windowFunctions().size(); + results.reserve(numFunctions); + if (fullPartitionCountStreaming_) { + for (size_t i = 0; i < numFunctions; ++i) { + results.push_back( + std::make_unique>( + activeRowCount_, true, stateStream_, get_output_mr())); + } + } else { + VELOX_CHECK(rangeSumStreaming_); + VELOX_CHECK_EQ(activeSums_.size(), numFunctions); + VELOX_CHECK_EQ(activeValidCounts_.size(), numFunctions); + for (size_t i = 0; i < numFunctions; ++i) { + const auto hasCumulative = + i < cumulativeValidCounts_.size() && cumulativeValidCounts_[i] > 0; + const auto hasActive = activeValidCounts_[i] > 0; + if (hasCumulative && hasActive) { + results.push_back(addValidScalars( + *cumulativeSums_[i], + *activeSums_[i], + stateStream_, + get_output_mr())); + } else if (hasCumulative) { + results.push_back( + copyScalar(*cumulativeSums_[i], stateStream_, get_output_mr())); + } else { + results.push_back( + copyScalar(*activeSums_[i], stateStream_, get_output_mr())); + } + } + + std::vector partitionPositions( + partitionKeyChannels_.size()); + std::iota(partitionPositions.begin(), partitionPositions.end(), 0); + auto partitionColumns = + selectColumns(activeKey_->view(), partitionPositions); + cumulativePartitionKey_ = std::make_unique( + cudf::table_view(partitionColumns), stateStream_, get_output_mr()); + cumulativeSums_.clear(); + cumulativeValidCounts_.clear(); + cumulativeSums_.reserve(numFunctions); + cumulativeValidCounts_.reserve(numFunctions); + for (const auto& result : results) { + cumulativeValidCounts_.push_back(result->is_valid(stateStream_) ? 1 : 0); + cumulativeSums_.push_back( + copyScalar(*result, stateStream_, get_output_mr())); + } + } + + streamingReplay_ = std::make_unique(); + streamingReplay_->paths = std::exchange(activeSpillPaths_, {}); + streamingReplay_->memoryRows = std::exchange(activeRows_, nullptr); + streamingReplay_->results = std::move(results); + activeRowsBytes_ = 0; + activeKey_.reset(); + activeRowCount_ = 0; + activeSums_.clear(); + activeValidCounts_.clear(); +} + +void CudfWindow::prepareNextStreamingReplayOutput() { + while (streamingReplay_ && pendingOutput_ == nullptr) { + if (streamingReplay_->reader) { + if (streamingReplay_->reader->has_next()) { + auto chunk = streamingReplay_->reader->read_chunk(); + if (chunk.tbl->num_rows() > 0) { + setPendingOutput(appendConstantResults( + std::move(chunk.tbl), streamingReplay_->results)); + return; + } + continue; + } + streamingReplay_->reader.reset(); + ++streamingReplay_->nextPath; + continue; + } + + if (streamingReplay_->nextPath < streamingReplay_->paths.size()) { + auto options = + cudf::io::parquet_reader_options::builder( + cudf::io::source_info{ + streamingReplay_->paths[streamingReplay_->nextPath]}) + .build(); + streamingReplay_->reader = + std::make_unique( + replayChunkLimit_, 0, options, stateStream_, get_output_mr()); + continue; + } + + if (streamingReplay_->memoryRows) { + auto output = appendConstantResults( + std::exchange(streamingReplay_->memoryRows, nullptr), + streamingReplay_->results); + streamingReplay_.reset(); + setPendingOutput(std::move(output)); + return; + } + streamingReplay_.reset(); + } +} + +void CudfWindow::processDeferredStreamingInput() { + VELOX_CHECK_NOT_NULL(deferredInput_); + VELOX_CHECK_GT(deferredInput_->num_rows(), 0); + const auto channels = streamingGroupChannels(); + const auto orders = streamingGroupOrders(); + const auto nullOrders = streamingGroupNullOrders(); + + if (activeKey_) { + const auto continuingRows = continuingPrefixSize( + deferredInput_->view(), + channels, + activeKey_, + orders, + nullOrders, + stateStream_, + get_output_mr()); + if (continuingRows > 0) { + auto prefix = copyTableSlice( + deferredInput_->view(), + 0, + continuingRows, + stateStream_, + get_output_mr()); + auto suffix = copyTableSlice( + deferredInput_->view(), + continuingRows, + deferredInput_->num_rows(), + stateStream_, + get_output_mr()); + deferredInput_ = std::move(suffix); + appendActiveGroup(std::move(prefix)); + if (!deferredInput_ || deferredInput_->num_rows() == 0) { + deferredInput_.reset(); + return; + } + } + finalizeActiveGroup(); + return; + } + + const auto trailingStart = trailingGroupStart( + deferredInput_->view(), + channels, + orders, + nullOrders, + stateStream_, + get_output_mr()); + auto trailing = copyTableSlice( + deferredInput_->view(), + trailingStart, + deferredInput_->num_rows(), + stateStream_, + get_output_mr()); + auto complete = copyTableSlice( + deferredInput_->view(), 0, trailingStart, stateStream_, get_output_mr()); + deferredInput_.reset(); + + std::unique_ptr output; + if (complete && complete->num_rows() > 0) { + output = fullPartitionCountStreaming_ + ? computeFullPartitionCountOutput( + std::move(complete), stateStream_, get_output_mr()) + : computeRangeRunningSumOutput( + std::move(complete), stateStream_, get_output_mr()); + } + startActiveGroup(std::move(trailing)); + if (output) { + setPendingOutput(std::move(output)); + } +} + +void CudfWindow::advanceBoundedStreaming() { + while (pendingOutput_ == nullptr) { + if (streamingReplay_) { + prepareNextStreamingReplayOutput(); + continue; + } + if (deferredInput_) { + processDeferredStreamingInput(); + continue; + } + if (noMoreInput_ && activeKey_) { + finalizeActiveGroup(); + continue; + } + return; + } +} + void CudfWindow::spillSortedRun() { if (inputs_.empty()) { return; @@ -672,17 +1481,22 @@ void CudfWindow::cleanupSpillFiles() { sortedRuns_.clear(); mergeCarry_.reset(); partitionCarry_.reset(); + activeSpillPaths_.clear(); if (spillDirectory_.empty()) { return; } std::error_code error; std::filesystem::remove_all(spillDirectory_, error); spillDirectory_.clear(); + if (streamingSpilled_ && !error) { + observedStreamingSpillCleanups.fetch_add(1); + } + streamingSpilled_ = false; ::malloc_trim(0); } void CudfWindow::doClose() { - if (rankLikeStreaming_) { + if (rankLikeStreaming_ || boundedStreaming_) { // close() can run while a task exception is unwinding. A poisoned CUDA // context must not replace that original failure, but owners are always // destroyed between best-effort pre/post drains. @@ -693,6 +1507,13 @@ void CudfWindow::doClose() { << error.what(); } pendingOutput_.reset(); + deferredInput_.reset(); + activeRows_.reset(); + activeKey_.reset(); + activeSums_.clear(); + cumulativePartitionKey_.reset(); + cumulativeSums_.clear(); + streamingReplay_.reset(); previousPartitionKey_.reset(); previousOrderKey_.reset(); try { @@ -757,9 +1578,8 @@ std::unique_ptr CudfWindow::fixRankLikeColumn( return localResult; } - auto addOffset = - [&](cudf::column_view values, - int64_t offset) -> std::unique_ptr { + auto addOffset = [&](cudf::column_view values, + int64_t offset) -> std::unique_ptr { switch (values.type().id()) { case cudf::type_id::INT32: { cudf::numeric_scalar scalar( @@ -819,13 +1639,13 @@ std::unique_ptr CudfWindow::fixRankLikeColumn( }; auto appendOffset = [&](cudf::size_type begin, cudf::size_type end, int64_t offset) { - if (begin == end) { - return; - } - auto slice = cudf::slice(localResult->view(), {begin, end}, stream); - ownedSegments.push_back(addOffset(slice.front(), offset)); - segments.push_back(ownedSegments.back()->view()); - }; + if (begin == end) { + return; + } + auto slice = cudf::slice(localResult->view(), {begin, end}, stream); + ownedSegments.push_back(addOffset(slice.front(), offset)); + segments.push_back(ownedSegments.back()->view()); + }; cudf::size_type fixedEnd{0}; if (functionName == "rank") { @@ -845,10 +1665,7 @@ std::unique_ptr CudfWindow::fixRankLikeColumn( } fixedEnd = sameOrderEnd; } - appendOffset( - fixedEnd, - samePartitionEnd, - previousPartitionRows_); + appendOffset(fixedEnd, samePartitionEnd, previousPartitionRows_); appendUnchanged(samePartitionEnd, localResult->size()); VELOX_CHECK(!segments.empty()); return segments.size() == 1 @@ -866,16 +1683,15 @@ void CudfWindow::updateRankLikeState( partitionKeyChannels_.size(), cudf::order::ASCENDING); const std::vector partitionNullOrders( partitionKeyChannels_.size(), cudf::null_order::BEFORE); - const auto continuedRows = hasRankLikeState_ - ? continuingPrefixSize( - sortedInput, - partitionKeyChannels_, - previousPartitionKey_, - partitionOrders, - partitionNullOrders, - stream, - mr) - : 0; + const auto continuedRows = hasRankLikeState_ ? continuingPrefixSize( + sortedInput, + partitionKeyChannels_, + previousPartitionKey_, + partitionOrders, + partitionNullOrders, + stream, + mr) + : 0; const bool continuedWholeBatch = continuedRows == numRows; auto partitionColumns = @@ -909,8 +1725,7 @@ void CudfWindow::updateRankLikeState( mr); const auto lastPeerBegin = firstSearchPosition(orderPositions->view(), stream); - const bool lastPeerContinued = - continuedWholeBatch && + const bool lastPeerContinued = continuedWholeBatch && continuingPrefixSize( sortedInput, sortKeyChannels_, @@ -920,13 +1735,11 @@ void CudfWindow::updateRankLikeState( stream, mr) == numRows; - const auto priorRows = - continuedWholeBatch ? previousPartitionRows_ : 0; + const auto priorRows = continuedWholeBatch ? previousPartitionRows_ : 0; if (!lastPeerContinued) { previousRank_ = priorRows + lastPeerBegin + 1; } - previousPartitionRows_ = - priorRows + numRows - trailingPartitionBegin; + previousPartitionRows_ = priorRows + numRows - trailingPartitionBegin; previousPartitionKey_ = std::make_unique(lastPartition.front(), stream, mr); previousOrderKey_ = @@ -1355,11 +2168,7 @@ std::unique_ptr CudfWindow::computeOutputTable( } if (rankLikeStreaming_) { result = fixRankLikeColumn( - std::move(result), - functionName, - sortedView, - stream, - mr); + std::move(result), functionName, sortedView, stream, mr); } resultColumns.push_back(std::move(result)); } diff --git a/velox/experimental/cudf/exec/CudfWindow.h b/velox/experimental/cudf/exec/CudfWindow.h index daaf2a12672..b7531636e07 100644 --- a/velox/experimental/cudf/exec/CudfWindow.h +++ b/velox/experimental/cudf/exec/CudfWindow.h @@ -21,9 +21,12 @@ #include "velox/core/PlanNode.h" #include +#include #include + #include +#include #include #include #include @@ -41,6 +44,9 @@ bool isSupportedCudfWindowNode( /// CURRENT ROW frame. /// - first(field) / first_value(field) over a partitioned ordered UNBOUNDED /// PRECEDING to CURRENT ROW frame. +/// - count(non-null constant) over a full partition ROWS frame. +/// - sum(field) over a partitioned ordered RANGE UNBOUNDED PRECEDING to +/// CURRENT ROW frame. class CudfWindow : public CudfOperatorBase { public: CudfWindow( @@ -48,8 +54,17 @@ class CudfWindow : public CudfOperatorBase { exec::DriverCtx* driverCtx, const std::shared_ptr& windowNode); + static void testingSetStreamingMemoryLimits( + uint64_t activeRowsBytes, + uint64_t replayChunkBytes); + static void testingResetStreamingMemoryLimits(); + static uint64_t testingStreamingSpillWrites(); + static uint64_t testingStreamingSpillCleanups(); + bool needsInput() const override { - return !noMoreInput_ && pendingOutput_ == nullptr; + return !noMoreInput_ && pendingOutput_ == nullptr && + (!boundedStreaming_ || + (deferredInput_ == nullptr && streamingReplay_ == nullptr)); } exec::BlockingReason isBlocked(ContinueFuture* /*future*/) override { @@ -106,6 +121,40 @@ class CudfWindow : public CudfOperatorBase { rmm::device_async_resource_ref mr, bool inputAlreadySorted = false); + std::unique_ptr computeFullPartitionCountOutput( + std::unique_ptr input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + + std::unique_ptr computeRangeRunningSumOutput( + std::unique_ptr input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + + std::vector streamingGroupChannels() const; + std::vector streamingGroupOrders() const; + std::vector streamingGroupNullOrders() const; + cudf::size_type trailingGroupStart( + cudf::table_view input, + const std::vector& channels, + const std::vector& orders, + const std::vector& nullOrders, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + void advanceBoundedStreaming(); + void processDeferredStreamingInput(); + void startActiveGroup(std::unique_ptr rows); + void appendActiveGroup(std::unique_ptr rows); + void updateActiveRangeSums(cudf::table_view rows); + void finalizeActiveGroup(); + void prepareNextStreamingReplayOutput(); + void spillActiveRows(); + uint64_t measureTableBytes(std::unique_ptr& table); + void setPendingOutput(std::unique_ptr output); + std::unique_ptr appendConstantResults( + std::unique_ptr rows, + const std::vector>& results); + std::unique_ptr fixRankLikeColumn( std::unique_ptr localResult, std::string_view functionName, @@ -141,6 +190,14 @@ class CudfWindow : public CudfOperatorBase { CudfVectorPtr computeNextSortedOutput(); void cleanupSpillFiles(); + struct StreamingReplay { + std::vector paths; + size_t nextPath{0}; + std::unique_ptr reader; + std::unique_ptr memoryRows; + std::vector> results; + }; + struct SortedRun { std::string path; std::unique_ptr reader; @@ -148,10 +205,15 @@ class CudfWindow : public CudfOperatorBase { const std::shared_ptr windowNode_; const RowTypePtr inputType_; - // Only partitioned rank-like windows whose child guarantees ordering use - // the bounded cross-batch fixer. Other Window shapes retain legacy behavior. + // Only Window shapes whose child guarantees the complete ordering contract + // use cross-batch streaming state. Other shapes retain legacy behavior. const bool rankLikeStreaming_; + const bool fullPartitionCountStreaming_; + const bool rangeSumStreaming_; + const bool boundedStreaming_; const rmm::cuda_stream_view stateStream_; + const uint64_t activeRowsLimit_; + const uint64_t replayChunkLimit_; std::vector inputs_; CudfVectorPtr pendingOutput_; uint64_t bufferedBytes_{0}; @@ -165,6 +227,22 @@ class CudfWindow : public CudfOperatorBase { bool mergeFinished_{false}; bool spilled_{false}; bool finished_{false}; + bool streamingSpilled_{false}; + + std::unique_ptr deferredInput_; + std::unique_ptr activeRows_; + uint64_t activeRowsBytes_{0}; + std::vector activeSpillPaths_; + std::unique_ptr activeKey_; + int64_t activeRowCount_{0}; + std::vector> activeSums_; + std::vector activeValidCounts_; + + std::unique_ptr cumulativePartitionKey_; + std::vector> cumulativeSums_; + std::vector cumulativeValidCounts_; + + std::unique_ptr streamingReplay_; std::vector partitionKeyChannels_; std::vector sortKeyChannels_; diff --git a/velox/experimental/cudf/tests/AdapterOperatorTest.cpp b/velox/experimental/cudf/tests/AdapterOperatorTest.cpp index dfa133348d8..0d5cd57183a 100644 --- a/velox/experimental/cudf/tests/AdapterOperatorTest.cpp +++ b/velox/experimental/cudf/tests/AdapterOperatorTest.cpp @@ -28,6 +28,8 @@ #include "velox/functions/prestosql/window/WindowFunctionsRegistration.h" #include "velox/functions/sparksql/window/WindowFunctionsRegistration.h" +#include + using namespace facebook::velox; using namespace facebook::velox::exec; using namespace facebook::velox::exec::test; @@ -60,6 +62,19 @@ class AdapterOperatorTest : public OperatorTestBase { } return false; } + + std::shared_ptr withRowsFrame( + const core::PlanNodePtr& plan) { + auto window = std::dynamic_pointer_cast(plan); + VELOX_CHECK_NOT_NULL(window); + auto functions = window->windowFunctions(); + for (auto& function : functions) { + function.frame.type = core::WindowNode::WindowType::kRows; + } + return core::WindowNode::Builder(*window) + .windowFunctions(std::move(functions)) + .build(); + } }; TEST_F(AdapterOperatorTest, adapterStatsMergedIntoPlanNode) { @@ -122,12 +137,189 @@ TEST_F(AdapterOperatorTest, fullPartitionWindowSumUsesCudfWindow) { EXPECT_TRUE(wasCudfWindowUsed(task)); } +TEST_F(AdapterOperatorTest, streamingFullPartitionCountCrossesInputBatches) { + std::vector data{ + makeRowVector( + {"k", "payload"}, + {makeNullableFlatVector({std::nullopt, std::nullopt, 1}), + makeFlatVector({0, 1, 2})}), + makeRowVector( + {"k", "payload"}, + {makeNullableFlatVector({1, 1, 2}), + makeFlatVector({3, 4, 5})}), + makeRowVector( + {"k", "payload"}, + {makeNullableFlatVector({2, 2, 3}), + makeFlatVector({6, 7, 8})})}; + createDuckDbTable(data); + + auto plan = withRowsFrame( + PlanBuilder() + .values(data) + .streamingWindow( + {"count(1) over (partition by k rows between unbounded preceding " + "and unbounded following) as c"}) + .planNode()); + auto task = AssertQueryBuilder(duckDbQueryRunner_) + .config("cudf.enabled", true) + .plan(plan) + .assertResults( + "SELECT k, payload, count(1) over (partition by k rows " + "between unbounded preceding and unbounded following) " + "FROM tmp"); + EXPECT_TRUE(wasCudfWindowUsed(task)); +} + +TEST_F( + AdapterOperatorTest, + streamingFullPartitionCountProducesOutputBeforeNoMoreInput) { + std::vector data{ + makeRowVector( + {"k", "payload"}, + {makeFlatVector({1, 1}), makeFlatVector({0, 1})}), + makeRowVector( + {"k", "payload"}, + {makeFlatVector({1, 2}), makeFlatVector({2, 3})})}; + auto plan = withRowsFrame( + PlanBuilder() + .values(data) + .streamingWindow( + {"count(1) over (partition by k rows between unbounded preceding " + "and unbounded following) as c"}) + .planNode()); + auto windowNode = std::dynamic_pointer_cast(plan); + ASSERT_NE(windowNode, nullptr); + auto valuesNode = std::dynamic_pointer_cast( + windowNode->sources()[0]); + ASSERT_NE(valuesNode, nullptr); + + auto task = Task::create( + "streaming-count-produces-output-before-no-more-input", + core::PlanFragment{plan}, + 0, + core::QueryCtx::create(driverExecutor_.get()), + Task::ExecutionMode::kParallel); + auto driver = Driver::testingCreate( + std::make_unique(task, 0, 0, kUngroupedGroupId, 0)); + cudf_velox::CudfValues values(0, driver->driverCtx(), valuesNode); + cudf_velox::CudfWindow window(1, driver->driverCtx(), windowNode); + values.initialize(); + window.initialize(); + + auto first = values.getOutput(); + ASSERT_NE(first, nullptr); + window.addInput(std::move(first)); + EXPECT_TRUE(window.needsInput()); + EXPECT_EQ(window.getOutput(), nullptr); + + auto second = values.getOutput(); + ASSERT_NE(second, nullptr); + window.addInput(std::move(second)); + EXPECT_FALSE(window.needsInput()); + auto completedPartition = window.getOutput(); + ASSERT_NE(completedPartition, nullptr); + EXPECT_EQ(completedPartition->size(), 3); + EXPECT_EQ(window.getOutput(), nullptr); + EXPECT_TRUE(window.needsInput()); + + window.noMoreInput(); + auto finalPartition = window.getOutput(); + ASSERT_NE(finalPartition, nullptr); + EXPECT_EQ(finalPartition->size(), 1); + EXPECT_TRUE(window.isFinished()); + EXPECT_EQ(window.getOutput(), nullptr); + window.close(); + values.close(); +} + +TEST_F(AdapterOperatorTest, streamingFullPartitionCountSpillsActivePartition) { + cudf_velox::CudfWindow::testingSetStreamingMemoryLimits(1, 1); + SCOPE_EXIT { + cudf_velox::CudfWindow::testingResetStreamingMemoryLimits(); + }; + + std::vector data{ + makeRowVector( + {"k", "payload"}, + {makeFlatVector({1, 1}), makeFlatVector({0, 1})}), + makeRowVector( + {"k", "payload"}, + {makeFlatVector({1, 1}), makeFlatVector({2, 3})}), + makeRowVector( + {"k", "payload"}, + {makeFlatVector({2, 2}), makeFlatVector({4, 5})})}; + auto plan = withRowsFrame( + PlanBuilder() + .values(data) + .streamingWindow( + {"count(1) over (partition by k rows between unbounded preceding " + "and unbounded following) as c"}) + .planNode()); + auto expected = makeRowVector( + {"k", "payload", "c"}, + {makeFlatVector({1, 1, 1, 1, 2, 2}), + makeFlatVector({0, 1, 2, 3, 4, 5}), + makeFlatVector({4, 4, 4, 4, 2, 2})}); + + auto task = AssertQueryBuilder(plan).maxDrivers(1).assertResults(expected); + EXPECT_TRUE(wasCudfWindowUsed(task)); + EXPECT_GE(cudf_velox::CudfWindow::testingStreamingSpillWrites(), 3); + EXPECT_EQ(cudf_velox::CudfWindow::testingStreamingSpillCleanups(), 1); +} + +TEST_F(AdapterOperatorTest, streamingWindowFeatureGapsFailClosed) { + auto data = makeRowVector( + {"k", "o", "v"}, + {makeFlatVector({1, 1}), + makeFlatVector({0, 1}), + makeNullableFlatVector({std::nullopt, 2})}); + + auto nullableCount = withRowsFrame( + PlanBuilder() + .values({data}) + .streamingWindow( + {"count(v) over (partition by k rows between unbounded preceding " + "and unbounded following) as c"}) + .planNode()); + EXPECT_FALSE(cudf_velox::isSupportedCudfWindowNode(nullableCount)); + + auto boundedRange = std::dynamic_pointer_cast( + PlanBuilder() + .values({data}) + .streamingWindow( + {"sum(v) over (partition by k order by o range between 1 " + "preceding and current row) as s"}) + .planNode()); + ASSERT_NE(boundedRange, nullptr); + EXPECT_FALSE(cudf_velox::isSupportedCudfWindowNode(boundedRange)); + + auto mixedFrames = std::dynamic_pointer_cast( + PlanBuilder() + .values({data}) + .streamingWindow( + {"sum(v) over (partition by k order by o range between " + "unbounded preceding and current row) as range_s", + "sum(v) over (partition by k order by o rows between " + "unbounded preceding and current row) as rows_s"}) + .planNode()); + ASSERT_NE(mixedFrames, nullptr); + EXPECT_FALSE(cudf_velox::isSupportedCudfWindowNode(mixedFrames)); + + auto unsortedRange = std::dynamic_pointer_cast( + PlanBuilder() + .values({data}) + .window({"sum(v) over (partition by k order by o range between " + "unbounded preceding and current row) as s"}) + .planNode()); + ASSERT_NE(unsortedRange, nullptr); + EXPECT_FALSE(cudf_velox::isSupportedCudfWindowNode(unsortedRange)); +} + TEST_F(AdapterOperatorTest, streamingRowNumberCrossesInputBatches) { std::vector data{ makeRowVector( {"k", "o", "payload"}, - {makeNullableFlatVector( - {std::nullopt, std::nullopt, 1, 1}), + {makeNullableFlatVector({std::nullopt, std::nullopt, 1, 1}), makeNullableFlatVector({std::nullopt, 0, 0, 1}), makeFlatVector({0, 1, 2, 3})}), makeRowVector( @@ -159,8 +351,7 @@ TEST_F(AdapterOperatorTest, streamingRankFixesContinuedTieAndLaterPeer) { makeRowVector( {"k", "o", "payload"}, {makeFlatVector({1, 1, 1, 1}), - makeNullableFlatVector( - {std::nullopt, std::nullopt, 1, 1}), + makeNullableFlatVector({std::nullopt, std::nullopt, 1, 1}), makeFlatVector({0, 1, 2, 3})}), makeRowVector( {"k", "o", "payload"}, @@ -198,21 +389,17 @@ TEST_F(AdapterOperatorTest, streamingRankProducesOutputBeforeNoMoreInput) { std::vector data{ makeRowVector( {"k", "o"}, - {makeFlatVector({1, 1}), - makeFlatVector({0, 1})}), + {makeFlatVector({1, 1}), makeFlatVector({0, 1})}), makeRowVector( {"k", "o"}, - {makeFlatVector({1, 1}), - makeFlatVector({2, 3})})}; + {makeFlatVector({1, 1}), makeFlatVector({2, 3})})}; auto plan = PlanBuilder() .values(data) - .streamingWindow( - {"rank() over (partition by k order by o) as rnk"}) + .streamingWindow({"rank() over (partition by k order by o) as rnk"}) .planNode(); - auto windowNode = - std::dynamic_pointer_cast(plan); + auto windowNode = std::dynamic_pointer_cast(plan); ASSERT_NE(windowNode, nullptr); auto valuesNode = std::dynamic_pointer_cast( windowNode->sources()[0]); @@ -224,8 +411,8 @@ TEST_F(AdapterOperatorTest, streamingRankProducesOutputBeforeNoMoreInput) { 0, core::QueryCtx::create(driverExecutor_.get()), Task::ExecutionMode::kParallel); - auto driver = Driver::testingCreate(std::make_unique( - task, 0, 0, kUngroupedGroupId, 0)); + auto driver = Driver::testingCreate( + std::make_unique(task, 0, 0, kUngroupedGroupId, 0)); cudf_velox::CudfValues values(0, driver->driverCtx(), valuesNode); cudf_velox::CudfWindow window(1, driver->driverCtx(), windowNode); values.initialize(); @@ -325,18 +512,16 @@ TEST_F(AdapterOperatorTest, streamingRankDoesNotRetainGiantPartition) { {"k", "o"}, {makeFlatVector( kRowsPerBatch, [](vector_size_t) { return 7; }), - makeFlatVector( - kRowsPerBatch, [batch](vector_size_t row) { - return static_cast(batch) * kRowsPerBatch + row; - })})); + makeFlatVector(kRowsPerBatch, [batch](vector_size_t row) { + return static_cast(batch) * kRowsPerBatch + row; + })})); } createDuckDbTable(data); auto plan = PlanBuilder() .values(data) - .streamingWindow( - {"rank() over (partition by k order by o) as rnk"}) + .streamingWindow({"rank() over (partition by k order by o) as rnk"}) .planNode(); auto task = assertQueryOrdered( plan, @@ -346,6 +531,169 @@ TEST_F(AdapterOperatorTest, streamingRankDoesNotRetainGiantPartition) { EXPECT_TRUE(wasCudfWindowUsed(task)); } +TEST_F(AdapterOperatorTest, streamingMultiKeyRangeSumCrossesPeersAndNulls) { + std::vector data{ + makeRowVector( + {"k", "a", "b", "v", "payload"}, + {makeFlatVector({1, 1}), + makeNullableFlatVector({3, 3}), + makeNullableFlatVector({std::nullopt, std::nullopt}), + makeNullableFlatVector({std::nullopt, 5}), + makeFlatVector({0, 1})}), + makeRowVector( + {"k", "a", "b", "v", "payload"}, + {makeFlatVector({1, 1, 1, 2}), + makeNullableFlatVector({3, 3, 2, 5}), + makeNullableFlatVector({std::nullopt, 1, 0, 1}), + makeNullableFlatVector({7, std::nullopt, 2, std::nullopt}), + makeFlatVector({2, 3, 4, 5})}), + makeRowVector( + {"k", "a", "b", "v", "payload"}, + {makeFlatVector({2, 2, 2, 2, 3, 3}), + makeNullableFlatVector( + {5, 4, 4, 4, std::nullopt, std::nullopt}), + makeNullableFlatVector( + {1, std::nullopt, std::nullopt, 2, std::nullopt, 1}), + makeNullableFlatVector( + {4, std::nullopt, 6, -1, std::nullopt, 8}), + makeFlatVector({6, 7, 8, 9, 10, 11})})}; + createDuckDbTable(data); + + auto plan = + PlanBuilder() + .values(data) + .streamingWindow( + {"sum(v) over (partition by k order by a desc nulls last, b asc " + "nulls first range between unbounded preceding and current row) " + "as s"}) + .planNode(); + auto task = AssertQueryBuilder(duckDbQueryRunner_) + .config("cudf.enabled", true) + .plan(plan) + .assertResults( + "SELECT k, a, b, v, payload, sum(v) over (partition by k " + "order by a desc nulls last, b asc nulls first range " + "between unbounded preceding and current row) FROM tmp"); + EXPECT_TRUE(wasCudfWindowUsed(task)); +} + +TEST_F(AdapterOperatorTest, streamingRangeSumProducesOutputBeforeNoMoreInput) { + std::vector data{ + makeRowVector( + {"k", "a", "b", "v"}, + {makeFlatVector({1, 1}), + makeFlatVector({3, 3}), + makeNullableFlatVector({std::nullopt, std::nullopt}), + makeNullableFlatVector({std::nullopt, 5})}), + makeRowVector( + {"k", "a", "b", "v"}, + {makeFlatVector({1, 1}), + makeFlatVector({3, 2}), + makeNullableFlatVector({std::nullopt, 0}), + makeNullableFlatVector({7, 2})})}; + auto plan = + PlanBuilder() + .values(data) + .streamingWindow( + {"sum(v) over (partition by k order by a desc nulls last, b asc " + "nulls first range between unbounded preceding and current row) " + "as s"}) + .planNode(); + auto windowNode = std::dynamic_pointer_cast(plan); + ASSERT_NE(windowNode, nullptr); + auto valuesNode = std::dynamic_pointer_cast( + windowNode->sources()[0]); + ASSERT_NE(valuesNode, nullptr); + + auto task = Task::create( + "streaming-range-sum-produces-output-before-no-more-input", + core::PlanFragment{plan}, + 0, + core::QueryCtx::create(driverExecutor_.get()), + Task::ExecutionMode::kParallel); + auto driver = Driver::testingCreate( + std::make_unique(task, 0, 0, kUngroupedGroupId, 0)); + cudf_velox::CudfValues values(0, driver->driverCtx(), valuesNode); + cudf_velox::CudfWindow window(1, driver->driverCtx(), windowNode); + values.initialize(); + window.initialize(); + + auto first = values.getOutput(); + ASSERT_NE(first, nullptr); + window.addInput(std::move(first)); + EXPECT_TRUE(window.needsInput()); + EXPECT_EQ(window.getOutput(), nullptr); + + auto second = values.getOutput(); + ASSERT_NE(second, nullptr); + window.addInput(std::move(second)); + EXPECT_FALSE(window.needsInput()); + auto completedPeer = window.getOutput(); + ASSERT_NE(completedPeer, nullptr); + EXPECT_EQ(completedPeer->size(), 3); + EXPECT_EQ(window.getOutput(), nullptr); + EXPECT_TRUE(window.needsInput()); + + window.noMoreInput(); + auto finalPeer = window.getOutput(); + ASSERT_NE(finalPeer, nullptr); + EXPECT_EQ(finalPeer->size(), 1); + EXPECT_TRUE(window.isFinished()); + EXPECT_EQ(window.getOutput(), nullptr); + window.close(); + values.close(); +} + +TEST_F(AdapterOperatorTest, streamingRangeSumSpillsActivePeer) { + cudf_velox::CudfWindow::testingSetStreamingMemoryLimits(1, 1); + SCOPE_EXIT { + cudf_velox::CudfWindow::testingResetStreamingMemoryLimits(); + }; + + std::vector data{ + makeRowVector( + {"k", "a", "b", "v", "payload"}, + {makeFlatVector({1, 1}), + makeFlatVector({10, 10}), + makeFlatVector({5, 5}), + makeNullableFlatVector({1, 2}), + makeFlatVector({0, 1})}), + makeRowVector( + {"k", "a", "b", "v", "payload"}, + {makeFlatVector({1, 1}), + makeFlatVector({10, 10}), + makeFlatVector({5, 5}), + makeNullableFlatVector({std::nullopt, 3}), + makeFlatVector({2, 3})}), + makeRowVector( + {"k", "a", "b", "v", "payload"}, + {makeFlatVector({1}), + makeFlatVector({9}), + makeFlatVector({0}), + makeNullableFlatVector({4}), + makeFlatVector({4})})}; + auto plan = + PlanBuilder() + .values(data) + .streamingWindow( + {"sum(v) over (partition by k order by a desc, b desc range " + "between unbounded preceding and current row) as s"}) + .planNode(); + auto expected = makeRowVector( + {"k", "a", "b", "v", "payload", "s"}, + {makeFlatVector({1, 1, 1, 1, 1}), + makeFlatVector({10, 10, 10, 10, 9}), + makeFlatVector({5, 5, 5, 5, 0}), + makeNullableFlatVector({1, 2, std::nullopt, 3, 4}), + makeFlatVector({0, 1, 2, 3, 4}), + makeFlatVector({6, 6, 6, 6, 10})}); + + auto task = AssertQueryBuilder(plan).maxDrivers(1).assertResults(expected); + EXPECT_TRUE(wasCudfWindowUsed(task)); + EXPECT_GE(cudf_velox::CudfWindow::testingStreamingSpillWrites(), 3); + EXPECT_EQ(cudf_velox::CudfWindow::testingStreamingSpillCleanups(), 1); +} + TEST_F(AdapterOperatorTest, orderedFirstValueUsesCudfWindow) { auto data = makeRowVector( {"k", "v", "o0", "o1"}, diff --git a/velox/experimental/cudf/tests/OrderByTest.cpp b/velox/experimental/cudf/tests/OrderByTest.cpp index d2a4a912fb6..9721e32501d 100644 --- a/velox/experimental/cudf/tests/OrderByTest.cpp +++ b/velox/experimental/cudf/tests/OrderByTest.cpp @@ -243,9 +243,10 @@ TEST_F(OrderByTest, externalSpillSchemaEligibility) { config.timestampUnit = cudf::type_id::TIMESTAMP_SECONDS; EXPECT_FALSE(cudf_velox::CudfOrderBy::isSupported(timestampOrderBy)); - for (const auto unit : {cudf::type_id::TIMESTAMP_MILLISECONDS, - cudf::type_id::TIMESTAMP_MICROSECONDS, - cudf::type_id::TIMESTAMP_NANOSECONDS}) { + for (const auto unit : + {cudf::type_id::TIMESTAMP_MILLISECONDS, + cudf::type_id::TIMESTAMP_MICROSECONDS, + cudf::type_id::TIMESTAMP_NANOSECONDS}) { config.timestampUnit = unit; EXPECT_TRUE(cudf_velox::CudfOrderBy::isSupported(timestampOrderBy)); } @@ -253,24 +254,42 @@ TEST_F(OrderByTest, externalSpillSchemaEligibility) { const auto safeNestedPayload = std::dynamic_pointer_cast( PlanBuilder() - .tableScan(ROW( - {"key", "payload"}, - {INTEGER(), ARRAY(ROW({BIGINT(), VARCHAR()}))})) + .tableScan( + ROW({"key", "payload"}, + {INTEGER(), ARRAY(ROW({BIGINT(), VARCHAR()}))})) .orderBy({"key ASC NULLS LAST"}, false) .planNode()); ASSERT_NE(safeNestedPayload, nullptr); EXPECT_TRUE(cudf_velox::CudfOrderBy::isSupported(safeNestedPayload)); + const auto safeMapPayload = + std::dynamic_pointer_cast( + PlanBuilder() + .tableScan(ROW( + {"key", "payload"}, {INTEGER(), MAP(VARCHAR(), VARCHAR())})) + .orderBy({"key ASC NULLS LAST"}, false) + .planNode()); + ASSERT_NE(safeMapPayload, nullptr); + EXPECT_TRUE(cudf_velox::CudfOrderBy::isSupported(safeMapPayload)); + + const auto unsupportedMapKey = + std::dynamic_pointer_cast( + PlanBuilder() + .tableScan(ROW( + {"key", "payload"}, {MAP(VARCHAR(), VARCHAR()), INTEGER()})) + .orderBy({"key ASC NULLS LAST"}, false) + .planNode()); + ASSERT_NE(unsupportedMapKey, nullptr); + EXPECT_FALSE(cudf_velox::CudfOrderBy::isSupported(unsupportedMapKey)); + const auto unsupportedBinaryPayload = std::dynamic_pointer_cast( PlanBuilder() - .tableScan( - ROW({"key", "payload"}, {INTEGER(), VARBINARY()})) + .tableScan(ROW({"key", "payload"}, {INTEGER(), VARBINARY()})) .orderBy({"key ASC NULLS LAST"}, false) .planNode()); ASSERT_NE(unsupportedBinaryPayload, nullptr); - EXPECT_FALSE( - cudf_velox::CudfOrderBy::isSupported(unsupportedBinaryPayload)); + EXPECT_FALSE(cudf_velox::CudfOrderBy::isSupported(unsupportedBinaryPayload)); } TEST_F(OrderByTest, externalSpillConstructorDefense) { @@ -474,11 +493,11 @@ TEST_F(OrderByTest, boundedExternalSortManyRuns) { nullEvery(29, (run * 2) % 29)); // A unique final key makes the expected global order deterministic even // though SQL does not define insertion order for complete sort-key ties. - auto c2 = makeFlatVector( - kRowsPerRun, - [run](vector_size_t row) { return run * kRowsPerRun + row; }); - auto payload = makeFlatVector( - kRowsPerRun, [run](vector_size_t row) { + auto c2 = makeFlatVector(kRowsPerRun, [run](vector_size_t row) { + return run * kRowsPerRun + row; + }); + auto payload = + makeFlatVector(kRowsPerRun, [run](vector_size_t row) { return fmt::format("run={};row={};", run, row) + std::string(512, static_cast('a' + run % 26)); }); @@ -486,14 +505,13 @@ TEST_F(OrderByTest, boundedExternalSortManyRuns) { } createDuckDbTable(vectors); - auto plan = PlanBuilder() - .values(vectors) - .orderBy( - {"c0 ASC NULLS FIRST", - "c1 DESC NULLS LAST", - "c2 ASC NULLS LAST"}, - false) - .planNode(); + auto plan = + PlanBuilder() + .values(vectors) + .orderBy( + {"c0 ASC NULLS FIRST", "c1 DESC NULLS LAST", "c2 ASC NULLS LAST"}, + false) + .planNode(); assertQueryOrdered( plan, "SELECT * FROM tmp ORDER BY c0 ASC NULLS FIRST, " @@ -509,6 +527,44 @@ TEST_F(OrderByTest, boundedExternalSortManyRuns) { EXPECT_EQ(OrderByTestHelper::spillCleanups(), 1); } +TEST_F(OrderByTest, boundedExternalSortMapPayload) { + constexpr int32_t kNumRuns = 7; + constexpr vector_size_t kRowsPerRun = 31; + OrderByTestHelper::setMemoryLimits(1, 2048, 2048, 13); + + std::vector vectors; + vectors.reserve(kNumRuns); + for (int32_t run = 0; run < kNumRuns; ++run) { + auto key = makeFlatVector(kRowsPerRun, [run](vector_size_t row) { + return run * kRowsPerRun + row; + }); + auto map = makeMapVector( + kRowsPerRun, + [](vector_size_t row) { return row % 4; }, + [run](vector_size_t index) { + return fmt::format("key-{}-{}", run, index); + }, + [run](vector_size_t index) { + return fmt::format("value-{}-{}", run, index); + }, + [run](vector_size_t row) { return (run + row) % 17 == 0; }, + [run](vector_size_t index) { return (run + index) % 11 == 0; }); + vectors.push_back(makeRowVector({key, map})); + } + createDuckDbTable(vectors); + + auto plan = PlanBuilder() + .values(vectors) + .orderBy({"c0 DESC NULLS LAST"}, false) + .planNode(); + assertQueryOrdered( + plan, "SELECT * FROM tmp ORDER BY c0 DESC NULLS LAST", {0}); + + EXPECT_GE(OrderByTestHelper::sourceChunks(), kNumRuns); + EXPECT_GT(OrderByTestHelper::mergeOutputBatches(), 0); + EXPECT_EQ(OrderByTestHelper::spillCleanups(), 1); +} + /// Verifies output batch rows of OrderBy TEST_F(OrderByTest, outputBatchRows) { struct {