Extract info from a Foreign Join plan node.

pyhalov · danolivo · commit 58ea474bd002 · 2022-10-14T14:32:47.000+05:00
diff --git a/expected/aqo_fdw.out b/expected/aqo_fdw.out
@@ -110,15 +110,15 @@ SELECT str FROM expln('
  JOINS: 0
 (6 rows)
 
--- TODO: Should learn on postgres_fdw nodes
+-- Should learn on postgres_fdw nodes
 SELECT str FROM expln('
   EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF, VERBOSE)
     SELECT * FROM frgn AS a, frgn AS b WHERE a.x=b.x;
 ') AS str WHERE str NOT LIKE '%Query Identifier%';
                                                   str                                                   
 --------------------------------------------------------------------------------------------------------
  Foreign Scan (actual rows=1 loops=1)
-   AQO not used
+   AQO: rows=1, error=0%
    Output: a.x, b.x
    Relations: (public.frgn a) INNER JOIN (public.frgn b)
    Remote SQL: SELECT r1.x, r2.x FROM (public.local r1 INNER JOIN public.local r2 ON (((r1.x = r2.x))))
@@ -127,6 +127,39 @@ SELECT str FROM expln('
  JOINS: 0
 (8 rows)
 
+CREATE TABLE local_a(aid int primary key, aval text);
+CREATE TABLE local_b(bid int primary key, aid int references local_a(aid), bval text);
+INSERT INTO local_a SELECT i, 'val_' || i FROM generate_series(1,100) i;
+INSERT INTO local_b SELECT i, mod((i+random()*10)::numeric, 10) + 1, 'val_' || i FROM generate_series(1,1000) i;
+ANALYZE local_a, local_b;
+CREATE FOREIGN TABLE frgn_a(aid int, aval text) SERVER loopback OPTIONS (table_name 'local_a');
+CREATE FOREIGN TABLE frgn_b(bid int, aid int, bval text) SERVER loopback OPTIONS (table_name 'local_b');
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+SELECT * from frgn_a AS a, frgn_b AS b
+WHERE a.aid = b.aid AND b.bval like 'val%';
+                  QUERY PLAN                   
+-----------------------------------------------
+ Foreign Scan (actual rows=1000 loops=1)
+   AQO not used
+   Relations: (frgn_a a) INNER JOIN (frgn_b b)
+ Using aqo: true
+ AQO mode: LEARN
+ JOINS: 0
+(6 rows)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+SELECT * from frgn_a AS a, frgn_b AS b
+WHERE a.aid = b.aid AND b.bval like 'val%';
+                  QUERY PLAN                   
+-----------------------------------------------
+ Foreign Scan (actual rows=1000 loops=1)
+   AQO: rows=1000, error=0%
+   Relations: (frgn_a a) INNER JOIN (frgn_b b)
+ Using aqo: true
+ AQO mode: LEARN
+ JOINS: 0
+(6 rows)
+
 -- TODO: Non-mergejoinable join condition.
 EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 SELECT * FROM frgn AS a, frgn AS b WHERE a.x<b.x;
@@ -147,7 +180,7 @@ SELECT str FROM expln('
                                                   str                                                   
 --------------------------------------------------------------------------------------------------------
  Foreign Scan (actual rows=0 loops=1)
-   AQO not used
+   AQO: rows=1, error=100%
    Output: a.x, b.x
    Relations: (public.frgn a) INNER JOIN (public.frgn b)
    Remote SQL: SELECT r1.x, r2.x FROM (public.local r1 INNER JOIN public.local r2 ON (((r1.x < r2.x))))
@@ -158,8 +191,12 @@ SELECT str FROM expln('
 
 DROP EXTENSION aqo CASCADE;
 DROP EXTENSION postgres_fdw CASCADE;
-NOTICE:  drop cascades to 3 other objects
+NOTICE:  drop cascades to 5 other objects
 DETAIL:  drop cascades to server loopback
 drop cascades to user mapping for public on server loopback
 drop cascades to foreign table frgn
+drop cascades to foreign table frgn_a
+drop cascades to foreign table frgn_b
 DROP TABLE local;
+DROP TABLE local_b;
+DROP TABLE local_a;
diff --git a/path_utils.c b/path_utils.c
@@ -23,6 +23,8 @@
 #include "aqo.h"
 #include "hash.h"
 
+#include "postgres_fdw.h"
+
 /*
  * Hook on creation of a plan node. We need to store AQO-specific data to
  * support learning stage.
@@ -60,6 +62,31 @@ create_aqo_plan_node()
 	return node;
 }
 
+
+/* Ensure that it's postgres_fdw's foreign server oid */
+static bool
+is_postgres_fdw_server(Oid serverid)
+{
+	ForeignServer *server;
+	ForeignDataWrapper *fdw;
+
+	if (!OidIsValid(serverid))
+		return false;
+
+	server = GetForeignServerExtended(serverid, FSV_MISSING_OK);
+	if (!server)
+		return false;
+
+	fdw = GetForeignDataWrapperExtended(server->fdwid, FDW_MISSING_OK);
+	if (!fdw || !fdw->fdwname)
+		return false;
+
+	if (strcmp(fdw->fdwname, "postgres_fdw") != 0)
+		return false;
+
+	return true;
+}
+
 /*
  * Extract an AQO node from the plan private field.
  * If no one node was found, return pointer to the default value or return NULL.
@@ -497,7 +524,8 @@ aqo_create_plan_hook(PlannerInfo *root, Path *src, Plan **dest)
 		return;
 
 	is_join_path = (src->type == T_NestPath || src->type == T_MergePath ||
-					src->type == T_HashPath);
+					src->type == T_HashPath ||
+					(src->type == T_ForeignPath && IS_JOIN_REL(src->parent)));
 
 	node = get_aqo_plan_node(plan, true);
 
@@ -513,8 +541,32 @@ aqo_create_plan_hook(PlannerInfo *root, Path *src, Plan **dest)
 
 	if (is_join_path)
 	{
-		node->clauses = aqo_get_clauses(root, ((JoinPath *) src)->joinrestrictinfo);
-		node->jointype = ((JoinPath *) src)->jointype;
+		if (IsA(src, ForeignPath))
+		{
+			PgFdwRelationInfo *fpinfo = (PgFdwRelationInfo *) src->parent->fdw_private;
+			List	*restrictclauses = NIL;
+
+			if (!fpinfo)
+				return;
+
+			/* We have to ensure that this is postgres_fdw ForeignPath */
+			if (!is_postgres_fdw_server(src->parent->serverid))
+				return;
+
+			restrictclauses = list_concat(restrictclauses, fpinfo->joinclauses);
+			restrictclauses = list_concat(restrictclauses, fpinfo->remote_conds);
+			restrictclauses = list_concat(restrictclauses, fpinfo->local_conds);
+
+			node->clauses = aqo_get_clauses(root, restrictclauses);
+			node->jointype = fpinfo->jointype;
+
+			list_free(restrictclauses);
+		}
+		else
+		{
+			node->clauses = aqo_get_clauses(root, ((JoinPath *) src)->joinrestrictinfo);
+			node->jointype = ((JoinPath *) src)->jointype;
+		}
 	}
 	else if (IsA(src, AggPath))
 	/* Aggregation node must store grouping clauses. */
diff --git a/sql/aqo_fdw.sql b/sql/aqo_fdw.sql
@@ -61,12 +61,29 @@ SELECT str FROM expln('
   SELECT * FROM frgn AS a, frgn AS b WHERE a.x=b.x;
 ') AS str WHERE str NOT LIKE '%Sort Method%';
 
--- TODO: Should learn on postgres_fdw nodes
+-- Should learn on postgres_fdw nodes
 SELECT str FROM expln('
   EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF, VERBOSE)
     SELECT * FROM frgn AS a, frgn AS b WHERE a.x=b.x;
 ') AS str WHERE str NOT LIKE '%Query Identifier%';
 
+CREATE TABLE local_a(aid int primary key, aval text);
+CREATE TABLE local_b(bid int primary key, aid int references local_a(aid), bval text);
+INSERT INTO local_a SELECT i, 'val_' || i FROM generate_series(1,100) i;
+INSERT INTO local_b SELECT i, mod((i+random()*10)::numeric, 10) + 1, 'val_' || i FROM generate_series(1,1000) i;
+ANALYZE local_a, local_b;
+
+CREATE FOREIGN TABLE frgn_a(aid int, aval text) SERVER loopback OPTIONS (table_name 'local_a');
+CREATE FOREIGN TABLE frgn_b(bid int, aid int, bval text) SERVER loopback OPTIONS (table_name 'local_b');
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+SELECT * from frgn_a AS a, frgn_b AS b
+WHERE a.aid = b.aid AND b.bval like 'val%';
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+SELECT * from frgn_a AS a, frgn_b AS b
+WHERE a.aid = b.aid AND b.bval like 'val%';
+
 -- TODO: Non-mergejoinable join condition.
 EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 SELECT * FROM frgn AS a, frgn AS b WHERE a.x<b.x;
@@ -78,4 +95,6 @@ SELECT str FROM expln('
 DROP EXTENSION aqo CASCADE;
 DROP EXTENSION postgres_fdw CASCADE;
 DROP TABLE local;
+DROP TABLE local_b;
+DROP TABLE local_a;
 
