fix: Allow companion functions when result type is not resolvable given intermediate type

[rui-mo]

The registrations of partial and merge companion functions does not require the
result type is resolvable given intermediate type. This PR removes the
limitations for them.The registration of merge_extract companion function used
to depend on the resolving, and this PR uses the passed-in result type rather
than the resolved from intermediate type.

[netlify]

✅ Deploy Preview for meta-velox canceled.

Name	Link
🔨 Latest commit	3d2f9e5
🔍 Latest deploy log	https://app.netlify.com/sites/meta-velox/deploys/677f42032b1ccf0008a6a673

[zhztheplayer]

@rui-mo Thank you for taking this on. I remember we needed to bypass this check as well to make the decimal aggregate functions work. Or we don't need that change anymore?

[rui-mo]

@zhztheplayer Thanks for the pointer! In Gluten we expect that all the names will end with "_merge_extract," however according to this reasoning, the names may have other suffixes. For instance, we would obtain the names for sum and average as follows. I'll investigate further to see how to incorporate this section with Gluten.

avg_merge_extract_DECIMAL, avg_merge_extract_double
sum_merge_extract_bigint, sum_merge_extract_DECIMAL, sum_merge_extract_double

[zhztheplayer]

@zhztheplayer Thanks for the pointer! In Gluten we expect that all the names will end with "_merge_extract," however according to this reasoning, the names may have other suffixes. For instance, we would obtain the names for sum and average as follows. I'll investigate further to see how to incorporate this section with Gluten.

avg_merge_extract_DECIMAL, avg_merge_extract_double
sum_merge_extract_bigint, sum_merge_extract_DECIMAL, sum_merge_extract_double

Sounds reasonable. Look forward to a solution here. Thanks!

[kagamiori]

Hi @rui-mo, it's not universally correct to use resultType here. The reason is that resultType is an argument received by the factory of the merge-extract-function (i.e., the lambda starting at line 337). This factory is called in the HashAggregation constructor for individual aggregation nodes that can be the partial aggregation step or the intermediate aggregation step, etc. Suppose an aggregation node perform the intermediate aggregation step of the merge-extract-function, then both the argTypes and the resultType received by the factory at line 337 are the intermediate type of the original function. But when we do auto fn = func->factory(...), we're creating the original aggregation function, so the result type passed to this factory should be the result type of the original aggregation function.

(This change doesn't trigger any test error because the AggregationTestBase::testAggregationsWithCompanion() currently doesn't test the functions with the merge-extract companion function, which we should better add...)

[zhztheplayer]

Suppose an aggregation node perform the intermediate aggregation step of the merge-extract-function, then both the argTypes and the resultType received by the factory at line 337 are the intermediate type of the original function.

Can we break down the cases here? E.g., when step is single / final, we just skip the result type resolution?

[rui-mo]

Can we break down the cases here? E.g., when step is single / final, we just skip the result type resolution?

@kagamiori I apologize for the delayed response; I was on vacation. Do you believe that @zhztheplayer's suggestion above makes sense? We are proposing this change to allow more flexibility in the aggregate function registration especially for the Spark decimal average.

[NEUpanning]

@rui-mo Thank you for working on this. I'm developing a new Spark collect_list with the signature T -> varbinary -> array(T), for more details see #12023. I suppose if this PR is merged, Velox would support this signature. So, I cherry-picked this PR into my branch and ran the test. I found four issues:

1. CompanionSignatures::extractFunctionSignature still has limitations. It leads to register no extract functions.
2. While creating the merge_extract signature, Velox checks if all type variables should appear in the inputs arguments. For this signature, T does not appear in varbinary. See code link
3. exec::test::PlanBuilder::project uses resolveScalarFunctionType to resolve result type from intermediate type, but it cannot be resolved for this signature. Here is the error message:

C++ exception with description "Exception: VeloxUserError
Error Source: USER
Error Code: INVALID_ARGUMENT
Reason: Scalar function signature is not supported: spark_collect_list_extract(VARBINARY). Supported signatures: (varbinary) -> array(E).
Retriable: False
Function: resolveScalarFunctionType
File: /Users/panning/VSProjects/glutenNvelox_CE/velox2/velox/velox/parse/TypeResolver.cpp
Line: 102

4. Gluten fails to validate function type for collect_list_merge_extract, because Velox cannot use actual intermediate type to resolve return type.

  Validation failed at file:SubstraitToVeloxPlanValidator.cc, line:1055, function:validateAggRelFunctionType, reason:Validation failed for function collect_list_merge_extract resolve type in AggregateRel.

For the first and second issues, perhaps we can simply remove the limitation and check? For the fourth issue, maybe we can use actual return type to resolve return type.

cc @zhztheplayer