Add alter index set statistics support (#4455)

* Add alter index set statistics support

* Use attNum instead of attName

src/backend/distributed/commands/index.c

@@ -785,6 +785,7 @@ PostprocessIndexStmt(Node *node, const char *queryString)
  *
  * ALTER INDEX SET ()
  * ALTER INDEX RESET ()
+ * ALTER INDEX ALTER COLUMN SET STATISTICS
  */
 void
 ErrorIfUnsupportedAlterIndexStmt(AlterTableStmt *alterTableStatement)
@@ -801,6 +802,7 @@ ErrorIfUnsupportedAlterIndexStmt(AlterTableStmt *alterTableStatement)
 			case AT_SetRelOptions:  /* SET (...) */
 			case AT_ResetRelOptions:    /* RESET (...) */
 			case AT_ReplaceRelOptions:  /* replace entire option list */
+			case AT_SetStatistics:  /* SET STATISTICS */
 			{
 				/* this command is supported by Citus */
 				break;
@@ -814,8 +816,8 @@ ErrorIfUnsupportedAlterIndexStmt(AlterTableStmt *alterTableStatement)
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("alter index ... set tablespace ... "
 								"is currently unsupported"),
-						 errdetail("Only RENAME TO, SET (), and RESET () "
+						 errdetail("Only RENAME TO, SET (), RESET () "
-								   "are supported.")));
+								   "and SET STATISTICS are supported.")));
 				return; /* keep compiler happy */
 			}
 		}

src/backend/distributed/commands/statistics.c

@@ -41,10 +41,19 @@
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
 #include "utils/lsyscache.h"
+#include "utils/relcache.h"
 #include "utils/ruleutils.h"
 #include "utils/syscache.h"
 
+#define DEFAULT_STATISTICS_TARGET -1
+#define ALTER_INDEX_COLUMN_SET_STATS_COMMAND \
+	"ALTER INDEX %s ALTER COLUMN %d SET STATISTICS %d"
+
+static List * GetAlterIndexStatisticsCommands(Oid indexOid);
 static List * GetExplicitStatisticsIdList(Oid relationId);
+static char * GenerateAlterIndexColumnSetStatsCommand(char *indexNameWithSchema,
+													  int16 attnum,
+													  int32 attstattarget);
 static Oid GetRelIdByStatsOid(Oid statsOid);
 static char * CreateAlterCommandIfOwnerNotDefault(Oid statsOid);
 #if PG_VERSION_NUM >= PG_VERSION_13
@@ -403,14 +412,14 @@ PreprocessAlterStatisticsOwnerStmt(Node *node, const char *queryString)
 
 /*
  * GetExplicitStatisticsCommandList returns the list of DDL commands to create
- * statistics that are explicitly created for the table with relationId. See
+ * or alter statistics that are explicitly created for the table with relationId.
- * comment of GetExplicitStatisticsIdList function.
+ * This function gets called when distributing the table with relationId.
+ * See comment of GetExplicitStatisticsIdList function.
  */
 List *
 GetExplicitStatisticsCommandList(Oid relationId)
 {
-	List *createStatisticsCommandList = NIL;
+	List *explicitStatisticsCommandList = NIL;
-	List *alterStatisticsCommandList = NIL;
 
 	PushOverrideEmptySearchPath(CurrentMemoryContext);
 
@@ -419,11 +428,12 @@ GetExplicitStatisticsCommandList(Oid relationId)
 	Oid statisticsId = InvalidOid;
 	foreach_oid(statisticsId, statisticsIdList)
 	{
+		/* we need create commands for already created stats before distribution */
 		char *createStatisticsCommand = pg_get_statisticsobj_worker(statisticsId,
 																	false);
 
-		createStatisticsCommandList =
+		explicitStatisticsCommandList =
-			lappend(createStatisticsCommandList,
+			lappend(explicitStatisticsCommandList,
 					makeTableDDLCommandString(createStatisticsCommand));
 #if PG_VERSION_NUM >= PG_VERSION_13
 
@@ -433,8 +443,8 @@ GetExplicitStatisticsCommandList(Oid relationId)
 
 		if (alterStatisticsTargetCommand != NULL)
 		{
-			alterStatisticsCommandList =
+			explicitStatisticsCommandList =
-				lappend(alterStatisticsCommandList,
+				lappend(explicitStatisticsCommandList,
 						makeTableDDLCommandString(alterStatisticsTargetCommand));
 		}
 #endif
@@ -445,19 +455,31 @@ GetExplicitStatisticsCommandList(Oid relationId)
 
 		if (alterStatisticsOwnerCommand != NULL)
 		{
-			alterStatisticsCommandList =
+			explicitStatisticsCommandList =
-				lappend(alterStatisticsCommandList,
+				lappend(explicitStatisticsCommandList,
 						makeTableDDLCommandString(alterStatisticsOwnerCommand));
 		}
 	}
 
+	/* find indexes on current relation, in case of modified stats target */
+	Relation relation = relation_open(relationId, AccessShareLock);
+	List *indexOidList = RelationGetIndexList(relation);
+	relation_close(relation, NoLock);
+
+	Oid indexId = InvalidOid;
+	foreach_oid(indexId, indexOidList)
+	{
+		/* we need alter index commands for altered targets on expression indexes */
+		List *alterIndexStatisticsCommands = GetAlterIndexStatisticsCommands(indexId);
+
+		explicitStatisticsCommandList =
+			list_concat(explicitStatisticsCommandList, alterIndexStatisticsCommands);
+	}
+
 	/* revert back to original search_path */
 	PopOverrideSearchPath();
 
-	createStatisticsCommandList = list_concat(createStatisticsCommandList,
+	return explicitStatisticsCommandList;
-											  alterStatisticsCommandList);
-
-	return createStatisticsCommandList;
 }
 
 
@@ -506,6 +528,49 @@ GetExplicitStatisticsSchemaIdList(Oid relationId)
 }
 
 
+/*
+ * GetAlterIndexStatisticsCommands returns the list of ALTER INDEX .. ALTER COLUMN ..
+ * SET STATISTICS commands, based on non default targets of the index with given id.
+ * Note that this function only looks for expression indexes, since this command is
+ * valid for only expression indexes.
+ */
+static List *
+GetAlterIndexStatisticsCommands(Oid indexOid)
+{
+	List *alterIndexStatisticsCommandList = NIL;
+	int16 exprCount = 1;
+	while (true)
+	{
+		HeapTuple attTuple = SearchSysCacheAttNum(indexOid, exprCount);
+
+		if (!HeapTupleIsValid(attTuple))
+		{
+			break;
+		}
+
+		Form_pg_attribute targetAttr = (Form_pg_attribute) GETSTRUCT(attTuple);
+		if (targetAttr->attstattarget != DEFAULT_STATISTICS_TARGET)
+		{
+			char *indexNameWithSchema = generate_qualified_relation_name(indexOid);
+
+			char *command =
+				GenerateAlterIndexColumnSetStatsCommand(indexNameWithSchema,
+														targetAttr->attnum,
+														targetAttr->attstattarget);
+
+			alterIndexStatisticsCommandList =
+				lappend(alterIndexStatisticsCommandList,
+						makeTableDDLCommandString(command));
+		}
+
+		ReleaseSysCache(attTuple);
+		exprCount++;
+	}
+
+	return alterIndexStatisticsCommandList;
+}
+
+
 /*
  * GetExplicitStatisticsIdList returns a list of OIDs corresponding to the statistics
  * that are explicitly created on the relation with relationId. That means,
@@ -553,6 +618,28 @@ GetExplicitStatisticsIdList(Oid relationId)
 }
 
 
+/*
+ * GenerateAlterIndexColumnSetStatsCommand returns a string in form of 'ALTER INDEX ..
+ * ALTER COLUMN .. SET STATISTICS ..' which will be used to create a DDL command to
+ * send to workers.
+ */
+static char *
+GenerateAlterIndexColumnSetStatsCommand(char *indexNameWithSchema,
+										int16 attnum,
+										int32 attstattarget)
+{
+	StringInfoData command;
+	initStringInfo(&command);
+	appendStringInfo(&command,
+					 ALTER_INDEX_COLUMN_SET_STATS_COMMAND,
+					 indexNameWithSchema,
+					 attnum,
+					 attstattarget);
+
+	return command.data;
+}
+
+
 /*
  * GetRelIdByStatsOid returns the relation id for given statistics oid.
  * If statistics doesn't exist, returns InvalidOid.

src/test/regress/expected/alter_index.out

@@ -0,0 +1,66 @@
+CREATE SCHEMA alterindex;
+SET search_path TO "alterindex";
+SET citus.next_shard_id TO 980000;
+SET client_min_messages TO WARNING;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+-- test alter index set statistics
+CREATE TABLE t1 (a int, b int);
+SELECT create_distributed_table('t1','a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE INDEX test_idx on t1 ((a+b));
+ALTER INDEX test_idx ALTER COLUMN 1 SET STATISTICS 4646;
+ALTER INDEX test_idx ALTER COLUMN 1 SET STATISTICS -4646;
+ERROR:  statistics target -4646 is too low
+ALTER INDEX test_idx ALTER COLUMN 3 SET STATISTICS 4646;
+ERROR:  column number 3 of relation "test_idx" does not exist
+-- test alter index set statistics before distribution
+CREATE TABLE t2 (a int, b int);
+CREATE INDEX test_idx2 on t2 ((a+b), (a-b), (a*b));
+ALTER INDEX test_idx2 ALTER COLUMN 2 SET STATISTICS 3737;
+ALTER INDEX test_idx2 ALTER COLUMN 3 SET STATISTICS 3737;
+ALTER INDEX test_idx2 ALTER COLUMN 2 SET STATISTICS 99999;
+WARNING:  lowering statistics target to 10000
+SELECT create_distributed_table('t2','a');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- verify statistics is set
+SELECT c.relname, a.attstattarget
+FROM pg_attribute a
+JOIN pg_class c ON a.attrelid = c.oid AND c.relname LIKE 'test\_idx%'
+ORDER BY c.relname, a.attnum;
+  relname  | attstattarget
+---------------------------------------------------------------------
+ test_idx  |          4646
+ test_idx2 |            -1
+ test_idx2 |         10000
+ test_idx2 |          3737
+(4 rows)
+
+\c - - - :worker_1_port
+SELECT c.relname, a.attstattarget
+FROM pg_attribute a
+JOIN pg_class c ON a.attrelid = c.oid AND c.relname LIKE 'test\_idx%'
+ORDER BY c.relname, a.attnum;
+     relname      | attstattarget
+---------------------------------------------------------------------
+ test_idx2_980004 |            -1
+ test_idx2_980004 |         10000
+ test_idx2_980004 |          3737
+ test_idx2_980006 |            -1
+ test_idx2_980006 |         10000
+ test_idx2_980006 |          3737
+ test_idx_980000  |          4646
+ test_idx_980002  |          4646
+(8 rows)
+
+\c - - - :master_port
+SET client_min_messages TO WARNING;
+DROP SCHEMA alterindex CASCADE;

src/test/regress/multi_schedule

@@ -47,7 +47,7 @@ test: multi_behavioral_analytics_basics multi_behavioral_analytics_single_shard_
 test: multi_shard_update_delete recursive_dml_with_different_planners_executors
 test: insert_select_repartition window_functions dml_recursive multi_insert_select_window
 test: multi_insert_select_conflict citus_table_triggers
-test: multi_row_insert insert_select_into_local_table multi_create_table_new_features
+test: multi_row_insert insert_select_into_local_table multi_create_table_new_features alter_index
 
 # following should not run in parallel because it relies on connection counts to workers
 test: insert_select_connection_leak

src/test/regress/output/multi_alter_table_statements.source

@@ -1132,7 +1132,7 @@ SELECT relname, reloptions FROM pg_class WHERE relname LIKE 'hash_dist_pkey%' OR
 -- verify error message on ALTER INDEX, SET TABLESPACE is unsupported
 ALTER INDEX hash_dist_pkey SET TABLESPACE foo;
 ERROR:  alter index ... set tablespace ... is currently unsupported
-DETAIL:  Only RENAME TO, SET (), and RESET () are supported.
+DETAIL:  Only RENAME TO, SET (), RESET () and SET STATISTICS are supported.
 -- verify that we can add indexes with new storage options
 CREATE UNIQUE INDEX another_index ON hash_dist(id) WITH (fillfactor=50);
 -- show the index and its storage options on coordinator, then workers

src/test/regress/sql/alter_index.sql

@@ -0,0 +1,39 @@
+CREATE SCHEMA alterindex;
+
+SET search_path TO "alterindex";
+SET citus.next_shard_id TO 980000;
+SET client_min_messages TO WARNING;
+SET citus.shard_count TO 4;
+SET citus.shard_replication_factor TO 1;
+
+-- test alter index set statistics
+CREATE TABLE t1 (a int, b int);
+SELECT create_distributed_table('t1','a');
+CREATE INDEX test_idx on t1 ((a+b));
+ALTER INDEX test_idx ALTER COLUMN 1 SET STATISTICS 4646;
+ALTER INDEX test_idx ALTER COLUMN 1 SET STATISTICS -4646;
+ALTER INDEX test_idx ALTER COLUMN 3 SET STATISTICS 4646;
+
+-- test alter index set statistics before distribution
+CREATE TABLE t2 (a int, b int);
+CREATE INDEX test_idx2 on t2 ((a+b), (a-b), (a*b));
+ALTER INDEX test_idx2 ALTER COLUMN 2 SET STATISTICS 3737;
+ALTER INDEX test_idx2 ALTER COLUMN 3 SET STATISTICS 3737;
+ALTER INDEX test_idx2 ALTER COLUMN 2 SET STATISTICS 99999;
+SELECT create_distributed_table('t2','a');
+
+-- verify statistics is set
+SELECT c.relname, a.attstattarget
+FROM pg_attribute a
+JOIN pg_class c ON a.attrelid = c.oid AND c.relname LIKE 'test\_idx%'
+ORDER BY c.relname, a.attnum;
+
+\c - - - :worker_1_port
+SELECT c.relname, a.attstattarget
+FROM pg_attribute a
+JOIN pg_class c ON a.attrelid = c.oid AND c.relname LIKE 'test\_idx%'
+ORDER BY c.relname, a.attnum;
+\c - - - :master_port
+
+SET client_min_messages TO WARNING;
+DROP SCHEMA alterindex CASCADE;