Single Shard Partition Column UDFs (#6964)

This PR fixes and tests:
- debug_equality_expression
- partition_column_id

src/backend/distributed/test/distribution_metadata.c

@@ -170,6 +170,10 @@ partition_column_id(PG_FUNCTION_ARGS)
 {
 	Oid distributedTableId = PG_GETARG_OID(0);
 	uint32 rangeTableId = 1;
+	if (!IsCitusTableType(distributedTableId, HASH_DISTRIBUTED))
+	{
+		ereport(ERROR, (errmsg("table needs to be hash distributed")));
+	}
 	Var *partitionColumn = PartitionColumn(distributedTableId, rangeTableId);
 
 	PG_RETURN_INT16((int16) partitionColumn->varattno);

src/backend/distributed/test/prune_shard_list.c

@@ -139,6 +139,10 @@ debug_equality_expression(PG_FUNCTION_ARGS)
 {
 	Oid distributedTableId = PG_GETARG_OID(0);
 	uint32 rangeTableId = 1;
+	if (!IsCitusTableType(distributedTableId, HASH_DISTRIBUTED))
+	{
+		ereport(ERROR, (errmsg("table needs to be hash distributed")));
+	}
 	Var *partitionColumn = PartitionColumn(distributedTableId, rangeTableId);
 	OpExpr *equalityExpression = MakeOpExpression(partitionColumn, BTEqualStrategyNumber);
 

src/test/regress/expected/single_shard_table_udfs.out

@@ -1068,5 +1068,33 @@ SELECT create_distributed_table('rep_tbl', NULL, colocate_with:='none');
 
 SELECT replicate_table_shards('rep_tbl');
 ERROR:  cannot replicate single shard tables' shards
+-- test debug_equality_expression
+CREATE FUNCTION debug_equality_expression(regclass)
+RETURNS cstring
+AS 'citus'
+LANGUAGE C STRICT;
+CREATE TABLE debug_tbl (a INT);
+SELECT create_distributed_table ('debug_tbl', NULL, colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT debug_equality_expression('debug_tbl'::regclass);
+ERROR:  table needs to be hash distributed
+-- test partition_column_id
+CREATE FUNCTION partition_column_id(regclass)
+RETURNS smallint
+AS 'citus'
+LANGUAGE C STRICT;
+CREATE TABLE partcol_tbl (a INT);
+SELECT create_distributed_table ('partcol_tbl', NULL, colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT partition_column_id('partcol_tbl'::regclass);
+ERROR:  table needs to be hash distributed
 SET client_min_messages TO WARNING;
 DROP SCHEMA null_dist_key_udfs CASCADE;

src/test/regress/sql/single_shard_table_udfs.sql

@@ -503,5 +503,25 @@ CREATE TABLE rep_tbl (a INT);
 SELECT create_distributed_table('rep_tbl', NULL, colocate_with:='none');
 SELECT replicate_table_shards('rep_tbl');
 
+-- test debug_equality_expression
+CREATE FUNCTION debug_equality_expression(regclass)
+RETURNS cstring
+AS 'citus'
+LANGUAGE C STRICT;
+
+CREATE TABLE debug_tbl (a INT);
+SELECT create_distributed_table ('debug_tbl', NULL, colocate_with:='none');
+SELECT debug_equality_expression('debug_tbl'::regclass);
+
+-- test partition_column_id
+CREATE FUNCTION partition_column_id(regclass)
+RETURNS smallint
+AS 'citus'
+LANGUAGE C STRICT;
+
+CREATE TABLE partcol_tbl (a INT);
+SELECT create_distributed_table ('partcol_tbl', NULL, colocate_with:='none');
+SELECT partition_column_id('partcol_tbl'::regclass);
+
 SET client_min_messages TO WARNING;
 DROP SCHEMA null_dist_key_udfs CASCADE;