v5-0001-support-COPY-partitioned_table-TO.patch

From a4c643ac3a9f40bbdf07dcacc38527ef6e86f1bc Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Fri, 28 Mar 2025 11:05:53 +0800
Subject: [PATCH v5 1/1] support COPY partitioned_table TO

CREATE TABLE pp (id  INT, val int ) PARTITION BY RANGE (id);
create table pp_1 (val int, id int);
create table pp_2 (val int, id int);
ALTER TABLE pp ATTACH PARTITION pp_1 FOR VALUES FROM (1) TO (5);
ALTER TABLE pp ATTACH PARTITION pp_2 FOR VALUES FROM (5) TO (10);
insert into pp select g, 10 + g from generate_series(1,9) g;
copy pp to stdout(header);

the above case is much slower (around 25% some case) than
``COPY (select * from pp) to stdout(header);``,
because of column remaping.  but this is still a new
feature, since master does not support ``COPY (partitioned_table)``.

reivewed by: vignesh C <vignesh21@gmail.com>
reivewed by: David Rowley <dgrowleyml@gmail.com>
reivewed by: Melih Mutlu <m.melihmutlu@gmail.com>

discussion: https://postgr.es/m/CACJufxEZt+G19Ors3bQUq-42-61__C=y5k2wk=sHEFRusu7=iQ@mail.gmail.com
commitfest entry: https://commitfest.postgresql.org/patch/5467/
---
 doc/src/sgml/ref/copy.sgml          |  8 +--
 src/backend/commands/copyto.c       | 89 +++++++++++++++++++++++++++--
 src/test/regress/expected/copy2.out | 16 ++++++
 src/test/regress/sql/copy2.sql      | 11 ++++
 4 files changed, 114 insertions(+), 10 deletions(-)

diff --git a/doc/src/sgml/ref/copy.sgml b/doc/src/sgml/ref/copy.sgml
index df093da97c5..f86e0b7ec35 100644
--- a/doc/src/sgml/ref/copy.sgml
+++ b/doc/src/sgml/ref/copy.sgml
@@ -521,15 +521,15 @@ COPY <replaceable class="parameter">count</replaceable>
 
    <para>
     <command>COPY TO</command> can be used only with plain
-    tables, not views, and does not copy rows from child tables
-    or child partitions.  For example, <literal>COPY <replaceable
+    tables, not views, and does not copy rows from child tables,
+    however <command>COPY TO</command> can be used with partitioned tables.
+    For example, in a table inheritance hierarchy, <literal>COPY <replaceable
     class="parameter">table</replaceable> TO</literal> copies
     the same rows as <literal>SELECT * FROM ONLY <replaceable
     class="parameter">table</replaceable></literal>.
     The syntax <literal>COPY (SELECT * FROM <replaceable
     class="parameter">table</replaceable>) TO ...</literal> can be used to
-    dump all of the rows in an inheritance hierarchy, partitioned table,
-    or view.
+    dump all of the rows in an inheritance hierarchy, or view.
    </para>
 
    <para>
diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c
index 84a3f3879a8..0973dc9c14b 100644
--- a/src/backend/commands/copyto.c
+++ b/src/backend/commands/copyto.c
@@ -19,6 +19,8 @@
 #include <sys/stat.h>
 
 #include "access/tableam.h"
+#include "access/table.h"
+#include "catalog/pg_inherits.h"
 #include "commands/copyapi.h"
 #include "commands/progress.h"
 #include "executor/execdesc.h"
@@ -82,6 +84,7 @@ typedef struct CopyToStateData
 	List	   *attnumlist;		/* integer list of attnums to copy */
 	char	   *filename;		/* filename, or NULL for STDOUT */
 	bool		is_program;		/* is 'filename' a program to popen? */
+	List	   *partitions;		/* oid list of partition oid for copy to */
 	copy_data_dest_cb data_dest_cb; /* function for writing data */
 
 	CopyFormatOptions opts;
@@ -643,6 +646,8 @@ BeginCopyTo(ParseState *pstate,
 		PROGRESS_COPY_COMMAND_TO,
 		0
 	};
+	List	   *children = NIL;
+	List	   *scan_oids = NIL;
 
 	if (rel != NULL && rel->rd_rel->relkind != RELKIND_RELATION)
 	{
@@ -670,11 +675,28 @@ BeginCopyTo(ParseState *pstate,
 					 errmsg("cannot copy from sequence \"%s\"",
 							RelationGetRelationName(rel))));
 		else if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
-			ereport(ERROR,
-					(errcode(ERRCODE_WRONG_OBJECT_TYPE),
-					 errmsg("cannot copy from partitioned table \"%s\"",
-							RelationGetRelationName(rel)),
-					 errhint("Try the COPY (SELECT ...) TO variant.")));
+		{
+			children = find_all_inheritors(RelationGetRelid(rel),
+										   AccessShareLock,
+										   NULL);
+			foreach_oid(childreloid, children)
+			{
+				char		 relkind = get_rel_relkind(childreloid);
+
+				if (relkind == RELKIND_FOREIGN_TABLE)
+					ereport(ERROR,
+							errcode(ERRCODE_WRONG_OBJECT_TYPE),
+							errmsg("cannot copy from foreign table \"%s\"",
+									RelationGetRelationName(rel)),
+							errdetail("partition \"%s\" is a foreign table", RelationGetRelationName(rel)),
+							errhint("Try the COPY (SELECT ...) TO variant."));
+
+				if (RELKIND_HAS_PARTITIONS(relkind))
+					continue;
+
+				scan_oids = lappend_oid(scan_oids, childreloid);
+			}
+		}
 		else
 			ereport(ERROR,
 					(errcode(ERRCODE_WRONG_OBJECT_TYPE),
@@ -710,6 +732,7 @@ BeginCopyTo(ParseState *pstate,
 		cstate->rel = rel;
 
 		tupDesc = RelationGetDescr(cstate->rel);
+		cstate->partitions = list_copy(scan_oids);
 	}
 	else
 	{
@@ -1066,7 +1089,61 @@ DoCopyTo(CopyToState cstate)
 
 	cstate->routine->CopyToStart(cstate, tupDesc);
 
-	if (cstate->rel)
+	/*
+	 * if COPY TO source table is a partitioned table, then open each
+	 * partition and process each individual partition.
+	 */
+	if (cstate->rel && cstate->rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+	{
+		processed = 0;
+
+		foreach_oid(scan_oid, cstate->partitions)
+		{
+			TupleTableSlot *slot;
+			TableScanDesc scandesc;
+			Relation		scan_rel;
+			TupleDesc		scan_tupdesc;
+			AttrMap    		*map;
+			TupleTableSlot  *root_slot = NULL;
+			TupleTableSlot  *original_slot = NULL;
+
+			scan_rel = table_open(scan_oid, AccessShareLock);
+			scan_tupdesc = RelationGetDescr(scan_rel);
+			map = build_attrmap_by_name_if_req(tupDesc, scan_tupdesc, false);
+
+			scandesc = table_beginscan(scan_rel, GetActiveSnapshot(), 0, NULL);
+			slot = table_slot_create(scan_rel, NULL);
+
+			while (table_scan_getnextslot(scandesc, ForwardScanDirection, slot))
+			{
+				CHECK_FOR_INTERRUPTS();
+
+				/* Deconstruct the tuple ... */
+				if (map != NULL)
+				{
+					original_slot = slot;
+					root_slot = MakeSingleTupleTableSlot(tupDesc, &TTSOpsBufferHeapTuple);
+					slot = execute_attr_map_slot(map, slot, root_slot);
+				}
+				else
+					slot_getallattrs(slot);
+
+				/* Format and send the data */
+				CopyOneRowTo(cstate, slot);
+
+				pgstat_progress_update_param(PROGRESS_COPY_TUPLES_PROCESSED,
+											 ++processed);
+
+				if (original_slot != NULL)
+					ExecDropSingleTupleTableSlot(original_slot);
+			};
+
+			ExecDropSingleTupleTableSlot(slot);
+			table_endscan(scandesc);
+			table_close(scan_rel, AccessShareLock);
+		}
+	}
+	else if (cstate->rel)
 	{
 		TupleTableSlot *slot;
 		TableScanDesc scandesc;
diff --git a/src/test/regress/expected/copy2.out b/src/test/regress/expected/copy2.out
index 64ea33aeae8..dcd97ae45b7 100644
--- a/src/test/regress/expected/copy2.out
+++ b/src/test/regress/expected/copy2.out
@@ -929,3 +929,19 @@ truncate copy_default;
 -- DEFAULT cannot be used in COPY TO
 copy (select 1 as test) TO stdout with (default '\D');
 ERROR:  COPY DEFAULT cannot be used with COPY TO
+-- COPY TO with partitioned table
+CREATE TABLE pp (id  INT, val int ) PARTITION BY RANGE (id);
+create table pp_1 (val int, id int);
+create table pp_2 (val int, id int);
+ALTER TABLE pp ATTACH PARTITION pp_1 FOR VALUES FROM (1) TO (3);
+ALTER TABLE pp ATTACH PARTITION pp_2 FOR VALUES FROM (3) TO (7);
+insert into pp select g, 10 + g from generate_series(1,6) g;
+copy pp to stdout(header);
+id	val
+1	11
+2	12
+3	13
+4	14
+5	15
+6	16
+DROP TABLE PP;
diff --git a/src/test/regress/sql/copy2.sql b/src/test/regress/sql/copy2.sql
index 45273557ce0..ba984388248 100644
--- a/src/test/regress/sql/copy2.sql
+++ b/src/test/regress/sql/copy2.sql
@@ -707,3 +707,14 @@ truncate copy_default;
 
 -- DEFAULT cannot be used in COPY TO
 copy (select 1 as test) TO stdout with (default '\D');
+
+-- COPY TO with partitioned table
+CREATE TABLE pp (id  INT, val int ) PARTITION BY RANGE (id);
+create table pp_1 (val int, id int);
+create table pp_2 (val int, id int);
+ALTER TABLE pp ATTACH PARTITION pp_1 FOR VALUES FROM (1) TO (3);
+ALTER TABLE pp ATTACH PARTITION pp_2 FOR VALUES FROM (3) TO (7);
+insert into pp select g, 10 + g from generate_series(1,6) g;
+copy pp to stdout(header);
+
+DROP TABLE PP;
-- 
2.34.1