Support num_segments option for foreign servers and tables

adam8157 · adam8157 · commit aa7c74cfa473 · 2021-12-21T16:30:54.000+08:00
Support specifying the num_segments option for foreign servers and
tables, Greenplum will create that amount of QEs if the `mpp_execute`
option's value is 'all segments'.

```
CREATE FOREIGN TABLE ft1 (
       c1 int
) SERVER s1 OPTIONS (mpp_execute 'all segments', num_segments '42');
```
diff --git a/src/backend/commands/foreigncmds.c b/src/backend/commands/foreigncmds.c
@@ -188,11 +188,14 @@ transformGenericOptions(Oid catalogId,
 	result = optionListToArray(resultOptions);
 
 	/*
-	 * Check and separate out the mpp_execute option, fdwvalidator doesn't have
+	 * Check and separate out the extra options, fdwvalidator doesn't have
 	 * to handle it. USER MAPPING doesn't have the mpp_execute option.
 	 */
 	if (catalogId != UserMappingRelationId)
+	{
 		SeparateOutMppExecute(&resultOptions);
+		SeparateOutNumSegments(&resultOptions);
+	}
 
 	if (OidIsValid(fdwvalidator))
 	{
diff --git a/src/backend/foreign/foreign.c b/src/backend/foreign/foreign.c
@@ -18,6 +18,7 @@
 #include "catalog/pg_foreign_server.h"
 #include "catalog/pg_foreign_table.h"
 #include "catalog/pg_user_mapping.h"
+#include "cdb/cdbutil.h"
 #include "commands/defrem.h"
 #include "foreign/fdwapi.h"
 #include "foreign/foreign.h"
@@ -74,6 +75,41 @@ SeparateOutMppExecute(List **options)
 	return exec_location;
 }
 
+/* Get and separate out the num_segments option */
+int32
+SeparateOutNumSegments(List **options)
+{
+	ListCell *lc = NULL;
+	ListCell *prev = NULL;
+	char *num_segments_str = NULL;
+	int32 num_segments = 0;
+
+	foreach(lc, *options)
+	{
+		DefElem    *def = (DefElem *) lfirst(lc);
+
+		if (strcmp(def->defname, "num_segments") == 0)
+		{
+			num_segments_str = defGetString(def);
+			num_segments = pg_atoi(num_segments_str, sizeof(int32), 0);
+
+			if (num_segments <= 0)
+			{
+				ereport(ERROR,
+						(errcode(ERRCODE_SYNTAX_ERROR),
+						 errmsg("\"%d\" is not a valid num_segments value",
+								num_segments)));
+			}
+
+			*options = list_delete_cell(*options, lc, prev);
+			break;
+		}
+
+		prev = lc;
+	}
+	return num_segments;
+}
+
 /*
  * GetForeignDataWrapper -	look up the foreign-data wrapper by OID.
  */
@@ -224,6 +260,12 @@ GetForeignServerExtended(Oid serverid, bits16 flags)
 		server->exec_location = fdw->exec_location;
 	}
 
+	server->num_segments = SeparateOutNumSegments(&server->options);
+	if (server->num_segments <= 0)
+	{
+		server->num_segments = getgpsegmentCount();
+	}
+
 	ReleaseSysCache(tp);
 
 	return server;
@@ -329,13 +371,20 @@ GetForeignTable(Oid relid)
 	else
 		ft->options = untransformRelOptions(datum);
 
+	ForeignServer *server = GetForeignServer(ft->serverid);
+
 	ft->exec_location = SeparateOutMppExecute(&ft->options);
 	if (ft->exec_location == FTEXECLOCATION_NOT_DEFINED)
 	{
-		ForeignServer *server = GetForeignServer(ft->serverid);
 		ft->exec_location = server->exec_location;
 	}
 
+	ft->num_segments = SeparateOutNumSegments(&ft->options);
+	if (ft->num_segments <= 0)
+	{
+		ft->num_segments = server->num_segments;
+	}
+
 	ReleaseSysCache(tp);
 
 	return ft;
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
@@ -3318,7 +3318,7 @@ create_foreignscan_path(PlannerInfo *root, RelOptInfo *rel,
 			CdbPathLocus_MakeGeneral(&(pathnode->path.locus));
 			break;
 		case FTEXECLOCATION_ALL_SEGMENTS:
-			CdbPathLocus_MakeStrewn(&(pathnode->path.locus), getgpsegmentCount());
+			CdbPathLocus_MakeStrewn(&(pathnode->path.locus), rel->num_segments);
 			break;
 		case FTEXECLOCATION_COORDINATOR:
 			CdbPathLocus_MakeEntry(&(pathnode->path.locus));
@@ -3383,7 +3383,7 @@ create_foreign_join_path(PlannerInfo *root, RelOptInfo *rel,
 			CdbPathLocus_MakeGeneral(&(pathnode->path.locus));
 			break;
 		case FTEXECLOCATION_ALL_SEGMENTS:
-			CdbPathLocus_MakeStrewn(&(pathnode->path.locus), getgpsegmentCount());
+			CdbPathLocus_MakeStrewn(&(pathnode->path.locus), rel->num_segments);
 			break;
 		case FTEXECLOCATION_COORDINATOR:
 			CdbPathLocus_MakeEntry(&(pathnode->path.locus));
diff --git a/src/backend/optimizer/util/plancat.c b/src/backend/optimizer/util/plancat.c
@@ -58,6 +58,7 @@
 
 #include "cdb/cdbappendonlyam.h"
 #include "cdb/cdbrelsize.h"
+#include "cdb/cdbutil.h"
 #include "catalog/pg_appendonly.h"
 #include "catalog/pg_foreign_server.h"
 #include "catalog/pg_inherits.h"
@@ -456,12 +457,14 @@ get_relation_info(PlannerInfo *root, Oid relationObjectId, bool inhparent,
 		rel->serverid = GetForeignServerIdByRelId(RelationGetRelid(relation));
 		rel->fdwroutine = GetFdwRoutineForRelation(relation, true);
 		rel->exec_location = GetForeignTable(RelationGetRelid(relation))->exec_location;
+		rel->num_segments = GetForeignTable(RelationGetRelid(relation))->num_segments;
 	}
 	else
 	{
 		rel->serverid = InvalidOid;
 		rel->fdwroutine = NULL;
 		rel->exec_location = FTEXECLOCATION_NOT_DEFINED;
+		rel->num_segments = getgpsegmentCount();
 	}
 
 	/* Collect info about relation's foreign keys, if relevant */
diff --git a/src/include/foreign/foreign.h b/src/include/foreign/foreign.h
@@ -53,6 +53,7 @@ typedef struct ForeignServer
 	char	   *serverversion;	/* server version, optional */
 	List	   *options;		/* srvoptions as DefElem list */
 	char		exec_location;  /* execute on MASTER, ANY or ALL SEGMENTS, Greenplum MPP specific */
+	int32		num_segments;	/* the number of segments of the foreign cluster */
 } ForeignServer;
 
 typedef struct UserMapping
@@ -69,6 +70,7 @@ typedef struct ForeignTable
 	Oid			serverid;		/* server Oid */
 	List	   *options;		/* ftoptions as DefElem list */
 	char		exec_location;  /* execute on COORDINATOR, ANY or ALL SEGMENTS, Greenplum MPP specific */
+	int32		num_segments;	/* the number of segments of the foreign table */
 } ForeignTable;
 
 /* Flags for GetForeignServerExtended */
@@ -79,6 +81,7 @@ typedef struct ForeignTable
 
 
 extern char SeparateOutMppExecute(List **options);
+extern int32 SeparateOutNumSegments(List **options);
 extern ForeignServer *GetForeignServer(Oid serverid);
 extern ForeignServer *GetForeignServerExtended(Oid serverid,
 											   bits16 flags);
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
@@ -864,6 +864,7 @@ typedef struct RelOptInfo
 	Oid			userid;			/* identifies user to check access as */
 	bool		useridiscurrent;	/* join is only valid for current user */
 	char		exec_location;  /* execute on MASTER, ANY or ALL SEGMENTS, Greenplum MPP specific */
+	int32		num_segments;  /* number of segments, Greenplum MPP specific */
 	/* use "struct FdwRoutine" to avoid including fdwapi.h here */
 	struct FdwRoutine *fdwroutine;
 	void	   *fdw_private;
diff --git a/src/test/regress/expected/gp_foreign_data.out b/src/test/regress/expected/gp_foreign_data.out
@@ -26,3 +26,16 @@ CREATE FOREIGN TABLE ft3 (
 CREATE FOREIGN TABLE ft4 (
 	c1 int
 ) SERVER s0 OPTIONS (delimiter ',', mpp_execute 'all segments');
+-- Test num_segments option
+CREATE SERVER s1 FOREIGN DATA WRAPPER dummy OPTIONS (num_segments '3');
+CREATE FOREIGN TABLE ft5 (
+       c1 int
+) SERVER s1 OPTIONS (delimiter ',', mpp_execute 'all segments', num_segments '5');
+\d+ ft5
+                                       Foreign table "public.ft5"
+ Column |  Type   | Collation | Nullable | Default | FDW options | Storage | Stats target | Description 
+--------+---------+-----------+----------+---------+-------------+---------+--------------+-------------
+ c1     | integer |           |          |         |             | plain   |              | 
+Server: s1
+FDW options: (delimiter ',', mpp_execute 'all segments', num_segments '5')
+
diff --git a/src/test/regress/sql/gp_foreign_data.sql b/src/test/regress/sql/gp_foreign_data.sql
@@ -24,3 +24,10 @@ CREATE FOREIGN TABLE ft3 (
 CREATE FOREIGN TABLE ft4 (
 	c1 int
 ) SERVER s0 OPTIONS (delimiter ',', mpp_execute 'all segments');
+
+-- Test num_segments option
+CREATE SERVER s1 FOREIGN DATA WRAPPER dummy OPTIONS (num_segments '3');
+CREATE FOREIGN TABLE ft5 (
+       c1 int
+) SERVER s1 OPTIONS (delimiter ',', mpp_execute 'all segments', num_segments '5');
+\d+ ft5
