Propagate Drop Statistics

src/backend/distributed/commands/distribute_object_ops.c

@@ -414,6 +414,13 @@ static DistributeObjectOps Schema_Rename = {
 	.postprocess = NULL,
 	.address = AlterSchemaRenameStmtObjectAddress,
 };
+static DistributeObjectOps Statistics_Drop = {
+	.deparse = NULL,
+	.qualify = QualifyDropStatisticsStmt,
+	.preprocess = PreprocessDropStatisticsStmt,
+	.postprocess = NULL,
+	.address = NULL,
+};
 static DistributeObjectOps Table_AlterTable = {
 	.deparse = NULL,
 	.qualify = NULL,
@@ -806,6 +813,11 @@ GetDistributeObjectOps(Node *node)
 					return &Schema_Drop;
 				}
 
+				case OBJECT_STATISTIC_EXT:
+				{
+					return &Statistics_Drop;
+				}
+
 				case OBJECT_TABLE:
 				{
 					return &Table_Drop;

src/backend/distributed/commands/statistics.c

@@ -8,6 +8,10 @@
  *
  *    CREATE STATISTICS ... queries.
  *
+ *    We also support dropping statistics from all the worker nodes in form of
+ *
+ *    DROP STATISTICS ... queries.
+ *
  * Copyright (c) Citus Data, Inc.
  *
  *-------------------------------------------------------------------------
@@ -40,6 +44,7 @@
 #include "utils/syscache.h"
 
 static List * GetExplicitStatisticsIdList(Oid relationId);
+static Oid GetRelIdByStatsOid(Oid statsOid);
 
 
 /*
@@ -127,6 +132,64 @@ CreateStatisticsStmtObjectAddress(Node *node, bool missingOk)
 }
 
 
+/*
+ * PreprocessDropStatisticsStmt is called during the planning phase for
+ * DROP STATISTICS.
+ */
+List *
+PreprocessDropStatisticsStmt(Node *node, const char *queryString)
+{
+	DropStmt *dropStatisticsStmt = castNode(DropStmt, node);
+	Assert(dropStatisticsStmt->removeType == OBJECT_STATISTIC_EXT);
+
+	if (!ShouldPropagate())
+	{
+		return NIL;
+	}
+
+	QualifyTreeNode((Node *) dropStatisticsStmt);
+
+	List *ddlJobs = NIL;
+	List *processedStatsOids = NIL;
+	List *objectNameList = NULL;
+	foreach_ptr(objectNameList, dropStatisticsStmt->objects)
+	{
+		Oid statsOid = get_statistics_object_oid(objectNameList,
+												 dropStatisticsStmt->missing_ok);
+
+		if (list_member_oid(processedStatsOids, statsOid))
+		{
+			/* skip duplicate entries in DROP STATISTICS */
+			continue;
+		}
+
+		processedStatsOids = lappend_oid(processedStatsOids, statsOid);
+
+		Oid relationId = GetRelIdByStatsOid(statsOid);
+
+		if (!OidIsValid(relationId) || !IsCitusTable(relationId))
+		{
+			continue;
+		}
+
+		char *ddlCommand = DeparseDropStatisticsStmt(objectNameList,
+													 dropStatisticsStmt->missing_ok);
+
+		DDLJob *ddlJob = palloc0(sizeof(DDLJob));
+
+		ddlJob->targetRelationId = relationId;
+		ddlJob->concurrentIndexCmd = false;
+		ddlJob->startNewTransaction = false;
+		ddlJob->commandString = ddlCommand;
+		ddlJob->taskList = DDLTaskList(relationId, ddlCommand);
+
+		ddlJobs = lappend(ddlJobs, ddlJob);
+	}
+
+	return ddlJobs;
+}
+
+
 /*
  * GetExplicitStatisticsCommandList returns the list of DDL commands to create
  * statistics that are explicitly created for the table with relationId. See
@@ -248,3 +311,24 @@ GetExplicitStatisticsIdList(Oid relationId)
 
 	return statisticsIdList;
 }
+
+
+/*
+ * GetRelIdByStatsOid returns the relation id for given statistics oid.
+ * If statistics doesn't exist, returns InvalidOid.
+ */
+static Oid
+GetRelIdByStatsOid(Oid statsOid)
+{
+	HeapTuple tup = SearchSysCache1(STATEXTOID, ObjectIdGetDatum(statsOid));
+
+	if (!HeapTupleIsValid(tup))
+	{
+		return InvalidOid;
+	}
+
+	Form_pg_statistic_ext statisticsForm = (Form_pg_statistic_ext) GETSTRUCT(tup);
+	ReleaseSysCache(tup);
+
+	return statisticsForm->stxrelid;
+}

src/backend/distributed/deparser/deparse_statistics_stmts.c

@@ -21,6 +21,7 @@
 #include "utils/builtins.h"
 
 static void AppendCreateStatisticsStmt(StringInfo buf, CreateStatsStmt *stmt);
+static void AppendDropStatisticsStmt(StringInfo buf, List *nameList, bool ifExists);
 static void AppendStatisticsName(StringInfo buf, CreateStatsStmt *stmt);
 static void AppendStatTypes(StringInfo buf, CreateStatsStmt *stmt);
 static void AppendColumnNames(StringInfo buf, CreateStatsStmt *stmt);
@@ -40,6 +41,18 @@ DeparseCreateStatisticsStmt(Node *node)
 }
 
 
+char *
+DeparseDropStatisticsStmt(List *nameList, bool ifExists)
+{
+	StringInfoData str;
+	initStringInfo(&str);
+
+	AppendDropStatisticsStmt(&str, nameList, ifExists);
+
+	return str.data;
+}
+
+
 static void
 AppendCreateStatisticsStmt(StringInfo buf, CreateStatsStmt *stmt)
 {
@@ -66,6 +79,20 @@ AppendCreateStatisticsStmt(StringInfo buf, CreateStatsStmt *stmt)
 }
 
 
+static void
+AppendDropStatisticsStmt(StringInfo buf, List *nameList, bool ifExists)
+{
+	appendStringInfoString(buf, "DROP STATISTICS ");
+
+	if (ifExists)
+	{
+		appendStringInfoString(buf, "IF EXISTS ");
+	}
+
+	appendStringInfo(buf, "%s", NameListToQuotedString(nameList));
+}
+
+
 static void
 AppendStatisticsName(StringInfo buf, CreateStatsStmt *stmt)
 {

src/backend/distributed/deparser/qualify_statistics_stmt.c

@@ -48,3 +48,33 @@ QualifyCreateStatisticsStmt(Node *node)
 		stmt->defnames = MakeNameListFromRangeVar(stat);
 	}
 }
+
+
+/*
+ * QualifyDropStatisticsStmt qualifies DropStmt's with schema name for
+ * DROP STATISTICS statements.
+ */
+void
+QualifyDropStatisticsStmt(Node *node)
+{
+	DropStmt *dropStatisticsStmt = castNode(DropStmt, node);
+	Assert(dropStatisticsStmt->removeType == OBJECT_STATISTIC_EXT);
+
+	List *objectNameListWithSchema = NIL;
+	List *objectNameList = NULL;
+	foreach_ptr(objectNameList, dropStatisticsStmt->objects)
+	{
+		RangeVar *stat = makeRangeVarFromNameList(objectNameList);
+
+		if (stat->schemaname == NULL)
+		{
+			Oid schemaOid = RangeVarGetCreationNamespace(stat);
+			stat->schemaname = get_namespace_name(schemaOid);
+		}
+
+		objectNameListWithSchema = lappend(objectNameListWithSchema,
+										   MakeNameListFromRangeVar(stat));
+	}
+
+	dropStatisticsStmt->objects = objectNameListWithSchema;
+}

src/backend/distributed/relay/relay_event_utility.c

@@ -349,6 +349,23 @@ RelayEventExtendNames(Node *parseTree, char *schemaName, uint64 shardId)
 			{
 				DropTriggerEventExtendNames(dropStmt, schemaName, shardId);
 			}
+			else if (objectType == OBJECT_STATISTIC_EXT)
+			{
+				List *shardStatisticsList = NIL;
+				List *objectNameList = NULL;
+				foreach_ptr(objectNameList, dropStmt->objects)
+				{
+					RangeVar *stat = makeRangeVarFromNameList(objectNameList);
+
+					SetSchemaNameIfNotExist(&stat->schemaname, schemaName);
+
+					AppendShardIdToName(&stat->relname, shardId);
+					shardStatisticsList = lappend(shardStatisticsList,
+												  MakeNameListFromRangeVar(stat));
+				}
+
+				dropStmt->objects = shardStatisticsList;
+			}
 			else
 			{
 				ereport(WARNING, (errmsg("unsafe object type in drop statement"),

src/include/distributed/commands.h

@@ -243,6 +243,7 @@ extern void ErrorIfDistributedAlterSeqOwnedBy(AlterSeqStmt *alterSeqStmt);
 extern List * PreprocessCreateStatisticsStmt(Node *node, const char *queryString);
 extern List * PostprocessCreateStatisticsStmt(Node *node, const char *queryString);
 extern ObjectAddress CreateStatisticsStmtObjectAddress(Node *node, bool missingOk);
+extern List * PreprocessDropStatisticsStmt(Node *node, const char *queryString);
 extern List * GetExplicitStatisticsCommandList(Oid relationId);
 extern List * GetExplicitStatisticsSchemaIdList(Oid relationId);
 

src/include/distributed/deparser.h

@@ -57,8 +57,10 @@ extern char * DeparseAlterSchemaRenameStmt(Node *stmt);
 
 /* forward declarations for deparse_statistics_stmts.c */
 extern char * DeparseCreateStatisticsStmt(Node *node);
+extern char * DeparseDropStatisticsStmt(List *nameList, bool ifExists);
 
 extern void QualifyCreateStatisticsStmt(Node *node);
+extern void QualifyDropStatisticsStmt(Node *node);
 
 /* forward declarations for deparse_type_stmts.c */
 extern char * DeparseCompositeTypeStmt(Node *stmt);

src/test/regress/expected/propagate_statistics.out

@@ -51,6 +51,25 @@ SELECT create_distributed_table ('test_stats3','a');
 
 (1 row)
 
+-- test dropping statistics
+CREATE TABLE test_stats4 (
+    a int,
+    b int
+);
+SELECT create_reference_table ('test_stats4');
+ create_reference_table
+---------------------------------------------------------------------
+
+(1 row)
+
+CREATE STATISTICS s3 ON a,b FROM test_stats3;
+CREATE STATISTICS sc2.s4 ON a,b FROM test_stats3;
+CREATE STATISTICS s5 ON a,b FROM test_stats4;
+-- s6 doesn't exist
+DROP STATISTICS IF EXISTS s3, sc2.s4, s6;
+DROP STATISTICS s5,s6;
+ERROR:  statistics object "statistics'Test.s6" does not exist
+DROP STATISTICS IF EXISTS s5,s5,s6,s6;
 \c - - - :worker_1_port
 SELECT stxname
 FROM pg_statistic_ext

src/test/regress/sql/propagate_statistics.sql

@@ -42,6 +42,21 @@ CREATE SCHEMA sc2;
 CREATE STATISTICS sc2."neW'Stat" ON a,b FROM test_stats3;
 SELECT create_distributed_table ('test_stats3','a');
 
+-- test dropping statistics
+CREATE TABLE test_stats4 (
+    a int,
+    b int
+);
+SELECT create_reference_table ('test_stats4');
+
+CREATE STATISTICS s3 ON a,b FROM test_stats3;
+CREATE STATISTICS sc2.s4 ON a,b FROM test_stats3;
+CREATE STATISTICS s5 ON a,b FROM test_stats4;
+-- s6 doesn't exist
+DROP STATISTICS IF EXISTS s3, sc2.s4, s6;
+DROP STATISTICS s5,s6;
+DROP STATISTICS IF EXISTS s5,s5,s6,s6;
+
 \c - - - :worker_1_port
 SELECT stxname
 FROM pg_statistic_ext