Show details Citus query plan instead of just 'Custom Scan'

Author: RoystonS

web/pgadmin/misc/static/explain/img/ex_citus.svg

@@ -41,10 +41,85 @@ const ImageMapper = {
     'image': 'ex_bmp_or.svg',
     'image_text': 'Bitmap OR',
   },
+  'Citus Job': function(data) {
+    // A 'Citus Job' represents a distributed query operation.
+    // The details of the distributed operation are in the sub-plans,
+    // but this node contains task count information, showing how many shards
+    // the query is being distributed to.
+
+    const taskCount = data['Task Count'];
+    const tasksShown = data['Tasks Shown'];
+    
+    // "Task Count" is the number of shard operations being run.
+    // "Tasks Shown" is either "All" or "One of N" depending on whether the returned query plan
+    // contains one sample task or all of them.
+
+    // We show single-shard or multi-shard with different images, and we show the
+    // literal value of 'Tasks Shown' as the image text.
+
+    const image = (taskCount === 1)
+      ? 'ex_citus_distributed_one_of_one.svg'
+      : 'ex_citus_distributed_one_of_many.svg';
+
+    return {
+      'image': image,
+      'image_text': tasksShown
+    };
+  },
+  'Citus Task': function(data) {
+    // A 'Citus Task' represents a Task executed on a particular worker node.
+    // The details of the Task are in the sub-plans, so for this node we just show
+    // some details of the worker node.
+
+    const node = data['Node'];
+    // "Node" has a value like "host=citus-worker-7 port=8394 dbname=postgres"
+    // That's a bit long to display, so we shrink it to 'citus-worker-7:8394 postgres'
+    const hostMatch = node.match(/host=(\S+)/);
+    const portMatch = node.match(/port=(\S+)/);
+    const dbnameMatch = node.match(/dbname=(\S+)/);
+
+    const host = hostMatch[1] || '';
+    let port = portMatch[1] || '';
+    if (port === '5432') {
+      // Default port. Don't bother showing.
+      port = '';
+    }
+    const dbname = dbnameMatch[1] || '';
+    let text = `Task ${host}`;
+    if (port) {
+      text += `:${port}`;
+    }
+    if (dbname) {
+      text += ` ${dbname}`;
+    }
+    return {
+      'image': 'ex_citus_worker_task.svg',
+      'image_text': text
+    };
+  },
   'CTE Scan': {
     'image': 'ex_cte_scan.svg',
     'image_text': 'CTE Scan',
   },
+  'Custom Scan': function(data) {
+    const customPlanProvider = data['Custom Plan Provider'];
+
+    let image;
+
+    switch (customPlanProvider) {
+    case 'Citus Adaptive':
+      image = 'ex_citus.svg';
+      break;
+    default:
+      image = 'ex_unknown.svg';
+      break;
+    }
+
+    return {
+      'image': image,
+      'image_text': data['Custom Plan Provider']
+    };
+  },
   'Function Scan': {
     'image': 'ex_result.svg',
     'image_text': 'Function Scan',

web/pgadmin/misc/static/explain/img/ex_citus_distributed_one_of_many.svg

@@ -350,6 +350,46 @@ function parsePlan(data, ctx) {
     }
   }
 
+  const citusDistributedQuery = data['Distributed Query'];
+  if (citusDistributedQuery) {
+    // This is a Citus Distributed Query plan.
+    // It contains a 'Job' with one or more 'Tasks' in it.
+    // We'll convert those Tasks into sub-Plans of this main plan and process it
+    // with the regular Plan layout code.
+    delete data['Distributed Query'];
+
+    // Convert the Job into a 'Citus Job' sub-plan.
+    // That allows us to show details of the Task count etc.
+    const citusJob = citusDistributedQuery['Job'];
+    const jobPlan = {
+      'Node Type': 'Citus Job',
+      ...citusJob
+    };
+    data['Plans'] = [jobPlan];
+
+    // Convert each of the Tasks into 'Citus Task' sub-plans of the Job plan.
+    const citusTasks = jobPlan['Tasks'];
+    if (citusTasks) {
+      delete jobPlan['Tasks'];
+
+      const citusTaskPlans = citusTasks.map(citusJobTask => {
+        const taskPlan = {
+          'Node Type': 'Citus Task',
+          ...citusJobTask
+        };
+
+        // A Citus Task contains a 'Remote Plan' which is the actual plan
+        // executed on the worker nodes. It's actually an array of arrays.
+        const remotePlan = taskPlan['Remote Plan'];
+        delete taskPlan['Remote Plan'];
+        // A Remote Plan is an array of arrays of Plans.
+        taskPlan['Plans'] = remotePlan.flatMap(arr => arr.map(planLevel1Entry => planLevel1Entry['Plan']));
+        return taskPlan;
+      });
+      jobPlan['Plans'] = citusTaskPlans;
+    }
+  }
+
   // Start calculating xpos, ypos, width and height for child plans if any
   if ('Plans' in data) {
     data['width'] += offsetX;