Skip to content

fix(cudf): fall back for unsupported partial identity aggs #42

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
winningsix wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

fix(cudf): fall back for unsupported partial identity aggs #42

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 17 additions & 13 deletions velox/experimental/cudf/exec/CudfGroupby.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1313,19 +1313,23 @@ void CudfGroupby::initialize() {
}

if (partialIdentityAggregationEnabled_) {
VELOX_USER_CHECK(
streamingEnabled_ &&
std::all_of(
aggregators_.begin(),
aggregators_.end(),
[](const auto& aggregator) {
return aggregator->supportsPartialIdentity();
}),
"partial identity aggregation currently supports only non-constant "
"SUM, MIN, and MAX companion aggregates");
LOG(INFO) << "CUDF_GROUPBY_PARTIAL_IDENTITY node=" << diagnosticNodeId_
<< " state=enabled keys=" << groupingKeyOutputChannels_.size()
<< " aggregates=" << aggregators_.size();
const bool supportsPartialIdentity = streamingEnabled_ &&
std::all_of(
aggregators_.begin(),
aggregators_.end(),
[](const auto& aggregator) {
return aggregator->supportsPartialIdentity();
});
if (!supportsPartialIdentity) {
LOG(WARNING) << "CUDF_GROUPBY_PARTIAL_IDENTITY node="
<< diagnosticNodeId_
<< " state=disabled reason=unsupported_aggregation_shape";
partialIdentityAggregationEnabled_ = false;
} else {
LOG(INFO) << "CUDF_GROUPBY_PARTIAL_IDENTITY node=" << diagnosticNodeId_
<< " state=enabled keys=" << groupingKeyOutputChannels_.size()
<< " aggregates=" << aggregators_.size();
}
}

if (deviceMemoryDiagnosticsEnabled()) {
Expand Down
2 changes: 1 addition & 1 deletion velox/experimental/cudf/exec/CudfGroupby.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ class CudfGroupby : public CudfOperatorBase {
// compaction overrides that suffix with the intermediate step.
bool streamingEnabled_{true};
const int32_t groupbyStreamingMaxDistinctKeys_;
const bool partialIdentityAggregationEnabled_;
bool partialIdentityAggregationEnabled_;
const int64_t maxPartialAggregationMemoryUsage_;
int64_t numInputRows_ = 0;

Expand Down
20 changes: 20 additions & 0 deletions velox/experimental/cudf/tests/AggregationTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1019,6 +1019,26 @@ TEST_F(AggregationTest, partialAggregationMemoryLimit) {
.customStats.count("flushRowCount"));
}

TEST_F(AggregationTest, partialIdentityAggregationFallsBackForUnsupportedAggs) {
auto data = makeRowVector({
makeFlatVector<int64_t>({1, 1, 2, 2}),
makeFlatVector<int64_t>({10, 20, 30, 40}),
});
createDuckDbTable({data});

AssertQueryBuilder(duckDbQueryRunner_)
.config(
cudf_velox::CudfConfig::kCudfPartialIdentityAggregation, true)
.plan(
PlanBuilder()
.values({data})
.partialAggregation({"c0"}, {"sum(c1)", "count(0)"})
.finalAggregation()
.planNode())
.assertResults(
"SELECT c0, sum(c1), count(*) FROM tmp GROUP BY c0");
}

TEST_F(AggregationTest, partialAggregationUsesBalancedRunMerges) {
std::vector<RowVectorPtr> vectors;
constexpr int32_t kBatches = 8;
Expand Down
Loading