Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
14 changes: 13 additions & 1 deletion velox/experimental/cudf/exec/CudfAggregation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,12 @@ std::vector<ResolvedAggregateInfo> resolveAggregateInfos(
const auto isDecimalAggregate = aggregate.rawInputTypes.size() == 1 &&
aggregate.rawInputTypes[0]->isDecimal();

std::vector<core::TypedExprPtr> extraInputs;
if (aggregate.call->inputs().size() > 1) {
extraInputs.assign(
aggregate.call->inputs().begin() + 1, aggregate.call->inputs().end());
}

params.emplace_back(
companionStep,
aggregate.call->name(),
Expand All @@ -215,7 +221,8 @@ std::vector<ResolvedAggregateInfo> resolveAggregateInfos(
isCountFunctionName(aggregate.call->name())
? std::make_optional(getCountInputKind(aggregate, constants[i]))
: std::nullopt,
isDecimalAggregate);
isDecimalAggregate,
std::move(extraInputs));
}
return params;
}
Expand Down Expand Up @@ -279,6 +286,11 @@ AggregationInputChannels buildAggregationInputChannels(
} else if (
auto constant =
dynamic_cast<const core::ConstantTypedExpr*>(arg.get())) {
if (!aggInputs.empty()) {
// Extra constant arguments (bloom_filter_agg estimatedNumItems and
// numBits) are not aggregation input columns.
continue;
}
result.constants[i] = constant->toConstantVector(operatorCtx.pool());
aggInputs.push_back(fallbackChannel);
} else {
Expand Down
5 changes: 5 additions & 0 deletions velox/experimental/cudf/exec/CudfAggregation.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,13 @@
#include "velox/experimental/cudf/expression/ExpressionEvaluator.h"
#include "velox/experimental/cudf/vector/CudfVector.h"

#include "velox/core/ITypedExpr.h"
#include "velox/exec/Operator.h"
#include "velox/expression/FunctionSignature.h"

#include <optional>
#include <string_view>
#include <vector>

namespace facebook::velox::cudf_velox {

Expand Down Expand Up @@ -67,6 +69,9 @@ struct ResolvedAggregateInfo {
// Routing keys off the function family, not the physical batch type (which is
// VARBINARY/STRING on intermediate and final steps).
bool isDecimalAggregate;
// Extra call arguments after the first aggregation input, e.g. constant
// estimatedNumItems/numBits for bloom_filter_agg.
std::vector<core::TypedExprPtr> extraInputs;
};

// Parse aggregate inputs from the aggregation node and resolve companion steps,
Expand Down
Loading
Loading