PG17 compatibility: add/fix tests with correlated subqueries that can be pulled to a join

[colm-mchugh]

Fix Test Failure in subquery_in_where, set_operations, dml_recursive in PG17 #7741

The test failures are caused by this commit in PG17, which enables correlated subqueries to be pulled up to a join. Prior to this, the correlated subquery was implemented as a subplan. In citus, it is not possible to pushdown a correlated subplan, but with a different plan in PG17 the query can be executed, per the test diff from subquery_in_where:

37,39c37,41
< DEBUG:  generating subplan XXX_1 for CTE event_id: SELECT user_id AS events_user_id, "time" AS events_time, event_type FROM public.events_table
< DEBUG:  Plan XXX query after replacing subqueries and CTEs: SELECT count(*) AS count FROM (SELECT intermediate_result.events_user_id, intermediate_result.events_time, intermediate_result.event_type FROM read_intermediate_result('XXX_1'::text, 'binary'::citus_copy_format) intermediate_result(events_user_id integer, events_time timestamp without time zone, event_type integer)) event_id WHERE (events_user_id OPERATOR(pg_catalog.=) ANY (SELECT users_table.user_id FROM public.users_table WHERE (users_table."time" OPERATOR(pg_catalog.=) event_id.events_time)))
< ERROR:  correlated subqueries are not supported when the FROM clause contains a CTE or subquery
---
>  count
> ---------------------------------------------------------------------
>      0
> (1 row)
> 

This is because with pg17 = ANY subquery in the queries can be implemented as a join, instead of as a subplan filter on a table scan. For example, SELECT * FROM test a WHERE x IN (SELECT x FROM test b UNION SELECT y FROM test c WHERE a.x = c.x) ORDER BY 1,2 (from set_operations) has this plan in pg17; note that the subquery is the inner side of a nested loop join:

┌───────────────────────────────────────────────────┐
│                    QUERY PLAN                     │
├───────────────────────────────────────────────────┤
│ Sort                                              │
│   Sort Key: a.x, a.y                              │
│   ->  Nested Loop                                 │
│         ->  Seq Scan on test a                    │
│         ->  Subquery Scan on "ANY_subquery"       │
│               Filter: (a.x = "ANY_subquery".x)    │
│               ->  HashAggregate                   │
│                     Group Key: b.x                │
│                     ->  Append                    │
│                           ->  Seq Scan on test b  │
│                           ->  Seq Scan on test c  │
│                                 Filter: (a.x = x) │
└───────────────────────────────────────────────────┘

and this plan in pg16 (and previous pg versions); the subquery is a correlated subplan filter on a table scan:

┌───────────────────────────────────────────────┐
│                  QUERY PLAN                   │
├───────────────────────────────────────────────┤
│ Sort                                          │
│   Sort Key: a.x, a.y                          │
│   ->  Seq Scan on test a                      │
│         Filter: (SubPlan 1)                   │
│         SubPlan 1                             │
│           ->  HashAggregate                   │
│                 Group Key: b.x                │
│                 ->  Append                    │
│                       ->  Seq Scan on test b  │
│                       ->  Seq Scan on test c  │
│                             Filter: (a.x = x) │
└───────────────────────────────────────────────┘

The fix Modifies the queries causing the test failures so that an ANY subquery is not folded to a join, preserving the expected output of the tests. A similar approach was taken for existing regress tests in the postgres commit. See the join regress test, for example.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (release-13.0@0fed87a). Learn more about missing BASE report.

Additional details and impacted files

@@               Coverage Diff               @@
##             release-13.0    #7745   +/-   ##
===============================================
  Coverage                ?   89.64%           
===============================================
  Files                   ?      274           
  Lines                   ?    59583           
  Branches                ?     7436           
===============================================
  Hits                    ?    53413           
  Misses                  ?     4037           
  Partials                ?     2133           

---- 🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests

[naisila]

A similar approach was taken for existing regress tests in the postgres commit.

They followed this approach in the postgres tests because they were having EXPLAIN diffs, and they wanted to avoid adding a new alternative test output file for PG17. In Citus, note that in these two tests, we are trying to run the query, not to explain it. So, we try to run these queries, both of them unexpectedly work.

My point is, we also need to understand what changed in the Citus planner path, in the codebase, and make sure that Citus is running these queries correctly.

Current fix is great, by the way, no extra output file, but we may need to test this more extensively in Citus through this PR.

[colm-mchugh]

Got it. I think that the Citus planner is running the queries correctly (in pg17) because it is getting a different plan from the pg planner, but I will verify, and see what tests can be added (maybe to pg17 regress test?) to test the new behavior in pg17.

[naisila]

Thank you, that sounds great.

maybe to pg17 regress test

Yes, makes sense.

[colm-mchugh]

The latest push contains a pg17 regress test that tests the pg17 feature of pulling up correlated ANY subqueries. It can be extended to test other 17-related functionality as appropriate.

[naisila]

By the way, can we add a similar fix to dml_recursive test to avoid the extra output? #7727

[colm-mchugh]

By the way, can we add a similar fix to dml_recursive test to avoid the extra output? #7727

Yes, it looks like dml_recursive can have a similar fix. In all three cases - set_operations, subquery_in_where and dml_recursive - the plan created by the Postgres planner pre-pg17 implemented the correlated subquery as a SubPlan filter. In all three cases with pg17 the pg optimizer can fold the correlated subquery to a join, so the pg plan does not have any correlated SubPlans, which seems to avoid the limitations in Citus.

[naisila]

Beautiful PR, thank you.
We can merge after the team sync on the queries that work, given that we don't discover any issues in that meeting.

[colm-mchugh]

@microsoft-github-policy-service agree company="Microsoft"

[naisila]

The reason why I was holding off on merging is that I forgot to test the PR with PG16 😆Sorry about that.
So, I realized we miss pg17_0.sql file, which is the alternative file output for pg16/pg15/pg14 runs. Thats why I am requesting changes with this PR review.
Also, check-style test is failing, it looks like there are some whitespaces around https://github.com/citusdata/citus/actions/runs/11861227624/job/33058148915?pr=7745

Additionally, I really like that you provided the query version rewritten with subquery pulled up to a join, which Citus can execute in all PG versions. So, I was thinking, we can include these outputs in pg17_0.sql file
Usually pgxx_0.sql file only has the following lines as we don't execute in previous versions:

--
-- PG16
--
SHOW server_version \gset
SELECT substring(:'server_version', '\d+')::int >= 16 AS server_version_ge_16
\gset
\if :server_version_ge_16
\else
\q

However, we might let it execute in this case. What do you think?

[colm-mchugh]

The reason why I was holding off on merging is that I forgot to test the PR with PG16 😆Sorry about that. So, I realized we miss pg17_0.sql file, which is the alternative file output for pg16/pg15/pg14 runs. Thats why I am requesting changes with this PR review. Also, check-style test is failing, it looks like there are some whitespaces around https://github.com/citusdata/citus/actions/runs/11861227624/job/33058148915?pr=7745

Ah, I was not aware of the pgxx_0.sql convention, let me address, and also check-style

However, we might let it execute in this case. What do you think?

I think that's reasonable! (include queries that Citus can run with pg < pg17)

[naisila]

Looks good for merging to me, thanks! A couple of things:

• check-style is failing because of some whitespaces I think
• before merging we need to change the base PR to release-13.0

[colm-mchugh]

• check-style is failing because of some whitespaces I think

Fixed; forgot to run after changing the test table population

• before merging we need to change the base PR to release-13.0

Just want to sanity-check how change the base PR to release-13.0 is done; is it:

git checkout dev-branch
git rebase -i --onto release-13.0 naisila/pg17_support dev-branch
< drop irrelevant commits >
git push -f

?

Thanks!

[colm-mchugh]

The PR has been rebased to release-13.0, should be good to merge pending any relevant checks