From c72d2b479b45ab7bc6d3e9ccce92dac7b90deba9 Mon Sep 17 00:00:00 2001
From: SaitTalhaNisanci <s.talhanisanci@gmail.com>
Date: Mon, 31 May 2021 21:02:20 +0300
Subject: [PATCH] Add tests for union pushdown workaround (#5005)

---
 src/test/regress/expected/union_pushdown.out | 115 ++++++++++++++++++-
 src/test/regress/sql/union_pushdown.sql      |  62 ++++++++++
 2 files changed, 176 insertions(+), 1 deletion(-)

diff --git a/src/test/regress/expected/union_pushdown.out b/src/test/regress/expected/union_pushdown.out
index 2691ff461..ec56f3949 100644
--- a/src/test/regress/expected/union_pushdown.out
+++ b/src/test/regress/expected/union_pushdown.out
@@ -859,10 +859,123 @@ $$);
  t
 (1 row)
 
+-- #4781
+CREATE TABLE test_a (id int, k int);
+CREATE TABLE test_b (id int, k int);
+SELECT create_distributed_table('test_a','id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT create_distributed_table('test_b','id');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE OR REPLACE VIEW v AS SELECT * from test_a where k>1 UNION ALL SELECT * from test_b where k<1;
+-- the followings can be pushed down since dist_key is used in the aggregation
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT COUNT(id) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ f
+(1 row)
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT AVG(id) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ f
+(1 row)
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT SUM(id) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ f
+(1 row)
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT MAX(id) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ f
+(1 row)
+
+-- cannot pushed down because postgres optimizes fields, needs to be fixed with #4781
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT COUNT(k) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ t
+(1 row)
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT AVG(k) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ t
+(1 row)
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT SUM(k) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ t
+(1 row)
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT MAX(k) FROM v;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ t
+(1 row)
+
+-- order by prevents postgres from optimizing fields so can be pushed down
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT id, COUNT(*) FROM v GROUP BY id ORDER BY id;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ f
+(1 row)
+
+-- order by is not on dist_key so can't pushed down, needs to be fixed with #4781
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT k, COUNT(*) FROM v GROUP BY k ORDER BY k;
+$$);
+ explain_has_distributed_subplan
+---------------------------------------------------------------------
+ t
+(1 row)
+
 RESET client_min_messages;
 DROP SCHEMA union_pushdown CASCADE;
-NOTICE:  drop cascades to 4 other objects
+NOTICE:  drop cascades to 7 other objects
 DETAIL:  drop cascades to table users_table_part
 drop cascades to table events_table_part
 drop cascades to table events_table_ref
 drop cascades to table events_table_local
+drop cascades to table test_a
+drop cascades to table test_b
+drop cascades to view v
diff --git a/src/test/regress/sql/union_pushdown.sql b/src/test/regress/sql/union_pushdown.sql
index 11655a701..c40f7c6fd 100644
--- a/src/test/regress/sql/union_pushdown.sql
+++ b/src/test/regress/sql/union_pushdown.sql
@@ -654,6 +654,68 @@ JOIN users_table_part USING(user_id)
 LIMIT 1;
 $$);
 
+-- #4781
+CREATE TABLE test_a (id int, k int);
+CREATE TABLE test_b (id int, k int);
+SELECT create_distributed_table('test_a','id');
+SELECT create_distributed_table('test_b','id');
+
+CREATE OR REPLACE VIEW v AS SELECT * from test_a where k>1 UNION ALL SELECT * from test_b where k<1;
+-- the followings can be pushed down since dist_key is used in the aggregation
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT COUNT(id) FROM v;
+$$);
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT AVG(id) FROM v;
+$$);
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT SUM(id) FROM v;
+$$);
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT MAX(id) FROM v;
+$$);
+
+-- cannot pushed down because postgres optimizes fields, needs to be fixed with #4781
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT COUNT(k) FROM v;
+$$);
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT AVG(k) FROM v;
+$$);
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT SUM(k) FROM v;
+$$);
+
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT MAX(k) FROM v;
+$$);
+
+-- order by prevents postgres from optimizing fields so can be pushed down
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT id, COUNT(*) FROM v GROUP BY id ORDER BY id;
+$$);
+
+-- order by is not on dist_key so can't pushed down, needs to be fixed with #4781
+SELECT public.explain_has_distributed_subplan($$
+EXPLAIN
+SELECT k, COUNT(*) FROM v GROUP BY k ORDER BY k;
+$$);
+
+
 
 RESET client_min_messages;
 DROP SCHEMA union_pushdown CASCADE;