Fixed execution bug

Author: mbakli

include/catalog/nodes.h

@@ -23,6 +23,9 @@ extern Datum GetDBName();
 /* Coordinator/worker node info (host, port, role) for the local backend. */
 extern TaskNode *GetNodeInfo();
 
+/* Node (host, port) actually hosting relationId's shard for rand_tile. */
+extern TaskNode *GetShardHostNode(Oid relationId, int rand_tile);
+
 /* Table id of the tile assigned to a randomly-picked worker for relationId. */
 extern char* GetRandomTileId(Oid relationId, ExecTaskType taskType, int rand_tile);
 

sql/helper_functions/colocation_utils.sql

@@ -34,47 +34,38 @@ END;
 $$;
 
 ----------------------------------------------------------------------------------------------------------------------
--- colocate_multirelation colocates one multirelation with another
+-- colocate_shards physically moves table2's shards so each one lands on the
+-- same node as table1's shard sharing the same tile_key range (shardminvalue).
+-- Citus' own colocate_with option doesn't support range-distributed tables,
+-- so this does the move explicitly with citus_move_shard_placement() instead
+-- of relying on Citus' colocation groups.
 ----------------------------------------------------------------------------------------------------------------------
-CREATE OR REPLACE FUNCTION colocate_multirelation(table1 text, table2 text)
+DROP FUNCTION IF EXISTS colocate_multirelation;
+CREATE OR REPLACE FUNCTION colocate_shards(table1 text, table2 text)
     RETURNS boolean
 LANGUAGE plpgsql
 AS $$
 DECLARE
-i integer;
-    shard_info record;
-    node_info record;
-    node text;
-    shardid_test bigint;
+    tile_pair record;
 BEGIN
-    --Move one of them to a new place because one node contains 10 and the other contains 11
-    --It is an enterpise feature
-    --SELECT master_move_shard_placement(102660,'pgxl4', 5432,'pgxl2', 5432);
-    --For every node, update shards information
-    FOR node_info in SELECT * FROM master_get_active_worker_nodes()
+    FOR tile_pair IN
+        SELECT s2.shardid AS moving_shard,
+               n1.nodename AS target_node, n1.nodeport AS target_port,
+               n2.nodename AS source_node, n2.nodeport AS source_port
+        FROM pg_dist_shard s1
+        JOIN pg_dist_placement p1 ON p1.shardid = s1.shardid
+        JOIN pg_dist_node n1 ON n1.groupid = p1.groupid AND n1.noderole = 'primary'
+        JOIN pg_dist_shard s2 ON s2.shardminvalue = s1.shardminvalue
+        JOIN pg_dist_placement p2 ON p2.shardid = s2.shardid
+        JOIN pg_dist_node n2 ON n2.groupid = p2.groupid AND n2.noderole = 'primary'
+        WHERE s1.logicalrelid = table1::regclass
+          AND s2.logicalrelid = table2::regclass
+          AND (n1.nodename, n1.nodeport) IS DISTINCT FROM (n2.nodename, n2.nodeport)
     LOOP
-        i = 0;
-        --Get the shards of every table in every node
-        FOR shard_info in SELECT shard.shardid, shard.shardminvalue, shard.shardmaxvalue
-                          FROM pg_dist_placement AS placement, pg_dist_node AS node, pg_dist_shard As shard
-                          WHERE placement.groupid = node.groupid
-                            AND shard.logicalrelid = table1::regclass
-                            AND placement.shardid = shard.shardid
-                            AND node.noderole = 'primary'
-                            AND nodename=node_info.node_name
-        LOOP
-                SELECT shard.shardid
-                FROM pg_dist_placement AS placement, pg_dist_node AS node, pg_dist_shard As shard
-                WHERE placement.groupid = node.groupid
-                    AND shard.logicalrelid = table2::regclass
-                    AND placement.shardid = shard.shardid
-                    AND node.noderole = 'primary'
-                    AND nodename=node_info.node_name offset i limit 1 INTO shardid_test;
-                --RAISE NOTICE 'Update:%',shardid_test;
-                UPDATE pg_dist_shard SET shardminvalue = shard_info.shardminvalue, shardmaxvalue=shard_info.shardmaxvalue
-                WHERE shardid = shardid_test;
-                i := i + 1;
-        END LOOP;
+        PERFORM citus_move_shard_placement(tile_pair.moving_shard,
+                                           tile_pair.source_node, tile_pair.source_port,
+                                           tile_pair.target_node, tile_pair.target_port,
+                                           'block_writes');
     END LOOP;
     RETURN TRUE;
 END;
@@ -104,15 +95,12 @@ BEGIN
     /* create_range_shards() above already assigns reshuffled_table's shards
      * the correct tile_key-aligned ranges (1..shards), matching tableName's
      * own tile numbering by convention -- both tables tile the same way.
-     * This used to also call colocate_multirelation() to try to physically
-     * co-locate the two tables' shard placements onto the same nodes, but
-     * that function's actual placement-move step is unimplemented (commented
-     * out) and its remaining code just re-copies shard ranges by pairing
-     * shards positionally per node -- which silently overwrites the correct
-     * ranges above with wrong, duplicated ones whenever the two tables'
-     * shards aren't distributed identically across nodes (the common case).
-     * Until shard co-location is implemented for real, skip it entirely
-     * rather than corrupt the ranges that are already correct. */
+     * Citus' shard placement for the newly-created shards is otherwise
+     * independent of tableName's placement, so without this, a tile-key
+     * join between the two tables is a cross-node repartition even though
+     * both sides cover the same tile_key range. colocate_shards() moves
+     * each of reshuffled_table's shards onto tableName's matching-tile node. */
+    PERFORM colocate_shards(tableName, reshuffled_table);
 RETURN true;
 END;
 $$;

src/catalog/nodes.c

@@ -17,6 +17,7 @@
 #include <executor/spi.h>
 #include "catalog/nodes.h"
 #include <utils/lsyscache.h>
+#include <utils/builtins.h>
 #include "executor/executor_tasks.h"
 #include "catalog/table_ops.h"
 
@@ -118,6 +119,55 @@ char * GetRandomTileId(Oid relationId, ExecTaskType taskType, int rand_tile)
     return NULL;
 }
 
+/*
+ * GetShardHostNode looks up the (nodename, nodeport) actually hosting
+ * relationId's shard whose shardminvalue matches rand_tile, for dispatching
+ * a command to the specific worker that holds that tile's data (unlike
+ * GetNodeInfo(), which just picks a random node from pg_dist_node).
+ */
+extern TaskNode *
+GetShardHostNode(Oid relationId, int rand_tile)
+{
+    TaskNode *taskNode = (TaskNode *) palloc0(sizeof(TaskNode));
+    int spi_result = SPI_connect();
+    if (spi_result != SPI_OK_CONNECT)
+    {
+        elog(ERROR, "Could not connect to database using SPI");
+    }
+
+    StringInfo logicalrel = makeStringInfo();
+    if (IsReshuffledTable(relationId))
+        appendStringInfo(logicalrel, "%s.%s", Var_Schema, get_rel_name(relationId));
+    else
+        appendStringInfo(logicalrel, "%s", get_rel_name(relationId));
+
+    StringInfo catalogQuery = makeStringInfo();
+    appendStringInfo(catalogQuery,
+                     "SELECT node.nodename, node.nodeport\n"
+                     "FROM pg_dist_shard shard\n"
+                     "JOIN pg_dist_placement placement ON placement.shardid = shard.shardid\n"
+                     "JOIN pg_dist_node node ON node.groupid = placement.groupid AND node.noderole = 'primary'\n"
+                     "WHERE shard.logicalrelid = '%s'::regclass\n"
+                     "  AND shard.shardminvalue = %d::text",
+                     logicalrel->data, rand_tile);
+    spi_result = SPI_execute(catalogQuery->data, true, 1);
+    if (spi_result == SPI_OK_SELECT && SPI_processed > 0)
+    {
+        bool isNull;
+        HeapTuple row = SPI_copytuple(SPI_tuptable->vals[0]);
+        TupleDesc rowDescriptor = SPI_tuptable->tupdesc;
+        char *nodename = SPI_getvalue(row, rowDescriptor, 1);
+        taskNode->node = CStringGetTextDatum(nodename);
+        taskNode->port = DatumGetInt32(SPI_getbinval(row, rowDescriptor, 2, &isNull));
+    }
+    spi_result = SPI_finish();
+    if (spi_result != SPI_OK_FINISH)
+    {
+        elog(ERROR, "Could not disconnect from database using SPI");
+    }
+    return taskNode;
+}
+
 /* GetDBName returns the name of the database the current backend is connected to. */
 extern Datum
 GetDBName()

src/executor/multi_phase_executor.c

@@ -21,6 +21,7 @@
 #include <catalog/namespace.h>
 #include <access/xact.h>
 #include "utils/planner_utils.h"
+#include "utils/helper_functions.h"
 #include "planner/planner_strategies.h"
 
 static void ConstructNeighborScanQuery(Rte *tbl, char * query_string, STMultirelation *base,
@@ -36,6 +37,7 @@ static GeneralScan *ConstructGeneralQuery(DistributedSpatiotemporalQueryPlan *di
                                            MultiPhaseExecutor *multiPhaseExecutor);
 static void IndexReshuffledData(Rte *reshuffledTable, MultiPhaseExecutor *multiPhaseExecutor);
 static void ConstructPostProcessingPhase(CoordinatorLevelOperator *coordOp, MultiPhaseExecutor *multiPhaseExecutor);
+static char *EliminateShapeSegmentDuplicates(char *query_string, bool hasDistributedAggregate);
 
 
 
@@ -137,9 +139,12 @@ ColocateRte(STMultirelation *base, Rte *other)
         char *reshuffled_table = get_rel_name(cell->relid);
         Var *distributionColumn = DistPartitionKey(base->catalogTableInfo.table_oid);
         int shardCount = ShardIntervalCount(base->catalogTableInfo.table_oid);
-        /* Citus' CreateDistributedTable() expects the literal string "default"
-         * (not NULL) to mean "no explicit colocation group" -- IsColocateWithDefault()
-         * dereferences it directly and crashes on NULL. */
+        /* Citus rejects colocate_with for range-distributed tables
+         * ("colocate_with option is not supported for append / range
+         * distributed tables"), so this stays "default"; physical
+         * co-location with base is instead done after the fact by
+         * colocate_shards() in create_reshuffled_multirelation, which
+         * explicitly moves each shard onto base's matching-tile node. */
         char *parentRelationName = "default";
 
         DropReshuffledTableIfExists(citusRteNode->reshuffledTable);
@@ -182,9 +187,12 @@ createReshuffledTable(STMultirelation *base, STMultirelation *other)
     char *reshuffled_table = get_rel_name(other->catalogTableInfo.table_oid);
     Var *distributionColumn = DistPartitionKey(other->catalogTableInfo.table_oid);
     int shardCount = ShardIntervalCount(base->catalogTableInfo.table_oid);
-    /* Citus' CreateDistributedTable() expects the literal string "default"
-     * (not NULL) to mean "no explicit colocation group" -- IsColocateWithDefault()
-     * dereferences it directly and crashes on NULL. */
+    /* Citus rejects colocate_with for range-distributed tables
+     * ("colocate_with option is not supported for append / range
+     * distributed tables"), so this stays "default"; physical
+     * co-location with base is instead done after the fact by
+     * colocate_shards() in create_reshuffled_multirelation, which
+     * explicitly moves each shard onto base's matching-tile node. */
     char *parentRelationName = "default";
 
     DropReshuffledTableIfExists(other->catalogTableInfo.reshuffledTable);
@@ -402,6 +410,28 @@ IndexReshuffledData(Rte *reshuffledTable, MultiPhaseExecutor *multiPhaseExecutor
 
 }
 
+/*
+ * EliminateShapeSegmentDuplicates removes shape-segmentation duplicates from
+ * query_string by wrapping it as `SELECT DISTINCT * FROM (query_string) AS
+ * x`, deduping on whatever the query actually projects. Only applied when
+ * no distributed aggregate is involved (see hasDistributedAggregate at the
+ * call site) -- for an aggregate like count(*), the duplicates are already
+ * consumed before this outer DISTINCT would ever see them, so this is left
+ * as a no-op (returns NULL) for that case rather than risking a rewrite
+ * Citus' planner may reject for non-colocated repartition joins.
+ */
+static char *
+EliminateShapeSegmentDuplicates(char *query_string, bool hasDistributedAggregate)
+{
+    if (hasDistributedAggregate)
+        return NULL;
+
+    char *innerQuery = replaceWord(query_string, ";", " ");
+    StringInfo dedupedQuery = makeStringInfo();
+    appendStringInfo(dedupedQuery, "SELECT DISTINCT * FROM (%s) AS dedup_result", innerQuery);
+    return dedupedQuery->data;
+}
+
 /*
  * ConstructGeneralQuery assembles the final SQL text to execute: it unions
  * together the worker-phase task query for each strategy used in the plan
@@ -449,6 +479,29 @@ ConstructGeneralQuery(DistributedSpatiotemporalQueryPlan *distPlan, MultiPhaseEx
         else
             elog(ERROR, "The query executor could not identify the planner strategy");
     }
+    /* NonColocation/Colocation strategies reshuffle onto tiles built from a
+     * spatiotemporal shape, which can place the same row's shape-segmented
+     * copy in more than one tile so a boundary-crossing match isn't missed
+     * -- see checkQueryType's dupRemOperator->active assignment. That makes
+     * the worker-phase join above contain the same logical match
+     * (base_row1, base_row2) more than once. This has to be resolved here,
+     * on the flat two-table query before the coordinator/aggregate-rewriter
+     * wrapping below nests it inside a "select sum(...) from (...) as fQ"
+     * subquery -- once that wrapping happens, ExtractRangeTableEntryList
+     * sees three range table entries (the subquery plus its two inner
+     * tables) instead of the two this rewrite expects, and any duplicate
+     * rows have already been consumed by the inner aggregate anyway. */
+    if (distPlan->postProcessing->coordinatorLevelOperator->dupRemOperator->active)
+    {
+        bool hasDistributedAggregate = list_length(distPlan->postProcessing->distfuns) > 0;
+        char *deduped = EliminateShapeSegmentDuplicates(generalScan->query_string->data,
+                                                        hasDistributedAggregate);
+        if (deduped != NULL)
+        {
+            resetStringInfo(generalScan->query_string);
+            appendStringInfo(generalScan->query_string, "%s", deduped);
+        }
+    }
     /* Loop though the post processing tasks */
     if (generalScan->length == 0)
         appendStringInfo(generalScan->query_string,"%s",