How to declare the signature of an aggregate function with indirect dependency between intermediate type and return type?

[liujiayi771]

The decimal average aggregate function in Spark and Presto differs in many ways. Spark requires a separate implementation of its corresponding average aggregate function in Velox. However, when implementing this aggregate function, I encountered issues with declaring the signature that I couldn't resolve.

The implementation of average in Spark is here. Simply put, when the input type is DECIMAL(p, s), the intermediate data's sum type is DECIMAL(p + 10, s), and the final returned average type is DECIMAL(p + 4, s + 4).

inputType = DECIMAL(p, s)
intermediateType = ROW(DECIMAL(p + 10, s), bigint)
resultType = DECIMAL(p + 4, s + 4)

The perfect signature declaration should be like this.

exec::AggregateFunctionSignatureBuilder()
    .integerVariable("a_precision")
    .integerVariable("a_scale")
    .integerVariable("i_precision", "min(38, a_precision + 10)")
    .integerVariable("r_precision", "min(38, a_precision + 4)")
    .integerVariable("r_scale", "min(38, a_scale + 4)")
    .argumentType("DECIMAL(a_precision, a_scale)")
    .intermediateType("ROW(DECIMAL(i_precision, a_scale), bigint)")
    .returnType("DECIMAL(r_precision, r_scale)")
    .build()

But Velox does not support this type of signature declaration. There are two issues.

1. Velox does not support having new variables in intermediate types. Declaring i_precision will cause an exception to be thrown during FunctionSignature validation.

   Reason: (4 vs. 5) Some integer variables are not used
   Retriable: False
   Expression: usedVariables.size() == variables.size()
   Function: validate
   File: ../velox/expression/FunctionSignature.cpp
   Line: 168
   

   The reason for the error here is that AggregateFunctionSignature holds intermediateType_ but does not override its own validation logic, causing an issue with the check for the number of variables.

2. This way of declaring intermediate and return types will cause the companion function to not be registered. Because Velox requires that the variables used in the return type must be a subset of the variables used in the intermediate type, so that the return type can be resolved using the intermediate type. The relevant code logic is here.

   bool isResultTypeResolvableGivenIntermediateType(
       const AggregateFunctionSignaturePtr& signature) {
     auto& allVariables = signature->variables();
     if (allVariables.empty()) {
       return true;
     }
   
     std::vector<TypeSignature> intermediateType{signature->intermediateType()};
     std::vector<TypeSignature> resultType{signature->returnType()};
     auto variablesInIntermediate =
         usedTypeVariables(intermediateType, allVariables);
     auto variablesInResult = usedTypeVariables(resultType, allVariables);
   
     for (const auto& [variable, _] : variablesInResult) {
       if (!variablesInIntermediate.count(variable)) {
         return false;
       }
     }
     return true;
   }

   The logic for checking if it's resolvable here is simple, it just checks if the variables used in the return type are all included in the variables used in the intermediate type.

   But in the above signature, the return type can also be resolved by the intermediate type, although it's an indirect relationship.

   r_precision = a_precision + 4, r_scale = a_scale + 4
   i_precision = a_precision + 10, a_scale
   r_precision = i_precision - 6, a_scale + 4
   

   The framework currently does not support such indirect dependency relationships.

[liujiayi771]

Hi @mbasmanova @kagamiori. Do you have any suggestions for adapting to this situation? Is it possible to add support for this in Velox's function signature framework?
cc @rui-mo

[mbasmanova]

@liujiayi771 It might be possible to extend the framework to support something like this, but before we spend time designing the solution, can you help us understand the use case. What's the logic behind the rules for computing intermediate and final types for sum(decimal)?

when the input type is DECIMAL(p, s), the intermediate data's sum type is DECIMAL(p + 10, s), and the final returned average type is DECIMAL(p + 4, s + 4).

[liujiayi771]

@mbasmanova The issue mainly pertains to the avg(decimal) in Spark, rather than sum(decimal). When I referred to "sum type", I was specifically addressing the "sum" within the avg accumulator.

Why are the precision and scale specified in this way for sum and avg, there is no specific explanation in the source code of Spark. You can refer to these codes.
https://github.com/apache/spark/blob/6c30525e0d71c707aaaed1be49600b866bced49b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala#L69
https://github.com/apache/spark/blob/6c30525e0d71c707aaaed1be49600b866bced49b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala#L76

However, I suspect this comes from the Hive code, you can refer to these codes.
https://github.com/apache/hive/blob/ffb1165f59defa66b31b4fd9cb6367b71050071b/ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDAFAverage.java#L301
https://github.com/apache/hive/blob/ffb1165f59defa66b31b4fd9cb6367b71050071b/ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDAFAverage.java#L599

It seems to be a static setting to maintain computation precision. This is because the sum value can be quite large and does not have scale changes, thus setting sum_p = p + 10 and sum_s = s. On the other hand, avg involves division and the result does not significantly increase in value, so avg_p = p + 4 is chosen. Since division involves scale changes, avg_s = s + 4 to maintain a certain level of precision.

[mbasmanova]

@liujiayi771 Thank you for clarifying. Would you like to take a closer look to see if you can figure out how to support this use case in Velox?

[liujiayi771]

@mbasmanova I have thought of some ways to support this scenario previously, but I haven't found a solution. Is @kagamiori more familiar with this part and able to provide some feasible support solutions?

[kagamiori]

Hi @liujiayi771, I'm taking a look.

[kagamiori]

Hi @liujiayi771, I have a clarification question. The ideal signature you give above says

...
    .integerVariable("i_precision", "min(38, a_precision + 10)")
    .integerVariable("r_precision", "min(38, a_precision + 4)")

So if the companion function only knows i_precision == 38, it can infer that a_precision >= 28, but couldn't know the exact value of a_precision and hence the value of r_precision. Do you know whether this is a valid concern or this situation would not happen?

[liujiayi771]

@kagamiori I hadn't noticed this before. It is possible that when p + 10 exceeds 38, we can't calculate the resultType through the intermediateType anymore. However, both can be determined by the raw input type. Is it possible for the merge extract companion functions to record the raw input type and determine their own type based on the raw input type?

[kagamiori]

Is it possible for the merge extract companion functions to record the raw input type and determine their own type based on the raw input type?

I don't think we would be able to do this. The reason is that at function registration time, we only know the raw input type is "DECIMAL(a_precision, a_scale)" but don't know the exact values of a_precision and a_scale. We can only know the exact decimal types when creating the aggregation function instance for a concrete query. But at that moment, the raw input type we see is the raw input type to the merge extract companion function (i.e., the intermediate type of the original aggregation function), not the raw input type of the original aggregation function.

[liujiayi771]

@kagamiori Thank you for helping to explain this issue.

Without changing the existing framework, do you know of any other methods that could work around this issue and support the scenario? If it cannot be resolved for the time being, we plan to first merge a version of Spark decimal average that does not support companion functions into the community, and later on, we'll see if there is a way to support an architecture with companion functions. We would like to first integrate the main logic of the Spark decimal average into the community.

[kagamiori]

Hi @liujiayi771, I think first merging the main logic of Spark decimal average function is fine. But I think the issue with resolving the result type from intermediate type is more fundamental than what is/isn't supported by AggregateCompanionAdapter.

My understanding is that in your use case, you need to use the avg_merge_extract function. I imagine that the original function is avg(decimal(a_precision, a_scale)) -> row(decimal(i_precision, a_scale), bigint) -> decimal(r_precision, r_scale), and there is a query SELECT avg_merge_extract(c0) from .... The c0 column is actually of the type row(decimal(i_precision, a_scale), bigint). Since the avg_merge_extract function only see the c0 input column, there is no way it could know the original input type decimal(a_precision, a_scale) and therefore the final result type decimal(r_precision, r_scale).

I feel that this is not just a problem with AggregateCompanionAdapter. If it was, you would be able to work around by manually implementing the companion functions (i.e., write the individual companion functions yourself instead of relying on the auto-generation through AggregateCompanionAdapter). But with the issue above, it seems not possible to determine the result type even if you manually implement avg_merge_extract.

Could you please share a description of your use case if I misunderstand?

[liujiayi771]

Hi @kagamiori. Your understanding is completely correct. We can first merge the Spark decimal average functions that do not support AggregateCompanionAdapter into the Velox. We have some workaround solutions that can adapt to the companion functions (Spark will pass the actual correct type to Velox, not relying on Velox's signature to deduce the correct decimal precision and scale).

I'm about to complete this aggregate average function development, which uses the FunctionState you mentioned in #8711. I will need your help to review it later, thanks.

[rui-mo]

@mbasmanova @kagamiori @liujiayi771 I am looking at the Spark decimal average signature issue, and these are my observations.

The table below lists the intermediate and final values' precision and scale. I find the major problem is that the registration of companion function requires that the result type could be inferred from the intermediate type. While for this case when the intermediate precision is 38, the result precision could be in the range of [32, 38] and cannot be determined.

variable	precision	scale
input	p	s
intermediate	std::min(p + 10, 38)	s
result	std::min(p + 4, 38)	std::min(s + 4, 38)

This requirement comes from the ExtractFunction which resolves the return type from argument types and register a new vector function.

velox/velox/exec/AggregateCompanionAdapter.cpp

Lines 437 to 443 in d1ac079

	auto resultType = resolveVectorFunction(name, argTypes);
	if (!resultType) {
	// TODO: limitation -- result type must be resolveable given
	// intermediate type of the original UDAF.
	VELOX_UNREACHABLE(
	"Signatures whose result types are not resolvable given intermediate types should have been excluded.");
	}

Functions are required to provide a way to statically resolve return type from input types but it is not always feasible for an extract companion function. Given that special forms do not have that requirement and they resolve return types dynamically at runtime, I believe it would be more appropriate to register ExtractFunction as a special form and the decimal avg signature issue shall be resolved.

This is a prototype #10985 and I will proceed if you think the idea makes sense. I notice @kagamiori is solving this issue by adding function-level states in #8710 but this PR is stale. If you prefer to move forward with that solution, please kindly let me know. Thanks!

[mbasmanova]

@rui-mo

registration of companion function requires that the result type could be inferred from the intermediate type

A companion function is just a regular function, which can be used in a SQL query. Hence, it should be possible to derive result type from the input types. Here this means that result type must be resolvable from the intermediate type (intermediate type for original function == input type for companion function).

I remember Gluten is using companion functions as a workaround for some problem. Can you remind me what that is? We may need to design a proper solution.

[rui-mo]

@mbasmanova Thanks for your reply. AFAIK, two Gluten issues have been solved with the use of companion functions.

1. In the Spark count distinct aggregation, companion functions are used to disable flushing in partial or intermediate aggregations and to enable mixed-phase aggregate functions. Details can be found in issue.
2. There were some discussions on 'enable spilling support for partial aggregation', and at last we chose to use Velox's FinalAggregation + Companion functions. Details can be found in PR.

cc: @zhztheplayer

[zhztheplayer]

I remember Gluten is using companion functions as a workaround for some problem.

To add to Rui's insights, we are basically using Velox companion functions to implement Spark's aggregate expression modes. See Spark's definition of AggregateExpression:

https://github.com/apache/spark/blob/f3c8d26eb0c3fd7f77950eb08c70bb2a9ab6493c/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala#L99-L104

In Spark, the aggregate modes (namely Velox's aggregate steps) are bound to aggregate functions than to aggregate operator. See Spark's code of HashAggregateExec:

https://github.com/apache/spark/blob/f3c8d26eb0c3fd7f77950eb08c70bb2a9ab6493c/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala#L52-L62

You can see the operator doesn't have a mode as property. So technically in Spark there can be a case that one aggregate operator handles different agg funcs for different agg modes at the same time.

As Velox aligns with Presto at the beginning, so Velox's design is one agg operator should have one certain agg step. This is where Presto and Spark make difference. So our workaround is to have all Spark's aggregates mapping to Velox's final aggs + companion functions.

To optimize this part of design, my proposal is - if we can manage to do - to remove step property from Velox's HashAggregation, add the property in aggregate function calls, and add a flushable flag or something to HashAggregation. Then Presto's partial aggregate operator can still map to Velox HashAggregation(flushable=true) + partial aggregate functions without much efforts. And in Spark we can remove the uses of companion functions and directly create partial / final aggregate functions too.

[liujiayi771]

@rui-mo I have a more complete implementation in #9167. But @mbasmanova suggested is that we can do it similar in the SimpleAggregateAdapter so that the UDAF authors doesn't have to keep function-level states.