Merge pull request #25 from mbakli/worker_level_optimization

Worker level optimization

Author: mbakli

include/catalog/nodes.h

@@ -29,6 +29,9 @@ 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);
 
+/* Physical shard name for a Citus reference table's single (every-node-replicated) shard. */
+extern char* GetReferenceTableShardName(Oid relationId);
+
 /* Looks up the tiling method used to distribute relationId; -1 if not distributed. */
 extern int TilingSearch(Oid relationId);
 

include/distributed_functions/distributed_function.h

@@ -61,6 +61,12 @@ extern DistributedFunction *addDistributedFunction(TargetEntry *operation);
 /* True if targetEntry's expression is a registered distributable function. */
 extern bool IsDistFunc(TargetEntry *targetEntry);
 
+/* Looks up workerFuncName's registered "final" combining op (e.g. "sum" for "length"), or NULL if unregistered. */
+extern char *LookupDistFuncFinalOp(const char *workerFuncName);
+
+/* Looks up workerFuncName's registered "combiner" op, or NULL if it has none (the common case today). */
+extern char *LookupDistFuncCombinerOp(const char *workerFuncName);
+
 /* Builds a QOperation pairing a distributed op (des) with its argument column (cur). */
 extern QOperation * AddQOperation(Datum des, Datum cur);
 #endif /* DISTRIBUTED_FUNCTION_H */

include/planner/distributed_mobilitydb_planner.h

@@ -62,6 +62,27 @@ typedef struct DistributedSpatiotemporalQueryPlan
     char *org_query_string;
     Datum range_bbox;
     PostProcessing *postProcessing;
+    /*
+     * Set when RewriteSegmentedDistFuncCalls rewrote the query (bare
+     * distributed-function call over a segmented table -> explicit grouped
+     * aggregate) and handed off to Citus' own planner directly. The
+     * EXPLAIN hook treats a non-NULL PlannedStmt from
+     * distributed_mobilitydb_planner_internal as "our custom planning
+     * bailed out, re-explain the original query" -- which is wrong here,
+     * since this *is* the plan that actually runs; without this, EXPLAIN
+     * would show the original bare (ungrouped) query shape instead of what
+     * was really executed.
+     */
+    char *segmentedRewriteQuery;
+    /*
+     * Set alongside segmentedRewriteQuery: one line per rewritten function
+     * naming its registered worker/combiner/final ops (from
+     * pg_dist_spatiotemporal_dist_functions) and the op actually applied in
+     * the rewrite -- for EXPLAIN to show *why* the rewrite looks the way it
+     * does in this extension's own worker/combiner/final vocabulary,
+     * without asserting anything about replication vs. true segmentation.
+     */
+    char *segmentedRewriteExplainNotes;
 } DistributedSpatiotemporalQueryPlan;
 
 /* Filter Operation */

include/planner/query_semantics.h

@@ -20,6 +20,35 @@
 /* Scans the SELECT targetlist for distributed aggregates and records them into postProcessing. */
 extern void analyseSelectClause(List *targetList, PostProcessing *postProcessing);
 
+/*
+ * Rewrites a bare distributed-function call over a segmented table into a
+ * grouped-by-trip aggregate query; NULL if nothing to rewrite. On success,
+ * *explainNotesOut is set to a human-readable, newline-joined description
+ * of each rewritten function's worker/combiner/final ops and the op
+ * actually applied (for EXPLAIN); left untouched on a NULL return.
+ */
+extern char *RewriteSegmentedDistFuncCalls(Query *parse, const char *query_string, STMultirelations *tablesList,
+                                           char **explainNotesOut);
+
+/*
+ * Rewrites a query with an explicit aggregate over a registered distributed
+ * function applied to a replicated (isMobilityDB, segmented) table's column
+ * -- e.g. `SUM(length(atTime(t.Trip, p.Period))) ... GROUP BY ...` -- into a
+ * two-level dedupe/aggregate query, so a trip replicated across N tiles
+ * contributes to the aggregate once instead of N times. NULL if nothing to
+ * rewrite.
+ */
+extern char *RewriteReplicatedAggregateQuery(Query *parse, const char *query_string, STMultirelations *tablesList);
+
+/*
+ * Same problem as RewriteReplicatedAggregateQuery, but for a query whose
+ * aggregate-over-distributed-function lives inside one of its own CTEs
+ * (e.g. `WITH x AS (SELECT ... SUM(length(atTime(...))) ... GROUP BY ...)
+ * SELECT ... FROM x`) rather than at the top level; NULL if no CTE needed
+ * rewriting (or there are no CTEs at all).
+ */
+extern char *RewriteReplicatedAggregateInCTEs(Query *parse, const char *query_string, STMultirelations *tablesList);
+
 /* Analyses fromExpr's predicates against tbl's catalog to derive its candidate-tile filter. */
 extern CatalogFilter *AnalyseCatalog(STMultirelation *tbl, FromExpr * fromExpr);
 

include/utils/helper_functions.h

@@ -36,6 +36,22 @@ extern char *change_sentence (char *sentence, char *find, char *replace);
 /* Returns a lowercased copy of str. */
 extern char *toLower(char *str);
 
+/*
+ * Finds the first case-insensitive, whitespace-bounded occurrence of
+ * `keyword` in `lowered` (already-lowercased haystack, already-lowercase
+ * keyword) -- tolerates newlines/tabs/multiple spaces around it, unlike a
+ * plain strstr(haystack, " keyword "). Returns a pointer to the start of
+ * the keyword itself, or NULL if not found.
+ */
+extern char *FindKeywordToken(const char *lowered, const char *keyword);
+
+/*
+ * Like FindKeywordToken, but only matches an occurrence at paren-depth 0
+ * (not nested inside a subquery/CTE's own parenthesized definition) -- for
+ * locating a keyword that belongs to a whole query's outermost SELECT.
+ */
+extern char *FindTopLevelKeywordToken(const char *lowered, const char *keyword);
+
 /* True if val is a NULL/zero Datum (no value set). */
 extern bool IsDatumEmpty(Datum val);
 

include/utils/planner_utils.h

@@ -29,7 +29,7 @@
 
 
 /* constants for tiles.options */
-#define Natts_MTS 14
+#define Natts_MTS 15
 #define Anum_MTS_oid 1
 #define Anum_MTS_numTiles 3
 #define Anum_MTS_method 4
@@ -42,7 +42,14 @@
 #define Anum_MTS_tileKey 11
 #define Anum_MTS_segmentation 12
 #define Anum_MTS_srid 13
-#define Anum_MTS_groupCol 5
+/* groupcol was appended as the table's 15th (0-indexed 14th) column --
+ * it did not exist when this catalog table was first designed, so unlike
+ * the constants above (which match physical column order), this one can't
+ * be slotted in without renumbering every constant after it. Previously
+ * defined as 5, colliding with Anum_MTS_type -- vestigial from a groupCol
+ * column that was never actually added to the table, so this was always
+ * dead/wrong (GetTilingSchemeInfo never read it). */
+#define Anum_MTS_groupCol 14
 
 /* MobilityDB and PostGIS variables */
 #define Var_MobilityDB_BBOX "mobdb_bbox"

sql/distributed_mobilitydb--1.0.sql

@@ -61,7 +61,8 @@ CREATE TABLE dist_mobilitydb.pg_dist_spatiotemporal_tables(
     distcoltype varchar(10),
     tilekey varchar(10),
     shapeSegmented boolean,
-    srid int
+    srid int,
+    groupcol varchar(50)
 );
 
 ALTER TABLE dist_mobilitydb.pg_dist_spatiotemporal_tables
@@ -91,14 +92,36 @@ CREATE TABLE dist_mobilitydb.pg_spatiotemporal_join_operations(
 );
 
 -- Add the OID for the distance and intersection query operations
+--
+-- GetPredicateOidAndArgs (src/planner/predicate_management.c) reads an
+-- OpExpr predicate's *operator* oid (opExpr->opno), not the oid of the
+-- function implementing it -- registering only pg_proc.oid values (as the
+-- two SELECTs below do) means an operator-written predicate like
+-- `t.Trip && p.Period` (the `&&` bbox/temporal-overlap operator,
+-- implemented by temporal_overlaps/span_overlaps/etc., but a *different*
+-- oid than those functions') never matched, even after the function
+-- itself was registered. For queries that use `&&` as their *only*
+-- spatiotemporal predicate (no accompanying eintersects/ST_Intersects/
+-- eDwithin, e.g. BerlinMOD Q8's "was this vehicle active during this
+-- period" check), that meant this extension's planner never engaged at
+-- all -- no tile pruning, none of its tile-boundary deduplication --
+-- reproduced as genuinely wrong (duplicated) results. The third SELECT
+-- below registers the `&&` *operator*'s own oid (from pg_operator) for
+-- every temporal/spatiotemporal type combination it's defined over.
 INSERT INTO dist_mobilitydb.pg_spatiotemporal_join_operations(op, opid, distance)
 SELECT proname,oid,true
 FROM pg_proc
 WHERE proname like ANY(ARRAY['%dwithin%', '%distance%'])
 union all
 SELECT proname,oid,false
 FROM pg_proc
-WHERE proname like ANY(ARRAY['%intersects%', '%contains%', '%disjoint%']);
+WHERE proname like ANY(ARRAY['%intersects%', '%contains%', '%disjoint%', '%overlaps%'])
+union all
+SELECT '&&', o.oid, false
+FROM pg_operator o
+WHERE o.oprname = '&&'
+  AND (o.oprleft::regtype::text ~ 'tgeompoint|tgeogpoint|tbool|tint|tfloat|ttext|tnpoint|tstz'
+       OR o.oprright::regtype::text ~ 'tgeompoint|tgeogpoint|tbool|tint|tfloat|ttext|tnpoint|tstz');
 
 ALTER TABLE dist_mobilitydb.pg_spatiotemporal_join_operations
 SET SCHEMA pg_catalog;

sql/helper_functions/write_to_catalog.sql

@@ -12,8 +12,8 @@ BEGIN
         column_type := 'geometry';
     END IF;
     EXECUTE format('%s', concat('' ||
-                            'INSERT INTO pg_dist_spatiotemporal_tables (tblOid, tableName, numTiles, tilingMethod, tilingType, granularity, disjoint, isMobilityDB, distcol, distcoltype, tilekey, shapeSegmented, srid) ' ||
-                            'VALUES (',0,',''',table_name_out,''',',tiling.numTiles,',''',tiling.method,''',''',tiling.type,''',''',tiling.granularity,''',''',tiling.disjointTiles,''',''',tiling.isMobilityDB,''', ''',tiling.distCol,''', ''',tiling.distColType,''', ''',tiling.tileKey,''', ''',tiling.segmentation,''', ',tiling.srid,') ' ||
+                            'INSERT INTO pg_dist_spatiotemporal_tables (tblOid, tableName, numTiles, tilingMethod, tilingType, granularity, disjoint, isMobilityDB, distcol, distcoltype, tilekey, shapeSegmented, srid, groupcol) ' ||
+                            'VALUES (',0,',''',table_name_out,''',',tiling.numTiles,',''',tiling.method,''',''',tiling.type,''',''',tiling.granularity,''',''',tiling.disjointTiles,''',''',tiling.isMobilityDB,''', ''',tiling.distCol,''', ''',tiling.distColType,''', ''',tiling.tileKey,''', ''',tiling.segmentation,''', ',tiling.srid,', ''',tiling.groupCol,''') ' ||
                                                                                                                                                                                                            'ON CONFLICT (tableName) ' ||
                                                                                                                                                                                                            'DO ' ||
                                                                                                                                                                                                            'UPDATE set tableName = EXCLUDED.tableName,' ||
@@ -27,7 +27,8 @@ BEGIN
                                                                                                                                                                                                            'distcoltype = EXCLUDED.distcoltype,' ||
                                                                                                                                                                                                            'tilekey = EXCLUDED.tilekey,' ||
                                                                                                                                                                                                            'shapeSegmented = EXCLUDED.shapeSegmented,' ||
-                                                                                                                                                                                                            'srid = EXCLUDED.srid'  ));
+                                                                                                                                                                                                            'srid = EXCLUDED.srid,' ||
+                                                                                                                                                                                                            'groupcol = EXCLUDED.groupcol'  ));
 
     EXECUTE format('%s', concat('SELECT id FROM pg_dist_spatiotemporal_tables WHERE tableName = ''',table_name_out,''''))
         INTO table_out_id;

sql/partitioning/CMakeLists.txt

@@ -3,6 +3,7 @@ SET(LOCAL_FILES
         shape_segmentation.sql
         search_in_dimensions.sql
         crange.sql
+        hierarchical.sql
         colocation.sql
         tiling.sql
         data_allocation.sql

sql/partitioning/data_allocation.sql

@@ -25,6 +25,21 @@ BEGIN
     IF not tiling.isMobilityDB and tiling.internaltype = 'linestring' THEN
         EXECUTE format('%s', concat('ALTER TABLE ', table_name_out,' ALTER COLUMN ', tiling.distCol,' TYPE geometry;'));
     END IF;
+    /*
+     * The distributed column holds large TOASTed values (whole/clipped
+     * trajectories for MobilityDB, geometries for PostGIS) -- switching its
+     * TOAST compression from the default pglz to lz4 (much faster to
+     * compress, at a modest space cost) cuts the segmentation/allocation
+     * INSERT's dominant cost, which is writing these rows, not scanning
+     * them (confirmed via EXPLAIN ANALYZE: the write phase alone accounts
+     * for the majority of this step's time, and that write can't be sped
+     * up by parallel workers -- Postgres disables parallel query entirely
+     * for any statement containing a write, regardless of GUCs). Measured
+     * ~2.75x faster (58.8s -> 21.4s) on a same-data before/after comparison
+     * of this exact INSERT. No effect on query results, only on-disk
+     * compression of this one column.
+     */
+    EXECUTE format('%s', concat('ALTER TABLE ', table_name_out,' ALTER COLUMN ', tiling.distCol,' SET COMPRESSION lz4;'));
     -- Add the distributed column
     EXECUTE format('%s', concat('ALTER TABLE ',table_name_out, ' ADD column ',tiling.tileKey,' integer'));
     -- Distribute the table using range multirelation
@@ -114,6 +129,24 @@ BEGIN
             SELECT ',org_table_columns, ',',tiling.tileKey,'
             FROM ',table_name_in,' t1, pg_dist_spatiotemporal_tiles
             WHERE table_id=',table_id,' and ', tiling.distCol,' && ',bbox_with_srid));
+    ELSIF tiling.internaltype in ('sequence','sequenceset') THEN
+        /*
+         * Mirrors the linestring/polygon branch above: replicate each trip
+         * whole (unclipped) into every tile it overlaps, rather than
+         * segmenting/clipping it (see segmentation_and_allocation for the
+         * clipping counterpart, selected instead whenever
+         * tiling.segmentation is true). Previously missing entirely --
+         * shape_segmentation => false (the "replicate" choice) fell into
+         * the catch-all ELSE below and errored out for any MobilityDB
+         * sequence/sequenceset table (e.g. BerlinMOD trips), even though
+         * the segmenting path already worked.
+         */
+        RAISE INFO 'Distributing the %s into the overlapping tiles without segmenting (i.e., replication) them:',tiling.internaltype;
+        EXECUTE format('%s', concat('
+            INSERT INTO ',table_name_out,'
+            SELECT ',org_table_columns, ',',tiling.tileKey,'
+            FROM ',table_name_in,' t1, pg_dist_spatiotemporal_tiles
+            WHERE table_id=',table_id,' and ', tiling.distCol,' && ',bbox_with_srid));
     ELSE
         RAISE Exception 'The column type is not detected!';
     END IF;