[SPARK-46625] CTE with Identifier clause as reference

[nebojsa-db]

What changes were proposed in this pull request?

DECLARE agg = 'max';
DECLARE col = 'c1';
DECLARE tab = 'T';

WITH S(c1, c2) AS (VALUES(1, 2), (2, 3)),
T(c1, c2) AS (VALUES ('a', 'b'), ('c', 'd'))
SELECT IDENTIFIER(agg)(IDENTIFIER(col)) FROM IDENTIFIER(tab);

-- OR

WITH S(c1, c2) AS (VALUES(1, 2), (2, 3)),
T(c1, c2) AS (VALUES ('a', 'b'), ('c', 'd'))
SELECT IDENTIFIER('max')(IDENTIFIER('c1')) FROM IDENTIFIER('T');

Currently we don't support Identifier clause as part of CTE reference.

Why are the changes needed?

Adding support for Identifier clause as part of CTE reference for both constant string expressions and session variables.

Does this PR introduce any user-facing change?

It contains user facing changes in sense that identifier clause as cte reference will now be supported.

How was this patch tested?

Added tests as part of this PR.

Was this patch authored or co-authored using generative AI tooling?

No.

[cloud-fan]

I'm not a big fan of this approach, as this duplicates the handling of IDENTIFIER clauses in CTESubstitution.

IMO, the root cause is we special-case CTE resolution and run CTESubstitution as an individual batch at the very beginning. The ideal solution is to look up CTE relations together with the normal table lookup.

My idea: let's split CTE resolution into two steps:

1. identify the available CTE relations for each UnresolvedRelation. Given the position of UnresolvedRelation, the available CTE relations can be very different (e.g. in the main query, in the CTE relations, in nested CTE, etc.). Then we wrap UnresolvedRelation with a new node WithCTERelations to hold available CTE relations.
2. In the analyzer main batch, we wait for the IDENTIFIER clause to be handled, then unwrap WithCTERelations by looking up CTE relations and resoving UnresolvedRelation. If the lookup fails, restore to UnresolvedRelation so that normal table lookup rule can handle it later.

[nebojsa-db]

I'm not a big fan of this approach, as this duplicates the handling of IDENTIFIER clauses in CTESubstitution.

IMO, the root cause is we special-case CTE resolution and run CTESubstitution as an individual batch at the very beginning. The ideal solution is to look up CTE relations together with the normal table lookup.

My idea: let's split CTE resolution into two steps:

1. identify the available CTE relations for each UnresolvedRelation. Given the position of UnresolvedRelation, the available CTE relations can be very different (e.g. in the main query, in the CTE relations, in nested CTE, etc.). Then we wrap UnresolvedRelation with a new node WithCTERelations to hold available CTE relations.
2. In the analyzer main batch, we wait for the IDENTIFIER clause to be handled, then unwrap WithCTERelations by looking up CTE relations and resoving UnresolvedRelation. If the lookup fails, restore to UnresolvedRelation so that normal table lookup rule can handle it later.

@cloud-fan Please take a look at the pushed changes now, I've created a rough draft changes which should work with your approach (if I understood correctly). I don't have deep understanding of all possible uses of CTEs and if changing the order of these few rules could cause some major issues?

[cloud-fan]

Yes this is the approach I was talking about, LGTM!

[cloud-fan]

Thinking more about it, I think we can still return SubqueryAlias(...) as before here, which is like a shortcut as we know we should look up CTE relations first when resolving UnresolvedRelation. There is no need to delay it.

[cloud-fan]

Let's keep the rule name unchanged and only make this change here. We should also add comments to explain it

// We must look up CTE relations first when resolving `UnresolvedRelation`s, but we can't do it here
// as `PlanWithUnresolvedIdentifier` is a leaf node and may produce `UnresolvedRelation` later. Here
// we wrap it with `UnresolvedWithCTERelations` so that we can delay the CTE relations lookup after
// `PlanWithUnresolvedIdentifier` is resolved.

[cloud-fan]

Another thing is that we should guarantee the UnresolvedRelation should be resolved by CTE relations lookup before the normal table lookup path. My proposal: UnresolvedWithCTERelations should be a leaf node so that the normal table lookup rule can't transform the UnresolvedRelation inside it. ResolveIdentifierClause should handle UnresolvedWithCTERelations specially and resolve PlanWithUnresolvedIdentifier inside it.

[nebojsa-db]

I have made changes suggested in these comments, hopefully I understood correctly what you suggested, please review :)

[cloud-fan]

Thinking about it more, now we need to special-case UnresolvedWithCTERelations twice: once in ResolveReferences to resolve session variables and once in ResolveIdentifierClause to resolve identifier and look up CTE relations.

How about we make UnresolvedWithCTERelations an unary code, and only special case it once in ResolveRelations that we should look up from CTE relations for UnresolvedRelations insideUnresolvedWithCTERelations? Sorry for the back and forth!

[nebojsa-db]

No worry!
Hm, issue with that approach is that ResolveRelations is traversing the tree in bottom up manner so we will first do table lookup instead of CTE relations lookup since it will first encounter UnresolvedRelation instead of UnresolvedWithCTERelations?

[cloud-fan]

you are right, it's better to keep the bottom-up resolotion.

[cloud-fan]

thanks, merging to master!