fixed tile allocation issue

Author: mbakli

include/multirelation/multirelation_utils.h

@@ -92,6 +92,11 @@ typedef struct STMultirelations
     int stCount;
     int nonStCount;
     int length;
+    /* Count of nonStCount entries that are Citus reference tables -- these
+     * are already replicated to every node, so they never need reshuffling
+     * and shouldn't count as a "different distributed table" when deciding
+     * whether a join needs the NonColocation strategy. */
+    int refCount;
 } STMultirelations;
 
 /* True if relationId is registered as a distributed spatiotemporal table. */

src/catalog/table_ops.c

@@ -258,18 +258,23 @@ GetShapeCol(Oid relationId)
                      get_rel_name(relationId));
 
     spi_result = SPI_execute(catalogQuery->data, true, 1);
-    /* Read back the PROJ text */
-    if (spi_result == SPI_OK_SELECT)
+    /* Read back the PROJ text. getDistributedCol() legitimately returns
+     * NULL for a plain (non-spatiotemporal) table, e.g. a reference table
+     * joined alongside a distributed one -- calling DatumToString on that
+     * NULL Datum crashed instead of just reporting "no shape column". */
+    char *result = NULL;
+    if (spi_result == SPI_OK_SELECT && SPI_processed > 0)
     {
         TupleDesc rowDescriptor = SPI_tuptable->tupdesc;
         HeapTuple row = SPI_copytuple(SPI_tuptable->vals[0]);
         Datum distcol = SPI_getbinval(row, rowDescriptor, 1, &isNull);
-        spi_result = SPI_finish();
-        if (spi_result != SPI_OK_FINISH)
-        {
-            elog(ERROR, "Could not disconnect from database using SPI");
-        }
-        return DatumToString(distcol, TEXTOID);;
+        if (!isNull)
+            result = DatumToString(distcol, TEXTOID);
     }
-    return NULL;
+    spi_result = SPI_finish();
+    if (spi_result != SPI_OK_FINISH)
+    {
+        elog(ERROR, "Could not disconnect from database using SPI");
+    }
+    return result;
 }

src/general/shared_library_init.c

@@ -15,6 +15,7 @@
 #include "postgres.h"
 #include "miscadmin.h"
 #include "utils/elog.h"
+#include "utils/guc.h"
 #include "commands/explain.h"
 #include "planner/distributed_mobilitydb_planner.h"
 #include "planner/distributed_mobilitydb_explain.h"
@@ -36,4 +37,12 @@ _PG_init(void)
     planner_hook = distributed_mobilitydb_planner;
 
     ExplainOneQuery_hook = distributed_mobilitydb_explain;
+
+    /* Cross-table joins built by this extension's NonColocation/Colocation
+     * strategies aren't guaranteed to be physically co-located (see
+     * colocate_shards()), so Citus plans them as repartition joins. Turn
+     * that on as a session default here so users don't have to discover and
+     * set it themselves before a join query works; they can still override
+     * it with their own SET/RESET afterwards. */
+    SetConfigOption("citus.enable_repartition_joins", "on", PGC_SUSET, PGC_S_SESSION);
 }

src/planner/distributed_mobilitydb_planner.c

@@ -122,6 +122,21 @@ distributed_mobilitydb_planner_internal(Query *parse, const char *query_string,
     /* Query rewriter */
     if (list_length(distPlan->postProcessing->distfuns) > 0 )
         RewriterDistFuncs(parse, distPlan->postProcessing, query_string);
+    if (needsSpatiotemporalPlanning && list_length(distPlan->strategies) == 0
+            && list_length(distPlan->postProcessing->distfuns) == 0)
+    {
+        /* needsDistributedSpatiotemporalPlanning() can return true purely
+         * from tablesList->diffCount > 1, even when checkQueryType() found
+         * no registered intersection/distance predicate to build a strategy
+         * for (e.g. a plain `@>` "contains" clause against a reference
+         * table isn't one of those). With no strategy and no distributed
+         * function, there is nothing for the custom executor below to
+         * build a plan from -- it would otherwise fall into
+         * ConstructGeneralQuery's `generalScan->length == 0` branch and use
+         * postProcessing->worker, which is NULL here, producing "(null)"
+         * as the query text. Let Citus plan the query directly instead. */
+        return distributed_planner(parse, query_string, cursorOptions, boundParams);
+    }
     if (needsSpatiotemporalPlanning)
     {
         if (!distPlan->activate_rewriter)
@@ -198,6 +213,13 @@ analyzeDistributedSpatiotemporalTables(List *rangeTableList,
 {
     ListCell *rangeTableCell = NULL;
     Oid curr_relid = -1;
+    /* diffCount/simCount need to know whether relid has appeared ANYWHERE
+     * earlier in the range table, not just in the immediately preceding
+     * entry -- comparing only to curr_relid miscounted a self-join like
+     * "Trips t1, Licences1 l1, Trips t2" as three different tables instead
+     * of recognizing t2 as a repeat of t1, since t2 is compared against
+     * l1's relid rather than t1's. */
+    List *seenRelids = NIL;
     bool shapeType;
     List *rtes = NIL;
     foreach(rangeTableCell, rangeTableList)
@@ -206,6 +228,21 @@ analyzeDistributedSpatiotemporalTables(List *rangeTableList,
         if (rangeTableEntry->rtekind != RTE_RELATION) {
             continue;
         }
+        /* A view reference (e.g. Licences1, a view over Licences) expands
+         * in the range table into the view's own subquery entry (rtekind
+         * != RTE_RELATION, already skipped above), PostgreSQL's internal
+         * rule-system "old"/"new" placeholder entries for that view, and
+         * the view's real underlying base table -- none of which the user
+         * actually joined except the last. inFromCl is false for exactly
+         * those non-user-visible placeholders (verified: "old"/"new" have
+         * inFromCl=0, the real base table has inFromCl=1), so without this
+         * check every view reference was triple-counted as extra distinct
+         * tables, inflating diffCount/length enough to make e.g. a query
+         * with one distributed table plus reference tables look like it
+         * had several more distributed tables than it really did. */
+        if (!rangeTableEntry->inFromCl) {
+            continue;
+        }
         if (IsDistributedSpatiotemporalTable(rangeTableEntry->relid))
         {
             if(IsReshuffledTable(rangeTableEntry->relid))
@@ -248,15 +285,29 @@ analyzeDistributedSpatiotemporalTables(List *rangeTableList,
             if (LookupCitusTableCacheEntry(rangeTableEntry->relid) != NULL)
             {
                 char partitioningMethod = PartitionMethodViaCatalog (rangeTableEntry->relid);
-                if (partitioningMethod == DISTRIBUTE_BY_HASH || partitioningMethod == DISTRIBUTE_BY_RANGE)
+                /* A reference table is replicated to every node, so joining
+                 * it alongside a distributed table needs no repartitioning.
+                 * Its reported partitioning method alone doesn't identify it
+                 * uniquely, so check its table type explicitly instead. */
+                bool isReferenceTable = IsCitusTableType(rangeTableEntry->relid, REFERENCE_TABLE);
+                if (partitioningMethod == DISTRIBUTE_BY_HASH || partitioningMethod == DISTRIBUTE_BY_RANGE
+                        || isReferenceTable)
                 {
                     /* Citus table processing */
                     CitusRteNode *citusNode = GetCitusRteInfo(rangeTableEntry,partitioningMethod);
                     citusNode->rangeTableCell = rangeTableCell;
-                    distPlan->tablesList->length++;
+                    /* length is already incremented unconditionally for
+                     * every range table entry below (after this if/else) --
+                     * incrementing it here too double-counted every Citus
+                     * table entry, inflating effectiveLength enough that a
+                     * query with one distributed table plus reference
+                     * tables was never recognized as an effective
+                     * single-table case. */
                     Rte *rteNode = GetRteNode((Node *) citusNode, CitusRte, rangeTableEntry->alias);
                     rtes = lappend(rtes , rteNode);
                     distPlan->tablesList->nonStCount++;
+                    if (isReferenceTable)
+                        distPlan->tablesList->refCount++;
                 }
                 else
                     elog(ERROR, "The %s table is not distributed using one of the supported partitioning methods",
@@ -270,9 +321,10 @@ analyzeDistributedSpatiotemporalTables(List *rangeTableList,
                 rtes = lappend(rtes , rteNode);
             }
         }
-        if (curr_relid != rangeTableEntry->relid)
+        if (!list_member_oid(seenRelids, rangeTableEntry->relid))
         {
             distPlan->tablesList->diffCount++;
+            seenRelids = lappend_oid(seenRelids, rangeTableEntry->relid);
         }
         else
             distPlan->tablesList->simCount++;
@@ -355,6 +407,14 @@ checkQueryType(Query *parse, DistributedSpatiotemporalQueryPlan *distPlan)
         /* TODO: subquery is excluded for now */
         ereport(ERROR, (errmsg("A sub query is not supported yet in Distributed MobilityDB!")));
     }
+    /* Reference tables are already replicated to every node, so a join
+     * against one never needs the NonColocation strategy's reshuffle --
+     * Citus can push the predicate down to each shard directly. Subtracting
+     * refCount here means a query joining one distributed spatiotemporal
+     * table with any number of reference tables is treated the same as a
+     * genuine single-table query below. */
+    int effectiveDiffCount = distPlan->tablesList->diffCount - distPlan->tablesList->refCount;
+    int effectiveLength = distPlan->tablesList->length - distPlan->tablesList->refCount;
     /* Iterate over the where clause conditions */
     foreach(clauseCell, whereClauseList)
     {
@@ -379,12 +439,12 @@ checkQueryType(Query *parse, DistributedSpatiotemporalQueryPlan *distPlan)
             {
                 if (IsIntersectionOperation(predicateOid))
                 {
-                    if (distPlan->tablesList->diffCount > 1)
+                    if (effectiveDiffCount > 1)
                     {
                         /* Intersection join between two distinct tables: must colocate them first. */
                         AddStrategy(distPlan, NonColocation);
                     }
-                    else if (distPlan->tablesList->length == 1)
+                    else if (effectiveLength == 1)
                     {
                         /* Single-table intersection: decide between rebalancing tiles to fit the
                          * query's search box or simply pushing the predicate to each worker. */
@@ -406,10 +466,22 @@ checkQueryType(Query *parse, DistributedSpatiotemporalQueryPlan *distPlan)
                 }
                 else if (IsDistanceOperation(predicateOid))
                 {
-                    /* The NonColocation strategy is triggered by default until the analysis changes it */
                     if (distPlan->tablesList->simCount >= 1)
                         AddStrategy(distPlan, Colocation);
-                    AddStrategy(distPlan, NonColocation);
+                    /* The NonColocation strategy is triggered by default until the analysis
+                     * changes it -- except when the only "different" tables besides one
+                     * distributed spatiotemporal table are reference tables (refCount > 0
+                     * guards this so behavior is untouched whenever no reference table is
+                     * involved), which Citus can push the predicate down to directly with
+                     * no reshuffle needed. */
+                    if (effectiveDiffCount > 1 || distPlan->tablesList->refCount == 0)
+                    {
+                        AddStrategy(distPlan, NonColocation);
+                    }
+                    else if (effectiveLength == 1)
+                    {
+                        AddStrategy(distPlan, PredicatePushDown);
+                    }
                     if(distPlan->predicatesList->predicateType == DISTANCE)
                     {
                         ereport(ERROR, (errmsg("Currently, we do not support using more than "

src/planner/planner_strategies.c

@@ -451,18 +451,40 @@ getReshuffledColumns(DistributedSpatiotemporalQueryPlan *distPlan, Oid oid)
 }
 
 /*
- * ColocationStrategyPlan records a Colocation PlanTask for the query's
- * first two range-table entries, joined on their shared tile key — no data
- * movement needed since both are already tiled the same way.
+ * ColocationStrategyPlan records a Colocation PlanTask for the query's two
+ * self-joined spatiotemporal (STRte) range-table entries, joined on their
+ * shared tile key — no data movement needed since both are already tiled
+ * the same way.
  */
 extern void
 ColocationStrategyPlan(DistributedSpatiotemporalQueryPlan *distPlan)
 {
     PlanTask * strategy = (PlanTask *) palloc0(sizeof(PlanTask));
     strategy->type = Colocation;
     /* TODO: Add the other cases */
-    strategy->tbl1 = (STMultirelation *) ((Rte *)list_nth(distPlan->tablesList->tables, 0))->rte;
-    strategy->tbl2 = (STMultirelation *) ((Rte *)list_nth(distPlan->tablesList->tables, 1))->rte;
+    /* This is a self-join, so the two tables to join are the two STRte
+     * entries specifically -- not just whichever entries happen to be
+     * first/second in tablesList->tables. Any reference or plain Citus
+     * tables also joined in the same query (CitusRte/LocalRte entries)
+     * can end up interleaved with them (e.g. "Trips t1, Licences1 l1,
+     * Trips t2" puts l1 at index 1), and blindly casting one of those to
+     * STMultirelation* read garbage through the wrong struct layout. */
+    ListCell *rangeTableCell = NULL;
+    STMultirelation *stTables[2] = {NULL, NULL};
+    int stTableCount = 0;
+    foreach(rangeTableCell, distPlan->tablesList->tables)
+    {
+        Rte *rteNode = (Rte *) lfirst(rangeTableCell);
+        if (rteNode->RteType == STRte && stTableCount < 2)
+        {
+            stTables[stTableCount] = (STMultirelation *) rteNode->rte;
+            stTableCount++;
+        }
+    }
+    if (stTableCount < 2)
+        ereport(ERROR, (errmsg("Colocation strategy requires two spatiotemporal tables to self-join")));
+    strategy->tbl1 = stTables[0];
+    strategy->tbl2 = stTables[1];
     strategy->tileKey = (Datum) Var_Catalog_Tile_Key;
     distPlan->strategyPlans = lappend(distPlan->strategyPlans, strategy);
 }