Support TABLESAMPLE in router queries

src/backend/distributed/utils/ruleutils_10.c

@@ -6916,6 +6916,7 @@ get_from_clause_item(Node *jtnode, Query *query, deparse_context *context)
 		deparse_columns *colinfo = deparse_columns_fetch(varno, dpns);
 		RangeTblFunction *rtfunc1 = NULL;
 		bool		printalias;
+		CitusRTEKind rteKind = GetRangeTblKind(rte);
 
 		if (rte->lateral)
 			appendStringInfoString(buf, "LATERAL ");
@@ -7135,8 +7136,11 @@ get_from_clause_item(Node *jtnode, Query *query, deparse_context *context)
 		}
 
 		/* Tablesample clause must go after any alias */
-		if (rte->rtekind == RTE_RELATION && rte->tablesample)
+		if ((rteKind == CITUS_RTE_RELATION || rteKind == CITUS_RTE_SHARD) &&
+			rte->tablesample)
+		{
 			get_tablesample_def(rte->tablesample, context);
+		}
 	}
 	else if (IsA(jtnode, JoinExpr))
 	{

src/backend/distributed/utils/ruleutils_11.c

@@ -6933,6 +6933,7 @@ get_from_clause_item(Node *jtnode, Query *query, deparse_context *context)
 		deparse_columns *colinfo = deparse_columns_fetch(varno, dpns);
 		RangeTblFunction *rtfunc1 = NULL;
 		bool		printalias;
+		CitusRTEKind rteKind = GetRangeTblKind(rte);
 
 		if (rte->lateral)
 			appendStringInfoString(buf, "LATERAL ");
@@ -7152,8 +7153,11 @@ get_from_clause_item(Node *jtnode, Query *query, deparse_context *context)
 		}
 
 		/* Tablesample clause must go after any alias */
-		if (rte->rtekind == RTE_RELATION && rte->tablesample)
+		if ((rteKind == CITUS_RTE_RELATION || rteKind == CITUS_RTE_SHARD) &&
+			rte->tablesample)
+		{
 			get_tablesample_def(rte->tablesample, context);
+		}
 	}
 	else if (IsA(jtnode, JoinExpr))
 	{

src/backend/distributed/utils/ruleutils_96.c

@@ -6660,6 +6660,7 @@ get_from_clause_item(Node *jtnode, Query *query, deparse_context *context)
 		deparse_columns *colinfo = deparse_columns_fetch(varno, dpns);
 		RangeTblFunction *rtfunc1 = NULL;
 		bool		printalias;
+		CitusRTEKind rteKind = GetRangeTblKind(rte);
 
 		if (rte->lateral)
 			appendStringInfoString(buf, "LATERAL ");
@@ -6876,8 +6877,11 @@ get_from_clause_item(Node *jtnode, Query *query, deparse_context *context)
 		}
 
 		/* Tablesample clause must go after any alias */
-		if (rte->rtekind == RTE_RELATION && rte->tablesample)
+		if ((rteKind == CITUS_RTE_RELATION || rteKind == CITUS_RTE_SHARD) &&
+			rte->tablesample)
+		{
 			get_tablesample_def(rte->tablesample, context);
+		}
 	}
 	else if (IsA(jtnode, JoinExpr))
 	{

src/test/regress/expected/multi_simple_queries.out

@@ -615,4 +615,47 @@ SELECT count(*) FROM (
     50
 (1 row)
 
+-- tablesample is supported
+SELECT * FROM articles TABLESAMPLE SYSTEM (0) WHERE author_id = 1;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id | title | word_count 
+----+-----------+-------+------------
+(0 rows)
+
+SELECT * FROM articles TABLESAMPLE BERNOULLI (0) WHERE author_id = 1;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id | title | word_count 
+----+-----------+-------+------------
+(0 rows)
+
+SELECT * FROM articles TABLESAMPLE SYSTEM (100) WHERE author_id = 1 ORDER BY id;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id |    title     | word_count 
+----+-----------+--------------+------------
+  1 |         1 | arsenous     |       9572
+ 11 |         1 | alamo        |       1347
+ 21 |         1 | arcading     |       5890
+ 31 |         1 | athwartships |       7271
+ 41 |         1 | aznavour     |      11814
+(5 rows)
+
+SELECT * FROM articles TABLESAMPLE BERNOULLI (100) WHERE author_id = 1 ORDER BY id;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id |    title     | word_count 
+----+-----------+--------------+------------
+  1 |         1 | arsenous     |       9572
+ 11 |         1 | alamo        |       1347
+ 21 |         1 | arcading     |       5890
+ 31 |         1 | athwartships |       7271
+ 41 |         1 | aznavour     |      11814
+(5 rows)
+
 SET client_min_messages to 'NOTICE';

src/test/regress/expected/multi_simple_queries_0.out

@@ -559,4 +559,47 @@ SELECT count(*) FROM (
     50
 (1 row)
 
+-- tablesample is supported
+SELECT * FROM articles TABLESAMPLE SYSTEM (0) WHERE author_id = 1;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id | title | word_count 
+----+-----------+-------+------------
+(0 rows)
+
+SELECT * FROM articles TABLESAMPLE BERNOULLI (0) WHERE author_id = 1;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id | title | word_count 
+----+-----------+-------+------------
+(0 rows)
+
+SELECT * FROM articles TABLESAMPLE SYSTEM (100) WHERE author_id = 1 ORDER BY id;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id |    title     | word_count 
+----+-----------+--------------+------------
+  1 |         1 | arsenous     |       9572
+ 11 |         1 | alamo        |       1347
+ 21 |         1 | arcading     |       5890
+ 31 |         1 | athwartships |       7271
+ 41 |         1 | aznavour     |      11814
+(5 rows)
+
+SELECT * FROM articles TABLESAMPLE BERNOULLI (100) WHERE author_id = 1 ORDER BY id;
+DEBUG:  Creating router plan
+DEBUG:  Plan is router executable
+DETAIL:  distribution column value: 1
+ id | author_id |    title     | word_count 
+----+-----------+--------------+------------
+  1 |         1 | arsenous     |       9572
+ 11 |         1 | alamo        |       1347
+ 21 |         1 | arcading     |       5890
+ 31 |         1 | athwartships |       7271
+ 41 |         1 | aznavour     |      11814
+(5 rows)
+
 SET client_min_messages to 'NOTICE';

src/test/regress/sql/multi_simple_queries.sql

@@ -296,4 +296,10 @@ SELECT count(*) FROM (
                   xmax IS NOT NULL
 ) x;
 
+-- tablesample is supported
+SELECT * FROM articles TABLESAMPLE SYSTEM (0) WHERE author_id = 1;
+SELECT * FROM articles TABLESAMPLE BERNOULLI (0) WHERE author_id = 1;
+SELECT * FROM articles TABLESAMPLE SYSTEM (100) WHERE author_id = 1 ORDER BY id;
+SELECT * FROM articles TABLESAMPLE BERNOULLI (100) WHERE author_id = 1 ORDER BY id;
+
 SET client_min_messages to 'NOTICE';