UDAF adapter supporting companion functions of different aggregation steps

[kagamiori]

(With @mbasmanova)

We are designing a UDAF adapter to allow more flexible combination of different aggregation steps in aggregation queries. Potential usages include rolling aggregation and the support for Spark query plans for aggregation with one distinct (#4412). We discuss our work in progress in this discussion.

Use Cases

Some Velox users are interested in computing aggregation on a rolling basis, e.g., aggregate events happening within a 30-day window and update the aggregation result every day.

timestamp	event
2023-01-03 04:00:00	15
2023-01-03 18:00:00	4
2023-01-04 11:00:00	20
2023-01-05 07:00:00	16
2023-01-05 20:00:00	7
…	…

A naive approach is to first filter records by date stamps and keep only those within the last 30 days and then aggregate events of them.

SELECT key, avg(event) 
FROM table 
WHERE date(timestamp) > '<DATEID-30>'
GROUP BY key;

However, computing such rolling aggregation every day from scratch is expensive and a waste of resources because the aggregation for events within the past few days won’t change as the window slides.

Alternatively, one can pre-aggregate events every day to intermediate states and then aggregate these daily states together when computing a 30-day aggregation. In the example below, avg_partial computes intermediate states from input events, then avg_merge aggregates partial states into another intermediate state, and finally, avg_extract converts the intermediate state into an aggregation result.

INSERT INTO daily
SELECT key, ds, avg_partial(event) as state
FROM (SELECT key, date(timestamp) as ds, event FROM table)
GROUP BY key, ds;

SELECT key, avg_extract(avg_merge(state))
FROM daily
WHERE ds > '<DATEID-30>'
GROUP BY key;

This approach requires more flexible computation of different phases of an aggregation. Presto supports this by providing companion aggregate and scalar functions for a UDAF that allows aggregating raw inputs to intermediate states, combining intermediate states, and extracting the aggregation result from an intermediate state separately. But this is only supported for a small set of Presto UDAFs right now, such as approx_distinct() via HyperLogLog functions and approx_percentile() via Q-Digest and T-Digest functions.

This alternative approach is better, but it still incurs unnecessary computation when aggregating daily states from DATEID-29 to DATEID-1 because these do not change from the previous 30-day aggregation to the current one. An even better approach is to subtract the daily aggregation of the oldest day from the previous 30-day aggregation and add the newest daily aggregation to it.

Below is an example query for rolling average over 30 days. In this example, window-size states are stored in the longterm table and slide-size states are stored in the daily table.

INSERT INTO daily
SELECT key, ds, avg_partial(event) 
FROM events
GROUP BY key, ds;

WITH u AS (
    SELECT key, state, FALSE AS to_add
    FROM daily
    WHERE ds = '<DATEID-30>'
    UNION ALL
    SELECT key, state, TRUE
    FROM longterm
    WHERE ds = '<DATEID-1>'
    UNION ALL
    SELECT key, state, TRUE
    FROM daily
    WHERE ds = '<DATEID>'
)
SELECT
    key,
    avg_extract(avg_retract(positive_merged, negative_merged))
FROM (
    SELECT
        key,
        avg_merge(state) FILTER (WHERE to_add = TRUE) AS positive_merged,
        avg_merge(state) FILTER (WHERE to_add = FALSE) AS negative_merged
    FROM u
    GROUP BY 1
)

Goal

We believe this is a general use case that would benefit from a generic solution. Therefore, we propose to enhance Velox to automatically generate companion functions for all aggregate functions. Given a Velox UDAF, e.g., avg, Velox will automatically provide two companion scalar functions and two companion aggregate functions, together with their registration code.

Function Name	Type	Example	Description
{AGGR_FUNC}	Aggregate	avg	The original UDAF.
{AGGR_FUNC}_partial	Aggregate	avg_partial	raw events -> intermediate state
{AGGR_FUNC}_merge	Aggregate	avg_merge	intermediate states -> intermediate state
{AGGR_FUNC}_extract	Scalar	avg_extract1	intermediate state -> final result
{AGGR_FUNC}_retract	Scalar	avg_retract	(intermediate state a, intermediate state b) -> intermediate state (a - b)

Design

This goal has two pieces:

1. Auto-generate companion functions.
2. Auto-generate registration code for them.

For the first piece, Velox will extend the vector-based UDAF authoring interface with two new optional APIs to allow defining how raw input and slide-size state can be retracted from a window-size state. UDAF authors will write UDAFs using this vector-based UDAF interface and Velox will provide an adapter to auto-generate companion functions for it.

For the second piece, the original registerAggregationFunction() API will be extended to accept an additional boolean flag indicating whether all companion functions will be registered together.

Current Velox UDAF interface and registration

The current Velox UDAF authoring interface consists of seven core methods for aggregation that UDAF authors need to define.

Methods	Descriptions
initializeNewGroups()	How to create and initialize accumulators before aggregation starts.
addRawInput	How to add raw input data to accumulators for the corresponding groups.
addSingleGroupRawInput	Similar to addRawInput but assumes there is only one group.
extractAccumulators	How to extract a vector of intermediate states from accumulators so that intermediate states can optionally be shuffled to different nodes for subsequent aggregation steps.
addIntermediateResults	How to add intermediate states from an input vector to accumulators. The input vector may be shuffled from a different node by the previous aggregation step.
addSingleGroupIntermediateResults	Similar to addIntermediateResults but assumes there is only one group.
extractValues	How to extract a vector of final aggregation results from accumulators.

In addition to these methods, the UDAF author defines all supported signatures for this UDAF, including their input types, intermediate types, and the result types. These signatures are used to resolve data types at runtime. The author also defines a factory function that takes the input type, result type, and the current aggregation step (e.g., core::AggregationNode::Step::kPartial), and returns an instance of the UDAF that supports this use case.

More details about how to implement a Velox UDAF can be found here: https://facebookincubator.github.io/velox/develop/aggregate-functions.html.

Planned extensions

To support creation of the retract function, Velox will introduce two optional method functions to the UDAF authoring interface, namely retractRawInput and retractIntermediateResults. retractIntermediateResults will be used to generate the retract companion function, and retractRawInput will be used to optimize the computation of window functions.

virtual void retractIntermediateResults(
     char** group,
     const SelectivityVector& rows,
     const std::vector<VectorPtr>& args) {
  VELOX_NYI();
}

virtual void retractRawInput(
     char** group,
     const SelectivityVector& rows,
     const std::vector<VectorPtr>& args) {
  VELOX_NYI();
}

The registration API will register all supported companion functions by default. The code that UDAF authors write is the same as before, but the partial, merge, and extract companion functions are registered automatically.

Not all UDAFs support the retract operation, e.g., min() and max(), hence we don’t register the retract companion function by default. If a UDAF supports retracting, the author should provide an AggregateFunctionMetadata object specifying this support when registering the UDAF. An extra retract function will then be registered together.

struct AggregateFunctionMetadata {
 bool supportsRetract = false;
 bool isOrderSensitive = false;
};

/// Register an aggregate function with the specified name and signatures. If registerCompanionFunctions is true, also register companion aggregate and scalar functions with it.
bool registerAggregateFunction(
   const std::string& name,
   std::vector<std::shared_ptr<AggregateFunctionSignature>> signatures,
   AggregateFunctionFactory factory,
   AggregateFunctionMetadata metadata = {},
   bool registerCompanionFunctions = true);

An additional API will be provided for retrieving signatures of companion functions of a given UDAF. If a UDAF with the given name doesn’t exist, this API returns std::nullopt.

struct CompanionSignatureEntry {
 std::string functionName;
 std::vector<FunctionSignaturePtr> signatures;
};

enum class CompanionType : int8_t { kPartial, kMerge, kExtract, kRetract };

std::optional<
   std::unordered_map<CompanionType, std::vector<CompanionSignatureEntry>>>
getCompanionFunctionSignatures(const std::string& name);

Known limitations

There are two known limitations with the current design. However, we expect most UDAFs to not have these issues or can avoid these issues through simple twists of their implementations.

1. If a signature of the original UDAF has type variables (e.g., T), its result type must be resolvable given solely the concrete intermediate type. E.g., T -> Varbinary -> T is not allowed, but T -> array(T) -> T is allowed.
2. If the original UDAF has multiple signatures with the same intermediate type but different result types, multiple extract functions will be generated, one for each distinct result type. The extract function names will be {AGGR_NAME}_extract_{RESULT_TYPE_NAME}. E.g., for avg(), there will be avg_extract_real() and avg_extract_double(). For a function in this situation that returns a complex type, e.g., Map<Varchar, Array>, we append the preorder traversal of the complex result type as the suffix to the extract function name, e.g., {AGGR_NAME}_extract_map_varchar_array_bigint.

Generate companion functions

The table below shows how the companion functions are composed of the member functions from the original UDAF. Let’s assume the original UDAF is avg() that is implemented through a class AverageAggregate.

Aggregate companion functions

Methods	avg_partial(arg)	avg_merge(arg)
initializeNewGroups	AverageAggregate::initializeNewGroups	AverageAggregate::initializeNewGroups
addRawInput	AverageAggregate::addRawInput	AverageAggregate::addIntermediateResults
addSingleGroupRawInput	AverageAggregate::addSingleGroupRawInput	AverageAggregate::addSingleGroupIntermediateResults
extractAccumulators	AverageAggregate::extractAccumulator	AverageAggregate::extractAccumulator
addIntermediateResults	AverageAggregate::addIntermediateResults	AverageAggregate::addIntermediateResults
addSingleGroupIntermediateResults	AverageAggregate::addSingleGroupIntermediateResults	AverageAggregate::addSingleGroupIntermediateResults
extractValues	AverageAggregate::extractAccumulator	AverageAggregate::extractAccumulator

Scalar companion function

Functions	Implementations
avg_retract(arg0, arg1)	char** groups = allocateGroups(allocator); AverageAggregate::initializeNewGroups(groups, range); AverageAggregate::addIntermediateResults(groups, rows, {args[0]}); AverageAggregate::retractIntermediateResults(groups, rows, {args[1]}); AverageAggregate::extractAccumulators(groups, rows.size(), args); freeGroups(groups);
avg_extract(arg)	char** groups = allocateGroups(allocator); AverageAggregate::initializeNewGroups(groups, range); AverageAggregate::addIntermediateResults(groups, rows, args); AverageAggregate::extractValues(groups, rows.size(), result); freeGroups(groups);

To generate these UDFs and UDAFs, Velox will build an adapter that takes the original UDAF and creates all the companion functions. A prototype of the adapter can be found here:

velox/velox/exec/AggregateFunctionAdapter.h

Lines 27 to 293 in fc41b12

	struct AggregateFunctionAdapter {
	class PartialFunction : public Aggregate {
	public:
	explicit PartialFunction(
	std::unique_ptr<Aggregate> fn,
	const TypePtr& resultType)
	: Aggregate{resultType}, fn_{std::move(fn)} {}
	void setOffsets(
	int32_t offset,
	int32_t nullByte,
	uint8_t nullMask,
	int32_t rowSizeOffset) override {
	Aggregate::setOffsets(offset, nullByte, nullMask, rowSizeOffset);
	fn_->setOffsets(offset, nullByte, nullMask, rowSizeOffset);
	}
	int32_t accumulatorFixedWidthSize() const override {
	return fn_->accumulatorFixedWidthSize();
	}
	void initializeNewGroups(
	char** groups,
	folly::Range<const vector_size_t*> indices) override {
	fn_->initializeNewGroups(groups, indices);
	}
	void addRawInput(
	char** groups,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addRawInput(groups, rows, args, mayPushdown);
	}
	void addSingleGroupRawInput(
	char* group,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addSingleGroupRawInput(group, rows, args, mayPushdown);
	}
	void addIntermediateResults(
	char** groups,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addIntermediateResults(groups, rows, args, mayPushdown);
	}
	void addSingleGroupIntermediateResults(
	char* group,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addSingleGroupIntermediateResults(group, rows, args, mayPushdown);
	}
	void extractAccumulators(
	char** groups,
	int32_t numGroups,
	VectorPtr* result) override {
	fn_->extractAccumulators(groups, numGroups, result);
	}
	void extractValues(char** groups, int32_t numGroups, VectorPtr* result)
	override {
	fn_->extractAccumulators(groups, numGroups, result);
	}
	private:
	std::unique_ptr<Aggregate> fn_;
	};
	class MergeFunction : public Aggregate {
	public:
	explicit MergeFunction(
	std::unique_ptr<Aggregate> fn,
	const TypePtr& resultType)
	: Aggregate{resultType}, fn_{std::move(fn)} {}
	void setOffsets(
	int32_t offset,
	int32_t nullByte,
	uint8_t nullMask,
	int32_t rowSizeOffset) override {
	Aggregate::setOffsets(offset, nullByte, nullMask, rowSizeOffset);
	fn_->setOffsets(offset, nullByte, nullMask, rowSizeOffset);
	}
	int32_t accumulatorFixedWidthSize() const override {
	return fn_->accumulatorFixedWidthSize();
	}
	void initializeNewGroups(
	char** groups,
	folly::Range<const vector_size_t*> indices) override {
	fn_->initializeNewGroups(groups, indices);
	}
	void addRawInput(
	char** groups,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addIntermediateResults(groups, rows, args, mayPushdown);
	}
	void addSingleGroupRawInput(
	char* group,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addSingleGroupIntermediateResults(group, rows, args, mayPushdown);
	}
	void addIntermediateResults(
	char** groups,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addIntermediateResults(groups, rows, args, mayPushdown);
	}
	void addSingleGroupIntermediateResults(
	char* group,
	const SelectivityVector& rows,
	const std::vector<VectorPtr>& args,
	bool mayPushdown) override {
	fn_->addSingleGroupIntermediateResults(group, rows, args, mayPushdown);
	}
	void extractAccumulators(
	char** groups,
	int32_t numGroups,
	VectorPtr* result) override {
	fn_->extractAccumulators(groups, numGroups, result);
	}
	void extractValues(char** groups, int32_t numGroups, VectorPtr* result)
	override {
	fn_->extractAccumulators(groups, numGroups, result);
	}
	private:
	std::unique_ptr<Aggregate> fn_;
	};
	class RetractFunction : public VectorFunction {
	public:
	explicit RetractFunction(std::unique_ptr<Aggregate> fn)
	: fn_{std::move(fn)} {}
	void apply(
	const SelectivityVector& rows,
	std::vector<VectorPtr>& args,
	const TypePtr& outputType,
	exec::EvalCtx& context,
	VectorPtr& result) const override {
	// Set up data members of fn_.
	HashStringAllocator stringAllocator{context.pool()};
	fn_->setAllocator(&stringAllocator);
	// Null byte.
	int32_t rowSizeOffset = bits::nbytes(1);
	int32_t offset = rowSizeOffset;
	offset = bits::roundUp(offset, fn_->accumulatorAlignmentSize());
	fn_->setOffsets(
	offset,
	RowContainer::nullByte(0),
	RowContainer::nullMask(0),
	rowSizeOffset);
	// Allocate groups.
	auto accumulatorsHeader =
	stringAllocator.allocate(sizeof(char*) * rows.size());
	auto accumulators = (char**)accumulatorsHeader->begin();
	std::vector<HashStringAllocator::Header*> headers;
	auto size = fn_->accumulatorFixedWidthSize();
	for (auto i = 0; i < rows.size(); ++i) {
	headers.push_back(stringAllocator.allocate(size + offset));
	accumulators[i] = headers.back()->begin();
	}
	// Perform per-row aggregation.
	VELOX_CHECK_EQ(args.size(), 2, "Expect two arguments");
	std::vector<vector_size_t> range;
	rows.applyToSelected([&](auto row) { range.push_back(row); });
	fn_->initializeNewGroups(accumulators, range);
	fn_->addIntermediateResults(accumulators, rows, {args[0]}, false);
	fn_->retractIntermediateResults(accumulators, rows, {args[1]}, false);
	if (!result) {
	result = BaseVector::create(outputType, rows.end(), context.pool());
	}
	fn_->extractAccumulators(accumulators, rows.size(), &result);
	// Free allocated space.
	for (auto i = 0; i < rows.size(); ++i) {
	stringAllocator.free(headers[i]);
	}
	stringAllocator.free(accumulatorsHeader);
	}
	private:
	std::unique_ptr<Aggregate> fn_;
	};
	class ExtractFunction : public VectorFunction {
	public:
	explicit ExtractFunction(std::unique_ptr<Aggregate> fn)
	: fn_{std::move(fn)} {}
	void apply(
	const SelectivityVector& rows,
	std::vector<VectorPtr>& args,
	const TypePtr& outputType,
	exec::EvalCtx& context,
	VectorPtr& result) const override {
	// Set up data members of fn_.
	HashStringAllocator stringAllocator{context.pool()};
	fn_->setAllocator(&stringAllocator);
	// Null byte.
	int32_t rowSizeOffset = bits::nbytes(1);
	int32_t offset = rowSizeOffset;
	offset = bits::roundUp(offset, fn_->accumulatorAlignmentSize());
	fn_->setOffsets(
	offset,
	RowContainer::nullByte(0),
	RowContainer::nullMask(0),
	rowSizeOffset);
	// Allocate groups.
	auto accumulatorsHeader =
	stringAllocator.allocate(sizeof(char*) * rows.size());
	auto accumulators = (char**)accumulatorsHeader->begin();
	std::vector<HashStringAllocator::Header*> headers;
	auto size = fn_->accumulatorFixedWidthSize();
	for (auto i = 0; i < rows.size(); ++i) {
	headers.push_back(stringAllocator.allocate(size + offset));
	accumulators[i] = headers.back()->begin();
	}
	// Perform per-row aggregation.
	std::vector<vector_size_t> range;
	rows.applyToSelected([&](auto row) { range.push_back(row); });
	fn_->initializeNewGroups(accumulators, range);
	fn_->addIntermediateResults(accumulators, rows, args, false);
	if (!result) {
	result = BaseVector::create(outputType, rows.end(), context.pool());
	}
	fn_->extractValues(accumulators, rows.size(), &result);
	// Free allocated space.
	for (auto i = 0; i < rows.size(); ++i) {
	stringAllocator.free(headers[i]);
	}
	stringAllocator.free(accumulatorsHeader);
	}
	private:
	std::unique_ptr<Aggregate> fn_;
	};
	};

Generate registration code

A UDAF author needs to write the registration code of the original UDAF and the rest registration of companion functions will be handled by a Velox adapter automatically. Specifically, the author needs to define all supported function signatures and a factory function that creates a std::unique_ptr of the original UDAF instance given the input type, result type, and the aggregation step (example)

We extend the registerAggregateFunction() API so that companion functions are registered together with the original UDAF:

velox/velox/exec/Aggregate.cpp

Lines 55 to 80 in fc41b12

	bool registerAggregateFunction(
	const std::string& name,
	std::vector<std::shared_ptr<AggregateFunctionSignature>> signatures,
	AggregateFunctionFactory factory,
	AggregateFunctionMetadata metadata,
	bool registerCompanionFunctions) {
	auto sanitizedName = sanitizeName(name);
	aggregateFunctions()[sanitizedName] = {
	signatures, metadata, std::move(factory)};
	// Register the aggregate as a window function also.
	registerAggregateWindowFunction(sanitizedName);
	// Register companion function if needed.
	if (registerCompanionFunctions) {
	RegisterAdapter::registerPartialFunction(name, signatures);
	RegisterAdapter::registerMergeFunction(name, signatures);
	RegisterAdapter::registerExtractFunction(name, signatures);
	if (metadata.supportsRetract) {
	RegisterAdapter::registerRetractFunction(name, signatures);
	}
	}
	return true;
	}

The new API for retrieving companion function signatures is added together:

velox/velox/exec/Aggregate.cpp

Lines 106 to 162 in fc41b12

	std::optional<
	std::unordered_map<CompanionType, std::vector<CompanionSignatureEntry>>>
	getCompanionFunctionSignatures(const std::string& name) {
	auto entry = getAggregateFunctionEntry(name);
	if (!entry.has_value()) {
	return std::nullopt;
	}
	auto signatures = entry.value()->signatures;
	std::unordered_map<CompanionType, std::vector<CompanionSignatureEntry>>
	companionSignatures;
	auto partialSignatures =
	CompanionSignatures::partialFunctionSignatures(signatures);
	companionSignatures.emplace(
	CompanionType::kPartial,
	std::vector<CompanionSignatureEntry>{
	{CompanionSignatures::partialFunctionName(name),
	std::vector<FunctionSignaturePtr>{
	partialSignatures.begin(), partialSignatures.end()}}});
	auto mergeSignatures =
	CompanionSignatures::mergeFunctionSignatures(signatures);
	companionSignatures.emplace(
	CompanionType::kMerge,
	std::vector<CompanionSignatureEntry>{
	{CompanionSignatures::mergeFunctionName(name),
	std::vector<FunctionSignaturePtr>{
	mergeSignatures.begin(), mergeSignatures.end()}}});
	if (entry.value()->metadata.supportsRetract) {
	companionSignatures.emplace(
	CompanionType::kRetract,
	std::vector<CompanionSignatureEntry>{
	{CompanionSignatures::retractFunctionName(name),
	CompanionSignatures::retractFunctionSignatures(signatures)}});
	}
	if (CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
	signatures)) {
	std::vector<CompanionSignatureEntry> entries;
	for (const auto& signature : signatures) {
	entries.push_back(
	{CompanionSignatures::extractFunctionNameWithSuffix(
	name, signature->returnType()),
	{CompanionSignatures::extractFunctionSignature(signature)}});
	}
	companionSignatures.emplace(CompanionType::kExtract, std::move(entries));
	} else {
	companionSignatures.emplace(
	CompanionType::kExtract,
	std::vector<CompanionSignatureEntry>{
	{CompanionSignatures::extractFunctionName(name),
	CompanionSignatures::extractFunctionSignatures(signatures)}});
	}
	return companionSignatures;
	}

The RegisterAdapter generates factories that create an instance of the original UDAF and pass it to the companion function constructors. A prototype can be found here:

velox/velox/exec/AggregateFunctionAdapter.h

Lines 295 to 480 in fc41b12

	class RegisterAdapter {
	public:
	static bool registerPartialFunction(
	const std::string& name,
	const std::vector<AggregateFunctionSignaturePtr>& originalSignatures) {
	auto signatures =
	CompanionSignatures::partialFunctionSignatures(originalSignatures);
	exec::registerAggregateFunction(
	CompanionSignatures::partialFunctionName(name),
	std::move(signatures),
	[name](
	core::AggregationNode::Step step,
	const std::vector<TypePtr>& argTypes,
	const TypePtr& resultType) -> std::unique_ptr<Aggregate> {
	if (auto func = getAggregateFunctionEntry(name)) {
	if (exec::isRawInput(step)) {
	auto fn = func.value()->factory(step, argTypes, resultType);
	return std::make_unique<
	AggregateFunctionAdapter::PartialFunction>(
	std::move(fn), resultType);
	} else {
	auto fn = func.value()->factory(
	core::AggregationNode::Step::kIntermediate,
	argTypes,
	resultType);
	return std::make_unique<
	AggregateFunctionAdapter::PartialFunction>(
	std::move(fn), argTypes[0]);
	}
	}
	VELOX_FAIL(
	"Original aggregation function {} not found: {}",
	name,
	CompanionSignatures::partialFunctionName(name));
	});
	return true;
	}
	static bool registerMergeFunction(
	const std::string& name,
	const std::vector<AggregateFunctionSignaturePtr>& originalSignatures) {
	auto signatures =
	CompanionSignatures::mergeFunctionSignatures(originalSignatures);
	exec::registerAggregateFunction(
	CompanionSignatures::mergeFunctionName(name),
	std::move(signatures),
	[name](
	core::AggregationNode::Step step,
	const std::vector<TypePtr>& argTypes,
	const TypePtr& resultType) -> std::unique_ptr<Aggregate> {
	if (auto func = getAggregateFunctionEntry(name)) {
	auto fn = func.value()->factory(
	core::AggregationNode::Step::kIntermediate,
	argTypes,
	resultType);
	return std::make_unique<AggregateFunctionAdapter::MergeFunction>(
	std::move(fn), argTypes[0]);
	}
	VELOX_FAIL(
	"Original aggregation function {} not found: {}",
	name,
	CompanionSignatures::mergeFunctionName(name));
	});
	return true;
	}
	static bool registerExtractFunctionWithSuffix(
	const std::string& originalName,
	const std::vector<AggregateFunctionSignaturePtr>& originalSignatures) {
	for (const auto& signature : originalSignatures) {
	auto extractSignature =
	CompanionSignatures::extractFunctionSignature(signature);
	auto factory = [extractSignature, originalName](
	const std::string& name,
	const std::vector<VectorFunctionArg>& inputArgs)
	-> std::shared_ptr<VectorFunction> {
	std::vector<TypePtr> argTypes{inputArgs.size()};
	std::transform(
	inputArgs.begin(),
	inputArgs.end(),
	argTypes.begin(),
	[](auto inputArg) { return inputArg.type; });
	SignatureBinder binder{*extractSignature, argTypes};
	binder.tryBind();
	auto resultType = binder.tryResolveReturnType();
	if (!resultType) {
	// TODO: limitation -- result type must be resolveable given
	// intermediate type of the original UDAF.
	VELOX_NYI();
	}
	if (auto func = getAggregateFunctionEntry(originalName)) {
	auto fn = func.value()->factory(
	core::AggregationNode::Step::kFinal, argTypes, resultType);
	return std::make_shared<AggregateFunctionAdapter::ExtractFunction>(
	std::move(fn));
	}
	return nullptr;
	};
	exec::registerStatefulVectorFunction(
	CompanionSignatures::extractFunctionNameWithSuffix(
	originalName, extractSignature->returnType()),
	{extractSignature},
	factory);
	}
	return true;
	}
	static bool registerExtractFunction(
	const std::string& originalName,
	const std::vector<AggregateFunctionSignaturePtr>& originalSignatures) {
	if (CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
	originalSignatures)) {
	return registerExtractFunctionWithSuffix(
	originalName, originalSignatures);
	}
	auto factory = [originalName](
	const std::string& name,
	const std::vector<VectorFunctionArg>& inputArgs)
	-> std::shared_ptr<VectorFunction> {
	std::vector<TypePtr> argTypes{inputArgs.size()};
	std::transform(
	inputArgs.begin(),
	inputArgs.end(),
	argTypes.begin(),
	[](auto inputArg) { return inputArg.type; });
	auto resultType = resolveVectorFunction(name, argTypes);
	if (!resultType) {
	VELOX_FAIL(
	"Result type should be resolveable given intermediate type of the original UDAF");
	}
	if (auto func = getAggregateFunctionEntry(originalName)) {
	auto fn = func.value()->factory(
	core::AggregationNode::Step::kFinal, argTypes, resultType);
	return std::make_shared<AggregateFunctionAdapter::ExtractFunction>(
	std::move(fn));
	}
	return nullptr;
	};
	exec::registerStatefulVectorFunction(
	CompanionSignatures::extractFunctionName(originalName),
	CompanionSignatures::extractFunctionSignatures(originalSignatures),
	factory);
	return true;
	}
	static bool registerRetractFunction(
	const std::string& originalName,
	const std::vector<AggregateFunctionSignaturePtr>& originalSignatures) {
	auto factory = [originalName](
	const std::string& name,
	const std::vector<VectorFunctionArg>& inputArgs)
	-> std::shared_ptr<VectorFunction> {
	VELOX_CHECK_EQ(inputArgs.size(), 2);
	std::vector<TypePtr> argTypes{inputArgs.size()};
	std::transform(
	inputArgs.begin(),
	inputArgs.end(),
	argTypes.begin(),
	[](auto inputArg) { return inputArg.type; });
	VELOX_CHECK(argTypes[0]->equivalent(*argTypes[1]));
	if (auto func = getAggregateFunctionEntry(originalName)) {
	auto fn = func.value()->factory(
	core::AggregationNode::Step::kIntermediate,
	{argTypes[0]},
	argTypes[0]);
	return std::make_shared<AggregateFunctionAdapter::RetractFunction>(
	std::move(fn));
	}
	return nullptr;
	};
	exec::registerStatefulVectorFunction(
	CompanionSignatures::retractFunctionName(originalName),
	CompanionSignatures::retractFunctionSignatures(originalSignatures),
	factory);
	return true;
	}
	};

Finally, #4489 gives an example of the aforementioned adapters together with unit tests for the auto-generated and auto-registered companion functions for avg().

Footnotes

1. The extract function name may have a suffix of the result type name if the original UDAF has multiple signatures with the same intermediate type but different result types. E.g., avg() has two signatures real -> row(double, bigint) -> real and double -> row(double, bigint) -> double. Their extract function signatures are row(double, bigint) -> real and row(double, bigint) -> double. However, Velox requires an expression result type to be infer-able from the function name and its argument types. Therefore, we append a suffix of the result type name to the extract functions to differentiate them, i.e., avg_extract_real() and avg_extract_double(). ↩

[kagamiori]

cc @mbasmanova @rui-mo

[pedroerp]

Thank for putting together the detailed design plan. +1 for adding the feature.

[laithsakka]

thanks wei for the well written design for the point:
``
If the original UDAF has multiple signatures with the same intermediate type but different result types, multiple extract functions will be generated, one for each distinct result type. The extract function names will be {AGGR_NAME}extract{RESULT_TYPE_NAME}.

I wonder if {AGGR_NAME}extract{INPUT_TYPE(s)_NAME}.
makes more sense?
basically naming the functions by its input type instead of output.

[kagamiori]

I think both are doable for us, so it depends on the use cases. The extract function takes input of the intermediate type and returns the result type of the original UDAF. So if we add the original input type into its function name (which can be irrelevant to the input and output type of the extract function itself), it might not be straightforward to understand what that type is.

Also, as @mbasmanova pointed out offline, one UDAF may have multiple inputs, so having input types in the function names may make function names too long.

[laithsakka]

would it make sense to have average_merge_final. (an aggregation that starts with partial result and up with final results with out having to call extract).

[kagamiori]

Yes, we can add this as an optimization after we have the basic functionalities to avoid the conversion such as accumulator --> intermediate result vector --> accumulator.

[laithsakka]

the document does not mention how addSingleGroupIntermediateResults is going to be created for the new agg function
@kagamiori

[kagamiori]

That would be similar to addIntermediateResults, but just call the corresponding "single group" APIs in the original UDAF. Let me add these to the table.

[laithsakka]

looks like we can also push
const std::string& name,
std::vector<std::shared_ptr> signatures,
into the meta data?

[kagamiori]

Could you elaborate a bit? Do you mean the name and signatures of the original UDAF or the companion functions? Also are you thinking about a new metadata abstraction or some existing ones?