For SqlColumnPrunerOptimizer - add support for CTEs in sub-queries

[plypaul]

Previously, the SqlColumnPrunerOptimizer only accounted for CTEs defined at the top level as in current use, the DataflowToSqlPlanConverter only generated SQL plans in that way. To allow for generation of SQL plans with CTEs defined in sub-queries (at any depth), this PR updates SqlColumnPrunerOptimizer to keep track of the CTE alias mapping / CTE context when traversing the SQL plan.

[courtneyholcomb]

Looks good assuming we do expect alias collisions in valid plans - but I'm wondering if the collision-handling is necessary if we can build unique aliases appropriately.

[courtneyholcomb]

Why would they be defined with the same alias? Shouldn't we use something similar to _next_unique_table_alias() to build CTE aliases? Or is the problem that both the DataflowToSqlQueryPlanConverter and the optimizer won't have access to that function to ensure unique aliases?

[plypaul]

If something like _next_unique_table_alias is used to generate the CTE aliases, then there would be no issue. However, it would be something that you would have to remember to do in DataflowToSqlQueryPlanConverter so handling this case is more defensive.

[courtneyholcomb]

Gotcha. I definitely was planning to use something similar to _next_unique_table_alias() in DataflowToSqlQueryPlanConverter for CTE aliases. If you have a strong opinion about handling this defensively, that's fine and we can leave it here. Just might be nice to simplify the logic if it's not necessary!

[plypaul]

My aim is to try to maintain the contract: if the optimizer is fed in valid SQL, then it should output valid SQL. Otherwise, it becomes harder to use / makes a trap.

[courtneyholcomb]

Ohhh ok I think I was misunderstanding. I was imagining the CTEs would always be bubbled up to the top level, even if there were defined when visiting a subquery. I didn't realize that you could have a CTE that lives inside a subquery. In that case, your logic makes sense.

[plypaul]

Yeah, I was thinking about your use case and thought CTE that lives inside a subquery was going to be a better fit since it was only locally needed? There's flexibility though, so let me merge this so that you can try it out.

[courtneyholcomb]

Yeah, either should be fine from an execution standpoint, but keeping the CTEs at the top level might be most readable!

[courtneyholcomb]

Is there any scenario where the select_node won't be unique in the dataflow plan at this point, and a valid plan will error?

[plypaul]

There's a call to cte_alias_mapping_exists before it's added, so it shouldn't throw an error. But a non-unique select_node also means that the generated SQL plan is no longer a tree?

[plypaul]

Actually, looks like some caching in DataflowNodeToSqlCteVisitor might cause select_node won't be unique. It doesn't cause issues at the moment, but let me taking a look on how to avoid that.

[courtneyholcomb]

Similar question here - why do we expect collisions?

[plypaul]

Currently, we do not. But once we start adding CTE nodes within DataflowToSqlQueryPlanConverter, there's the potential if unique aliases aren't enforced.

[courtneyholcomb]

Just checking my understanding - this works because the CTE nodes from the inner query will be added to the mapping before the nodes from the outer query?

[plypaul]

It's actually the other way around e.g. for

WITH cte_in_outer_query AS (
  SELECT 1
)

SELECT * FROM (
  WITH cte_in_inner_query AS (
    SELECT 1
  )
  SELECT * FROM cte_in_inner_query JOIN cte_in_outer_query
)

We add the mapping {"cte_in_outer_query": ..} then {cte_in_outer_query": ..., "cte_in_inner_query": ...}

[courtneyholcomb]

I'm confused then 🤔 With this conditional, won't we only add the mapping if the node isn't in the mapping yet? In which case, the outer node will override the inner node, which is the opposite of what the comment says?

[plypaul]

To clarify, the mapping refers to the whole {"cte_in_outer_query": ..., "cte_in_inner_query": ...} where as the lookup is like {"inner_select": {"cte_in_outer_query": ..., "cte_in_inner_query": ...}}. The conditional check is seeing if we have figured out the CTE-alias mapping for this SELECT statement. So for:

WITH cte_in_outer_query AS (
  SELECT 1
)

# outer_select
SELECT * FROM (
  WITH cte_in_inner_query AS (
    SELECT 1
  )
  # inner_select
  SELECT * FROM cte_in_inner_query JOIN cte_in_outer_query
)

We first visit outer_select. cte_alias_mapping_exists(outer_select) == False so we add the mapping to the lookup which then looks like:

{
  "outer_select": {"cte_in_outer_query": ...}`
}

Then we visit inner_select. cte_alias_mapping_exists(inner_select) == False as well so we add the mapping "outer_select": {"cte_in_outer_query": ...}. The lookup now looks like:

{
  "outer_select": {"cte_in_outer_query": ...},
  "inner_select": {"cte_in_outer_query": ..., "cte_in_inner_query": ...}
}

• To make this easier to follow, I've parted out the logic for building a SqlCteAliasMappingLookup to a separate class (see last commit).
• There is a issue with the way that a cache in DataflowNodeToSqlCteVisitor works right now that is related to this, but it doesn't cause errors (will have to address in a separate PR).

[courtneyholcomb]

Ohh helpful, thank you!!