From 3f7bc0cbf5b7858e9cd788726551c1495a2c9495 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Halil=20Ozan=20Akg=C3=BCl?= <hozanakgul@gmail.com>
Date: Tue, 6 Jun 2023 17:55:40 +0300
Subject: [PATCH] Single Shard Partition Column UDFs (#6964)

This PR fixes and tests:
- debug_equality_expression
- partition_column_id
---
 .../distributed/test/distribution_metadata.c  |  4 +++
 .../distributed/test/prune_shard_list.c       |  4 +++
 .../expected/single_shard_table_udfs.out      | 28 +++++++++++++++++++
 .../regress/sql/single_shard_table_udfs.sql   | 20 +++++++++++++
 4 files changed, 56 insertions(+)

diff --git a/src/backend/distributed/test/distribution_metadata.c b/src/backend/distributed/test/distribution_metadata.c
index cb378e9e3..c3bc7fb51 100644
--- a/src/backend/distributed/test/distribution_metadata.c
+++ b/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);
diff --git a/src/backend/distributed/test/prune_shard_list.c b/src/backend/distributed/test/prune_shard_list.c
index d83a645dc..a9f5e4a88 100644
--- a/src/backend/distributed/test/prune_shard_list.c
+++ b/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);
 
diff --git a/src/test/regress/expected/single_shard_table_udfs.out b/src/test/regress/expected/single_shard_table_udfs.out
index 2a4e59f82..861ba1477 100644
--- a/src/test/regress/expected/single_shard_table_udfs.out
+++ b/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;
diff --git a/src/test/regress/sql/single_shard_table_udfs.sql b/src/test/regress/sql/single_shard_table_udfs.sql
index 65b623a57..8d9057aca 100644
--- a/src/test/regress/sql/single_shard_table_udfs.sql
+++ b/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;