Support custom cast from / to timestamptz in time partition management UDFs (#6923)

This is to implement custom cast of table partition column
type from / to `timestamptz` in time partition management UDFs, as
proposed in ticket #6454

The general idea is for a time partition column with type other than
`date`, `timestamp`, or `timestamptz`, users can provide custom
bidirectional cast between the column type and `timestamptz`, the UDFs
then will be able to create and drop time partitions for such tables.

Fixes #6454

---------

Signed-off-by: Xin Li <xin@swirldslabs.com>
Co-authored-by: Marco Slot <marco.slot@microsoft.com>
Co-authored-by: Ahmet Gedemenli <afgedemenli@gmail.com>

src/backend/distributed/sql/citus--11.3-1--12.0-1.sql

@@ -39,3 +39,6 @@ DROP FUNCTION citus_shard_sizes;
 -- udfs to convert a regular/tenant schema to a tenant/regular schema
 #include "udfs/citus_schema_distribute/12.0-1.sql"
 #include "udfs/citus_schema_undistribute/12.0-1.sql"
+
+#include "udfs/drop_old_time_partitions/12.0-1.sql"
+#include "udfs/get_missing_time_partition_ranges/12.0-1.sql"

src/backend/distributed/sql/downgrades/citus--12.0-1--11.3-1.sql

@@ -73,3 +73,6 @@ DROP FUNCTION pg_catalog.citus_stat_tenants_local_internal(
     OUT DOUBLE PRECISION,
     OUT BIGINT);
 #include "../udfs/citus_stat_tenants_local/11.3-1.sql"
+
+#include "../udfs/drop_old_time_partitions/10.2-1.sql"
+#include "../udfs/get_missing_time_partition_ranges/10.2-1.sql"

src/backend/distributed/sql/udfs/drop_old_time_partitions/12.0-1.sql

@@ -0,0 +1,68 @@
+CREATE OR REPLACE PROCEDURE pg_catalog.drop_old_time_partitions(
+    table_name regclass,
+    older_than timestamptz)
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    -- properties of the partitioned table
+    number_of_partition_columns int;
+    partition_column_index int;
+    partition_column_type regtype;
+
+    -- used to support dynamic type casting between the partition column type and timestamptz
+    custom_cast text;
+    is_partition_column_castable boolean;
+    older_partitions_query text;
+
+    r record;
+BEGIN
+    -- check whether the table is time partitioned table, if not error out
+    SELECT partnatts, partattrs[0]
+    INTO number_of_partition_columns, partition_column_index
+    FROM pg_catalog.pg_partitioned_table
+    WHERE partrelid = table_name;
+
+    IF NOT FOUND THEN
+        RAISE '% is not partitioned', table_name::text;
+    ELSIF number_of_partition_columns <> 1 THEN
+        RAISE 'partitioned tables with multiple partition columns are not supported';
+    END IF;
+
+    -- get datatype here to check interval-table type
+    SELECT atttypid
+    INTO partition_column_type
+    FROM pg_attribute
+    WHERE attrelid = table_name::oid
+    AND attnum = partition_column_index;
+
+    -- we currently only support partitioning by date, timestamp, and timestamptz
+    custom_cast = '';
+    IF partition_column_type <> 'date'::regtype
+    AND partition_column_type <> 'timestamp'::regtype
+    AND partition_column_type <> 'timestamptz'::regtype  THEN
+      SELECT EXISTS(SELECT OID FROM pg_cast WHERE castsource = partition_column_type AND casttarget = 'timestamptz'::regtype) AND
+             EXISTS(SELECT OID FROM pg_cast WHERE castsource = 'timestamptz'::regtype AND casttarget = partition_column_type)
+      INTO is_partition_column_castable;
+      IF not is_partition_column_castable THEN
+        RAISE 'type of the partition column of the table % must be date, timestamp or timestamptz', table_name;
+      END IF;
+      custom_cast = format('::%s', partition_column_type);
+    END IF;
+
+    older_partitions_query = format('SELECT partition, nspname AS schema_name, relname AS table_name, from_value, to_value
+        FROM pg_catalog.time_partitions, pg_catalog.pg_class c, pg_catalog.pg_namespace n
+        WHERE parent_table = $1 AND partition = c.oid AND c.relnamespace = n.oid
+        AND to_value IS NOT NULL
+        AND to_value%1$s::timestamptz <= $2
+        ORDER BY to_value%1$s::timestamptz', custom_cast);
+    FOR r IN EXECUTE older_partitions_query USING table_name, older_than
+    LOOP
+        RAISE NOTICE 'dropping % with start time % and end time %', r.partition, r.from_value, r.to_value;
+        EXECUTE format('DROP TABLE %I.%I', r.schema_name, r.table_name);
+    END LOOP;
+END;
+$$;
+COMMENT ON PROCEDURE pg_catalog.drop_old_time_partitions(
+    table_name regclass,
+    older_than timestamptz)
+IS 'drop old partitions of a time-partitioned table';

src/backend/distributed/sql/udfs/drop_old_time_partitions/latest.sql

@@ -1,6 +1,6 @@
 CREATE OR REPLACE PROCEDURE pg_catalog.drop_old_time_partitions(
-	table_name regclass,
+    table_name regclass,
-	older_than timestamptz)
+    older_than timestamptz)
 LANGUAGE plpgsql
 AS $$
 DECLARE
@@ -9,6 +9,11 @@ DECLARE
     partition_column_index int;
     partition_column_type regtype;
 
+    -- used to support dynamic type casting between the partition column type and timestamptz
+    custom_cast text;
+    is_partition_column_castable boolean;
+    older_partitions_query text;
+
     r record;
 BEGIN
     -- check whether the table is time partitioned table, if not error out
@@ -31,19 +36,26 @@ BEGIN
     AND attnum = partition_column_index;
 
     -- we currently only support partitioning by date, timestamp, and timestamptz
+    custom_cast = '';
     IF partition_column_type <> 'date'::regtype
     AND partition_column_type <> 'timestamp'::regtype
     AND partition_column_type <> 'timestamptz'::regtype  THEN
+      SELECT EXISTS(SELECT OID FROM pg_cast WHERE castsource = partition_column_type AND casttarget = 'timestamptz'::regtype) AND
+             EXISTS(SELECT OID FROM pg_cast WHERE castsource = 'timestamptz'::regtype AND casttarget = partition_column_type)
+      INTO is_partition_column_castable;
+      IF not is_partition_column_castable THEN
         RAISE 'type of the partition column of the table % must be date, timestamp or timestamptz', table_name;
+      END IF;
+      custom_cast = format('::%s', partition_column_type);
     END IF;
 
-    FOR r IN
+    older_partitions_query = format('SELECT partition, nspname AS schema_name, relname AS table_name, from_value, to_value
-		SELECT partition, nspname AS schema_name, relname AS table_name, from_value, to_value
+        FROM pg_catalog.time_partitions, pg_catalog.pg_class c, pg_catalog.pg_namespace n
-		FROM pg_catalog.time_partitions, pg_catalog.pg_class c, pg_catalog.pg_namespace n
+        WHERE parent_table = $1 AND partition = c.oid AND c.relnamespace = n.oid
-		WHERE parent_table = table_name AND partition = c.oid AND c.relnamespace = n.oid
+        AND to_value IS NOT NULL
-		AND to_value IS NOT NULL
+        AND to_value%1$s::timestamptz <= $2
-		AND to_value::timestamptz <= older_than
+        ORDER BY to_value%1$s::timestamptz', custom_cast);
-		ORDER BY to_value::timestamptz
+    FOR r IN EXECUTE older_partitions_query USING table_name, older_than
     LOOP
         RAISE NOTICE 'dropping % with start time % and end time %', r.partition, r.from_value, r.to_value;
         EXECUTE format('DROP TABLE %I.%I', r.schema_name, r.table_name);
@@ -51,6 +63,6 @@ BEGIN
 END;
 $$;
 COMMENT ON PROCEDURE pg_catalog.drop_old_time_partitions(
-	table_name regclass,
+    table_name regclass,
-	older_than timestamptz)
+    older_than timestamptz)
 IS 'drop old partitions of a time-partitioned table';

src/backend/distributed/sql/udfs/get_missing_time_partition_ranges/12.0-1.sql

@@ -0,0 +1,239 @@
+CREATE OR REPLACE FUNCTION pg_catalog.get_missing_time_partition_ranges(
+    table_name regclass,
+    partition_interval INTERVAL,
+    to_value timestamptz,
+    from_value timestamptz DEFAULT now())
+returns table(
+    partition_name text,
+    range_from_value text,
+    range_to_value text)
+LANGUAGE plpgsql
+AS $$
+DECLARE
+    -- properties of the partitioned table
+    table_name_text text;
+    table_schema_text text;
+    number_of_partition_columns int;
+    partition_column_index int;
+    partition_column_type regtype;
+
+    -- used for generating time ranges
+    current_range_from_value timestamptz := NULL;
+    current_range_to_value timestamptz := NULL;
+    current_range_from_value_text text;
+    current_range_to_value_text text;
+
+    -- used to check whether there are misaligned (manually created) partitions
+    manual_partition regclass;
+    manual_partition_from_value_text text;
+    manual_partition_to_value_text text;
+
+    -- used for partition naming
+    partition_name_format text;
+    max_table_name_length int := current_setting('max_identifier_length');
+
+    -- used to determine whether the partition_interval is a day multiple
+    is_day_multiple boolean;
+
+    -- used to support dynamic type casting between the partition column type and timestamptz
+    custom_cast text;
+    is_partition_column_castable boolean;
+    partition regclass;
+    partition_covers_query text;
+    partition_exist_query text;
+BEGIN
+    -- check whether the table is time partitioned table, if not error out
+    SELECT relname, nspname, partnatts, partattrs[0]
+    INTO table_name_text, table_schema_text, number_of_partition_columns, partition_column_index
+    FROM pg_catalog.pg_partitioned_table, pg_catalog.pg_class c, pg_catalog.pg_namespace n
+    WHERE partrelid = c.oid AND c.oid = table_name
+    AND c.relnamespace = n.oid;
+    IF NOT FOUND THEN
+        RAISE '% is not partitioned', table_name;
+    ELSIF number_of_partition_columns <> 1 THEN
+        RAISE 'partitioned tables with multiple partition columns are not supported';
+    END IF;
+
+    -- to not to have partitions to be created in parallel
+    EXECUTE format('LOCK TABLE %I.%I IN SHARE UPDATE EXCLUSIVE MODE', table_schema_text, table_name_text);
+
+    -- get datatype here to check interval-table type alignment and generate range values in the right data format
+    SELECT atttypid
+    INTO partition_column_type
+    FROM pg_attribute
+    WHERE attrelid = table_name::oid
+    AND attnum = partition_column_index;
+
+    -- we currently only support partitioning by date, timestamp, and timestamptz
+    custom_cast = '';
+    IF partition_column_type <> 'date'::regtype
+    AND partition_column_type <> 'timestamp'::regtype
+    AND partition_column_type <> 'timestamptz'::regtype THEN
+      SELECT EXISTS(SELECT OID FROM pg_cast WHERE castsource = partition_column_type AND casttarget = 'timestamptz'::regtype) AND
+             EXISTS(SELECT OID FROM pg_cast WHERE castsource = 'timestamptz'::regtype AND casttarget = partition_column_type)
+      INTO is_partition_column_castable;
+      IF not is_partition_column_castable THEN
+        RAISE 'type of the partition column of the table % must be date, timestamp or timestamptz', table_name;
+      END IF;
+      custom_cast = format('::%s', partition_column_type);
+    END IF;
+
+    IF partition_column_type = 'date'::regtype AND partition_interval IS NOT NULL THEN
+        SELECT date_trunc('day', partition_interval) = partition_interval
+        INTO is_day_multiple;
+
+        IF NOT is_day_multiple THEN
+            RAISE 'partition interval of date partitioned table must be day or multiple days';
+        END IF;
+    END IF;
+
+    -- If no partition exists, truncate from_value to find intuitive initial value.
+    -- If any partition exist, use the initial partition as the pivot partition.
+    -- tp.to_value and tp.from_value are equal to '', if default partition exists.
+    EXECUTE format('SELECT tp.from_value%1$s::timestamptz, tp.to_value%1$s::timestamptz
+        FROM pg_catalog.time_partitions tp
+        WHERE parent_table = $1 AND tp.to_value <> '' AND tp.from_value <> ''
+        ORDER BY tp.from_value%1$s::timestamptz ASC
+        LIMIT 1', custom_cast)
+    INTO current_range_from_value, current_range_to_value
+    USING table_name;
+
+    IF current_range_from_value is NULL THEN
+        -- Decide on the current_range_from_value of the initial partition according to interval of the table.
+        -- Since we will create all other partitions by adding intervals, truncating given start time will provide
+        -- more intuitive interval ranges, instead of starting from from_value directly.
+        IF partition_interval < INTERVAL '1 hour' THEN
+            current_range_from_value = date_trunc('minute', from_value);
+        ELSIF partition_interval < INTERVAL '1 day' THEN
+            current_range_from_value = date_trunc('hour', from_value);
+        ELSIF partition_interval < INTERVAL '1 week' THEN
+            current_range_from_value = date_trunc('day', from_value);
+        ELSIF partition_interval < INTERVAL '1 month' THEN
+            current_range_from_value = date_trunc('week', from_value);
+        ELSIF partition_interval = INTERVAL '3 months' THEN
+            current_range_from_value = date_trunc('quarter', from_value);
+        ELSIF partition_interval < INTERVAL '1 year' THEN
+            current_range_from_value = date_trunc('month', from_value);
+        ELSE
+            current_range_from_value = date_trunc('year', from_value);
+        END IF;
+
+        current_range_to_value := current_range_from_value + partition_interval;
+    ELSE
+        -- if from_value is newer than pivot's from value, go forward, else go backward
+        IF from_value >= current_range_from_value THEN
+            WHILE current_range_from_value < from_value LOOP
+                    current_range_from_value := current_range_from_value + partition_interval;
+            END LOOP;
+        ELSE
+            WHILE current_range_from_value > from_value LOOP
+                    current_range_from_value := current_range_from_value - partition_interval;
+            END LOOP;
+        END IF;
+        current_range_to_value := current_range_from_value + partition_interval;
+    END IF;
+
+    -- reuse pg_partman naming scheme for back-and-forth migration
+    IF partition_interval = INTERVAL '3 months' THEN
+        -- include quarter in partition name
+        partition_name_format = 'YYYY"q"Q';
+    ELSIF partition_interval = INTERVAL '1 week' THEN
+        -- include week number in partition name
+        partition_name_format := 'IYYY"w"IW';
+    ELSE
+        -- always start with the year
+        partition_name_format := 'YYYY';
+
+        IF partition_interval < INTERVAL '1 year' THEN
+            -- include month in partition name
+            partition_name_format := partition_name_format || '_MM';
+        END IF;
+
+        IF partition_interval < INTERVAL '1 month' THEN
+            -- include day of month in partition name
+            partition_name_format := partition_name_format || '_DD';
+        END IF;
+
+        IF partition_interval < INTERVAL '1 day' THEN
+            -- include time of day in partition name
+            partition_name_format := partition_name_format || '_HH24MI';
+        END IF;
+
+        IF partition_interval < INTERVAL '1 minute' THEN
+             -- include seconds in time of day in partition name
+             partition_name_format := partition_name_format || 'SS';
+        END IF;
+    END IF;
+
+    partition_exist_query = format('SELECT partition FROM pg_catalog.time_partitions tp
+        WHERE tp.from_value%1$s::timestamptz = $1 AND tp.to_value%1$s::timestamptz = $2 AND parent_table = $3',
+        custom_cast);
+    partition_covers_query = format('SELECT partition, tp.from_value, tp.to_value
+        FROM pg_catalog.time_partitions tp
+        WHERE
+            (($1 >= tp.from_value%1$s::timestamptz AND $1 < tp.to_value%1$s::timestamptz) OR
+            ($2 > tp.from_value%1$s::timestamptz AND $2 < tp.to_value%1$s::timestamptz)) AND
+            parent_table = $3',
+        custom_cast);
+
+    WHILE current_range_from_value < to_value LOOP
+        -- Check whether partition with given range has already been created
+        -- Since partition interval can be given with different types, we are converting
+        -- all variables to timestamptz to make sure that we are comparing same type of parameters
+        EXECUTE partition_exist_query into partition using current_range_from_value, current_range_to_value, table_name;
+
+        IF partition is not NULL THEN
+            current_range_from_value := current_range_to_value;
+            current_range_to_value := current_range_to_value + partition_interval;
+            CONTINUE;
+        END IF;
+
+        -- Check whether any other partition covers from_value or to_value
+        -- That means some partitions doesn't align with the initial partition.
+        -- In other words, gap(s) exist between partitions which is not multiple of intervals.
+        EXECUTE partition_covers_query
+        INTO manual_partition, manual_partition_from_value_text, manual_partition_to_value_text
+        using current_range_from_value, current_range_to_value, table_name;
+
+        IF manual_partition is not NULL THEN
+            RAISE 'partition % with the range from % to % does not align with the initial partition given the partition interval',
+            manual_partition::text,
+            manual_partition_from_value_text,
+            manual_partition_to_value_text
+            USING HINT = 'Only use partitions of the same size, without gaps between partitions.';
+        END IF;
+
+        IF partition_column_type = 'date'::regtype THEN
+            SELECT current_range_from_value::date::text INTO current_range_from_value_text;
+            SELECT current_range_to_value::date::text INTO current_range_to_value_text;
+        ELSIF partition_column_type = 'timestamp without time zone'::regtype THEN
+            SELECT current_range_from_value::timestamp::text INTO current_range_from_value_text;
+            SELECT current_range_to_value::timestamp::text INTO current_range_to_value_text;
+        ELSIF partition_column_type = 'timestamp with time zone'::regtype THEN
+            SELECT current_range_from_value::timestamptz::text INTO current_range_from_value_text;
+            SELECT current_range_to_value::timestamptz::text INTO current_range_to_value_text;
+        ELSE
+            EXECUTE format('SELECT $1%s::text', custom_cast) INTO current_range_from_value_text using current_range_from_value;
+            EXECUTE format('SELECT $1%s::text', custom_cast) INTO current_range_to_value_text using current_range_to_value;
+        END IF;
+
+        -- use range values within the name of partition to have unique partition names
+        RETURN QUERY
+        SELECT
+            substring(table_name_text, 0, max_table_name_length - length(to_char(current_range_from_value, partition_name_format)) - 1) || '_p' ||
+            to_char(current_range_from_value, partition_name_format),
+            current_range_from_value_text,
+            current_range_to_value_text;
+
+        current_range_from_value := current_range_to_value;
+        current_range_to_value := current_range_to_value + partition_interval;
+    END LOOP;
+    RETURN;
+END;
+$$;
+COMMENT ON FUNCTION pg_catalog.get_missing_time_partition_ranges(
+    table_name regclass,
+    partition_interval INTERVAL,
+    to_value timestamptz,
+    from_value timestamptz)
+IS 'get missing partitions ranges for table within the range using the given interval';

src/backend/distributed/sql/udfs/get_missing_time_partition_ranges/latest.sql

@@ -34,6 +34,13 @@ DECLARE
 
     -- used to determine whether the partition_interval is a day multiple
     is_day_multiple boolean;
+
+    -- used to support dynamic type casting between the partition column type and timestamptz
+    custom_cast text;
+    is_partition_column_castable boolean;
+    partition regclass;
+    partition_covers_query text;
+    partition_exist_query text;
 BEGIN
     -- check whether the table is time partitioned table, if not error out
     SELECT relname, nspname, partnatts, partattrs[0]
@@ -58,10 +65,17 @@ BEGIN
     AND attnum = partition_column_index;
 
     -- we currently only support partitioning by date, timestamp, and timestamptz
+    custom_cast = '';
     IF partition_column_type <> 'date'::regtype
     AND partition_column_type <> 'timestamp'::regtype
-    AND partition_column_type <> 'timestamptz'::regtype  THEN
+    AND partition_column_type <> 'timestamptz'::regtype THEN
+      SELECT EXISTS(SELECT OID FROM pg_cast WHERE castsource = partition_column_type AND casttarget = 'timestamptz'::regtype) AND
+             EXISTS(SELECT OID FROM pg_cast WHERE castsource = 'timestamptz'::regtype AND casttarget = partition_column_type)
+      INTO is_partition_column_castable;
+      IF not is_partition_column_castable THEN
         RAISE 'type of the partition column of the table % must be date, timestamp or timestamptz', table_name;
+      END IF;
+      custom_cast = format('::%s', partition_column_type);
     END IF;
 
     IF partition_column_type = 'date'::regtype AND partition_interval IS NOT NULL THEN
@@ -76,14 +90,15 @@ BEGIN
     -- If no partition exists, truncate from_value to find intuitive initial value.
     -- If any partition exist, use the initial partition as the pivot partition.
     -- tp.to_value and tp.from_value are equal to '', if default partition exists.
-    SELECT tp.from_value::timestamptz, tp.to_value::timestamptz
+    EXECUTE format('SELECT tp.from_value%1$s::timestamptz, tp.to_value%1$s::timestamptz
+        FROM pg_catalog.time_partitions tp
+        WHERE parent_table = $1 AND tp.to_value <> '' AND tp.from_value <> ''
+        ORDER BY tp.from_value%1$s::timestamptz ASC
+        LIMIT 1', custom_cast)
     INTO current_range_from_value, current_range_to_value
-    FROM pg_catalog.time_partitions tp
+    USING table_name;
-    WHERE parent_table = table_name AND tp.to_value <> '' AND tp.from_value <> ''
-    ORDER BY tp.from_value::timestamptz ASC
-    LIMIT 1;
 
-    IF NOT FOUND THEN
+    IF current_range_from_value is NULL THEN
         -- Decide on the current_range_from_value of the initial partition according to interval of the table.
         -- Since we will create all other partitions by adding intervals, truncating given start time will provide
         -- more intuitive interval ranges, instead of starting from from_value directly.
@@ -150,16 +165,24 @@ BEGIN
         END IF;
     END IF;
 
+    partition_exist_query = format('SELECT partition FROM pg_catalog.time_partitions tp
+        WHERE tp.from_value%1$s::timestamptz = $1 AND tp.to_value%1$s::timestamptz = $2 AND parent_table = $3',
+        custom_cast);
+    partition_covers_query = format('SELECT partition, tp.from_value, tp.to_value
+        FROM pg_catalog.time_partitions tp
+        WHERE
+            (($1 >= tp.from_value%1$s::timestamptz AND $1 < tp.to_value%1$s::timestamptz) OR
+            ($2 > tp.from_value%1$s::timestamptz AND $2 < tp.to_value%1$s::timestamptz)) AND
+            parent_table = $3',
+        custom_cast);
+
     WHILE current_range_from_value < to_value LOOP
         -- Check whether partition with given range has already been created
         -- Since partition interval can be given with different types, we are converting
         -- all variables to timestamptz to make sure that we are comparing same type of parameters
-        PERFORM * FROM pg_catalog.time_partitions tp
+        EXECUTE partition_exist_query into partition using current_range_from_value, current_range_to_value, table_name;
-        WHERE
+
-            tp.from_value::timestamptz = current_range_from_value::timestamptz AND
+        IF partition is not NULL THEN
-            tp.to_value::timestamptz = current_range_to_value::timestamptz AND
-            parent_table = table_name;
-        IF found THEN
             current_range_from_value := current_range_to_value;
             current_range_to_value := current_range_to_value + partition_interval;
             CONTINUE;
@@ -168,20 +191,16 @@ BEGIN
         -- Check whether any other partition covers from_value or to_value
         -- That means some partitions doesn't align with the initial partition.
         -- In other words, gap(s) exist between partitions which is not multiple of intervals.
-        SELECT partition, tp.from_value::text, tp.to_value::text
+        EXECUTE partition_covers_query
         INTO manual_partition, manual_partition_from_value_text, manual_partition_to_value_text
-        FROM pg_catalog.time_partitions tp
+        using current_range_from_value, current_range_to_value, table_name;
-        WHERE
-            ((current_range_from_value::timestamptz >= tp.from_value::timestamptz AND current_range_from_value < tp.to_value::timestamptz) OR
-            (current_range_to_value::timestamptz > tp.from_value::timestamptz AND current_range_to_value::timestamptz < tp.to_value::timestamptz)) AND
-            parent_table = table_name;
 
-        IF found THEN
+        IF manual_partition is not NULL THEN
             RAISE 'partition % with the range from % to % does not align with the initial partition given the partition interval',
             manual_partition::text,
             manual_partition_from_value_text,
             manual_partition_to_value_text
-			USING HINT = 'Only use partitions of the same size, without gaps between partitions.';
+            USING HINT = 'Only use partitions of the same size, without gaps between partitions.';
         END IF;
 
         IF partition_column_type = 'date'::regtype THEN
@@ -194,7 +213,8 @@ BEGIN
             SELECT current_range_from_value::timestamptz::text INTO current_range_from_value_text;
             SELECT current_range_to_value::timestamptz::text INTO current_range_to_value_text;
         ELSE
-            RAISE 'type of the partition column of the table % must be date, timestamp or timestamptz', table_name;
+            EXECUTE format('SELECT $1%s::text', custom_cast) INTO current_range_from_value_text using current_range_from_value;
+            EXECUTE format('SELECT $1%s::text', custom_cast) INTO current_range_to_value_text using current_range_to_value;
         END IF;
 
         -- use range values within the name of partition to have unique partition names
@@ -212,7 +232,7 @@ BEGIN
 END;
 $$;
 COMMENT ON FUNCTION pg_catalog.get_missing_time_partition_ranges(
-	table_name regclass,
+    table_name regclass,
     partition_interval INTERVAL,
     to_value timestamptz,
     from_value timestamptz)

src/test/regress/expected/multi_partitioning.out

@@ -4059,6 +4059,47 @@ BEGIN;
 
 ROLLBACK;
 set client_min_messages to notice;
+-- 7) test with bigint partition column
+CREATE FUNCTION nanos_to_timestamptz(nanos bigint) RETURNS timestamptz LANGUAGE plpgsql AS
+$$
+DECLARE
+  value timestamptz;
+BEGIN
+  select to_timestamp(nanos * 1.0 / 1000000000) into value;
+  return value;
+END;
+$$;
+CREATE CAST (bigint AS timestamptz) WITH FUNCTION nanos_to_timestamptz(bigint);
+CREATE FUNCTION timestamptz_to_nanos(ts timestamptz) RETURNS bigint LANGUAGE plpgsql AS
+$$
+DECLARE
+  value bigint;
+BEGIN
+  select extract(epoch from ts) * 1000000000 into value;
+  return value;
+END;
+$$;
+CREATE CAST (timestamptz AS bigint) WITH FUNCTION timestamptz_to_nanos(timestamptz);
+CREATE TABLE bigint_partitioned_table (timestamp bigint, description text) partition by range (timestamp);
+BEGIN;
+  SELECT create_time_partitions('bigint_partitioned_table', INTERVAL '1 month', '2023-05-01', '2023-01-1');
+ create_time_partitions
+---------------------------------------------------------------------
+ t
+(1 row)
+
+  SELECT * FROM time_partitions WHERE parent_table = 'bigint_partitioned_table'::regclass ORDER BY 3;
+       parent_table       | partition_column |             partition             |     from_value      |      to_value       | access_method
+---------------------------------------------------------------------
+ bigint_partitioned_table | timestamp        | bigint_partitioned_table_p2023_01 | 1672560000000000000 | 1675238400000000000 | heap
+ bigint_partitioned_table | timestamp        | bigint_partitioned_table_p2023_02 | 1675238400000000000 | 1677657600000000000 | heap
+ bigint_partitioned_table | timestamp        | bigint_partitioned_table_p2023_03 | 1677657600000000000 | 1680332400000000000 | heap
+ bigint_partitioned_table | timestamp        | bigint_partitioned_table_p2023_04 | 1680332400000000000 | 1682924400000000000 | heap
+(4 rows)
+
+ROLLBACK;
+DROP CAST (bigint AS timestamptz);
+DROP CAST (timestamptz AS bigint);
 -- c) test drop_old_time_partitions
 -- 1) test with date partitioned table
 CREATE TABLE date_partitioned_table_to_exp (event_date date, event int) partition by range (event_date);
@@ -4359,10 +4400,13 @@ SELECT a, b FROM stxdinp GROUP BY 1, 2;
 (10 rows)
 
 DROP SCHEMA partitioning_schema CASCADE;
-NOTICE:  drop cascades to 5 other objects
+NOTICE:  drop cascades to 8 other objects
 DETAIL:  drop cascades to table "schema-test"
 drop cascades to table another_distributed_table
 drop cascades to table distributed_parent_table
+drop cascades to function nanos_to_timestamptz(bigint)
+drop cascades to function timestamptz_to_nanos(timestamp with time zone)
+drop cascades to table bigint_partitioned_table
 drop cascades to table part_table_with_very_long_name
 drop cascades to table stxdinp
 RESET search_path;

src/test/regress/sql/multi_partitioning.sql

@@ -1839,6 +1839,40 @@ BEGIN;
   SELECT * FROM time_partitions WHERE parent_table = 'date_partitioned_citus_local_table'::regclass ORDER BY 3;
 ROLLBACK;
 set client_min_messages to notice;
+
+-- 7) test with bigint partition column
+CREATE FUNCTION nanos_to_timestamptz(nanos bigint) RETURNS timestamptz LANGUAGE plpgsql AS
+$$
+DECLARE
+  value timestamptz;
+BEGIN
+  select to_timestamp(nanos * 1.0 / 1000000000) into value;
+  return value;
+END;
+$$;
+CREATE CAST (bigint AS timestamptz) WITH FUNCTION nanos_to_timestamptz(bigint);
+
+CREATE FUNCTION timestamptz_to_nanos(ts timestamptz) RETURNS bigint LANGUAGE plpgsql AS
+$$
+DECLARE
+  value bigint;
+BEGIN
+  select extract(epoch from ts) * 1000000000 into value;
+  return value;
+END;
+$$;
+CREATE CAST (timestamptz AS bigint) WITH FUNCTION timestamptz_to_nanos(timestamptz);
+
+CREATE TABLE bigint_partitioned_table (timestamp bigint, description text) partition by range (timestamp);
+
+BEGIN;
+  SELECT create_time_partitions('bigint_partitioned_table', INTERVAL '1 month', '2023-05-01', '2023-01-1');
+  SELECT * FROM time_partitions WHERE parent_table = 'bigint_partitioned_table'::regclass ORDER BY 3;
+ROLLBACK;
+
+DROP CAST (bigint AS timestamptz);
+DROP CAST (timestamptz AS bigint);
+
 -- c) test drop_old_time_partitions
 -- 1) test with date partitioned table
 CREATE TABLE date_partitioned_table_to_exp (event_date date, event int) partition by range (event_date);