Add more tests and prevent push down for where, group by, over, etc.

abramsh · abramsh · commit 2f27ecd46dd4 · 2025-09-05T13:02:32.000-07:00
diff --git a/src/multicorn.c b/src/multicorn.c
@@ -125,6 +125,10 @@ static void multicorn_xact_callback(XactEvent event, void *arg);
 void	   *serializePlanState(MulticornPlanState * planstate);
 MulticornExecState *initializeExecState(void *internal_plan_state);
 
+static void add_foreign_ordered_paths(PlannerInfo *root,
+                                    RelOptInfo *input_rel,
+                                    RelOptInfo *final_rel,
+                                    FinalPathExtraData *extra);
 static void add_foreign_final_paths(PlannerInfo *root,
                                     RelOptInfo *input_rel,
                                     RelOptInfo *final_rel,
@@ -454,8 +458,20 @@ static void multicornGetForeignUpperPaths(PlannerInfo *root,
                                           RelOptInfo *output_rel, 
                                           void *extra)
 {
+    // elog(WARNING, "Got input_rel private: %p", input_rel->fdw_private);
+    // elog(WARNING, "Got output_rel private: %p", output_rel->fdw_private);
+
+    // If the input_rel has no private, then pushdown wasn't supported for the previous stage
+    // Which means we can't pushdown anything for the the current stage (as least this is true for limit/offset)
+    if (!input_rel->fdw_private)
+        return;
+
+    // elog(WARNING, "Got stage: %d", stage);
     switch (stage)
 	{
+        case UPPERREL_ORDERED:
+            add_foreign_ordered_paths(root, input_rel, output_rel, (FinalPathExtraData *)extra);
+            break;
 		case UPPERREL_FINAL:
             add_foreign_final_paths(root, input_rel, output_rel, (FinalPathExtraData *)extra);
 			break;
@@ -464,6 +480,60 @@ static void multicornGetForeignUpperPaths(PlannerInfo *root,
 	}
 }
 
+/*
+ * add_foreign_ordered_paths
+ *		Add foreign paths for performing the sort processing remotely.
+ *
+ * Given input_rel contains the source-data Paths.  The paths are added to the
+ * given final_rel.
+ *
+ * Note: Since sorts are already taken care of in the base rel, we only check for pushdown here.
+ */
+ static void
+ add_foreign_ordered_paths(PlannerInfo *root, RelOptInfo *input_rel,
+                         RelOptInfo *final_rel,
+                         FinalPathExtraData *extra)
+ {
+    List *applied_pathkeys = NIL;
+    ListCell *lc;
+    MulticornPlanState *planstate = input_rel->fdw_private;
+    ForeignPath *cheapest_path = (ForeignPath *)input_rel->cheapest_total_path;
+    if ( planstate )
+    {
+        if (cheapest_path && IsA(cheapest_path, ForeignPath))
+        {
+            MulticornPathState *pathstate = (MulticornPathState *)cheapest_path->fdw_private;
+            if ( pathstate )
+            {
+                planstate->pathkeys = pathstate->pathkeys;
+            }
+        }
+
+        /* Extract the pathkeys from the input_rel */
+        foreach(lc, input_rel->pathlist)
+        {
+            Path *path = (Path *) lfirst(lc);
+            if (IsA(path, ForeignPath))
+            {
+                ForeignPath *fpath = (ForeignPath *) path;
+                if (fpath->path.pathkeys != NIL)
+                {
+                    applied_pathkeys = fpath->path.pathkeys;
+                    break;
+                }
+            }
+        }
+
+        /* We only support limit/offset if the sort is completely pushed down */
+        /* By bailing here, input_rel for the next state will not have planstate, which will cause no more pushdowns */
+        if (!pathkeys_contained_in(root->sort_pathkeys, applied_pathkeys))
+            return;
+
+        planstate->input_rel = input_rel;
+        final_rel->fdw_private = planstate;
+    }
+}
+
 /*
  * add_foreign_final_paths
  *		Add foreign paths for performing the final processing remotely.
@@ -482,10 +552,6 @@ static void multicornGetForeignUpperPaths(PlannerInfo *root,
     ForeignPath *final_path;
     int limitCount = -1;
     int limitOffset = -1;
-    Path *cheapest_path;
-    List *deparsed_pathkeys = NIL;
-    List *applied_pathkeys = NIL;
-    ListCell *lc;
 
     /* No work if there is no need to add a LIMIT node */
     if (!extra->limit_needed)
@@ -499,6 +565,10 @@ static void multicornGetForeignUpperPaths(PlannerInfo *root,
     if (parse->limitOption == LIMIT_OPTION_WITH_TIES)
         return;
 
+    /* We don't currently support pushing down limits with quals */
+    if (parse->jointree->quals)
+        return;
+
     /* only push down constant LIMITs... */
     if ((parse->limitCount && !IsA(parse->limitCount, Const)) || (parse->limitOffset && !IsA(parse->limitOffset, Const)))
         return;
@@ -516,57 +586,16 @@ static void multicornGetForeignUpperPaths(PlannerInfo *root,
 
     /* Get the current input_rel and it's planstate */
     planstate = input_rel->fdw_private;
-    // TODO: Maybe this isn't needed if we handle the previous stages correctly?
-    if (!planstate)
-    {
-        for (int i = 1; i < root->simple_rel_array_size; i++)
-        {
-            RelOptInfo *rel = root->simple_rel_array[i];
-            if (rel && rel->reloptkind == RELOPT_BASEREL)
-            {
-                planstate = rel->fdw_private;
-                input_rel->fdw_private = planstate;
-                input_rel = rel;
-                break;
-            }
-        }
-    }
-    
-    /* Extract pathkeys from the cheapest path's fdw_private if it exists */
-    cheapest_path = input_rel->cheapest_total_path;
-    if (cheapest_path && IsA(cheapest_path, ForeignPath))
-    {
-        ForeignPath *foreign_path = (ForeignPath *)cheapest_path;
-        if (foreign_path->fdw_private)
-        {
-            MulticornPathState *input_pathstate = (MulticornPathState *)foreign_path->fdw_private;
-            deparsed_pathkeys = input_pathstate->pathkeys;
-        }
-    }
-    
-    /* Extract the pathkeys from the input_rel */
-    foreach(lc, input_rel->pathlist)
-    {
-        Path *path = (Path *) lfirst(lc);
-        if (IsA(path, ForeignPath))
-        {
-            ForeignPath *fpath = (ForeignPath *) path;
-            if (fpath->path.pathkeys != NIL)
-                applied_pathkeys = fpath->path.pathkeys;
-        }
-    }
- 
-    /* We only support limit/offset if the sort is completely pushed down */
-    if (!pathkeys_contained_in(root->sort_pathkeys, applied_pathkeys))
-        return;
+    if ( planstate->input_rel )
+        input_rel = planstate->input_rel;
    
     /* Check if Python FWD can push down the LIMIT/OFFSET */
     if (!canLimit(planstate, limitCount, limitOffset))
         return;
 
     /* Create foreign final path with the correct number of rows, and include state for limit/offset pushdown */
     pathstate = (MulticornPathState *)calloc(1, sizeof(MulticornPathState));
-    pathstate->pathkeys = deparsed_pathkeys;
+    pathstate->pathkeys = planstate->pathkeys;
     pathstate->limit = limitCount;
     pathstate->offset = limitOffset;
     final_path = create_foreign_upper_path(root,
diff --git a/src/multicorn.h b/src/multicorn.h
@@ -93,6 +93,9 @@ typedef struct MulticornPlanState
     int offset;
     int limit;
 
+    /* Used for tracking the input_rel for upper plans*/
+    RelOptInfo *input_rel;
+
 }	MulticornPlanState;
 
 typedef struct MulticornExecState
diff --git a/test-3.9/expected/multicorn_test_limit.out b/test-3.9/expected/multicorn_test_limit.out
@@ -121,6 +121,85 @@ NOTICE:  ['test1', 'test2']
  01-01-2011 | Sun Jan 02 14:30:25 2011
 (1 row)
 
+-- Verify limit and offset are not pushed down with where (which we may eventually support in the future)
+EXPLAIN SELECT * FROM testmulticorn WHERE test1 = '02-03-2011' LIMIT 1 OFFSET 1;
+                                 QUERY PLAN                                 
+----------------------------------------------------------------------------
+ Limit  (cost=29.50..49.00 rows=1 width=20)
+   ->  Foreign Scan on testmulticorn  (cost=10.00..400.00 rows=20 width=20)
+         Filter: (test1 = '02-03-2011'::date)
+(3 rows)
+
+SELECT * FROM testmulticorn WHERE test1 = '02-03-2011' LIMIT 1 OFFSET 1;
+NOTICE:  [test1 = 2011-02-03]
+NOTICE:  ['test1', 'test2']
+   test1    |          test2           
+------------+--------------------------
+ 02-03-2011 | Tue Feb 01 14:30:25 2011
+(1 row)
+
+-- Verify limit and offset are not pushed down with group by
+EXPLAIN SELECT max(test1) FROM testmulticorn GROUP BY test1 LIMIT 1 OFFSET 1;
+                                    QUERY PLAN                                    
+----------------------------------------------------------------------------------
+ Limit  (cost=19.52..29.03 rows=1 width=8)
+   ->  GroupAggregate  (cost=10.00..200.30 rows=20 width=8)
+         Group Key: test1
+         ->  Foreign Scan on testmulticorn  (cost=10.00..200.00 rows=20 width=10)
+(4 rows)
+
+SELECT max(test1) FROM testmulticorn GROUP BY test1 LIMIT 1 OFFSET 1;
+NOTICE:  []
+NOTICE:  ['test1']
+NOTICE:  requested sort(s): 
+NOTICE:  SortKey(attname='test1', attnum=1, is_reversed=False, nulls_first=False, collate=None)
+    max     
+------------
+ 02-03-2011
+(1 row)
+
+-- Verify limit and offset are not pushed down with window functions
+EXPLAIN SELECT max(test1) OVER (ORDER BY test1) FROM testmulticorn LIMIT 1 OFFSET 1;
+                                    QUERY PLAN                                    
+----------------------------------------------------------------------------------
+ Limit  (cost=28.55..37.59 rows=1 width=8)
+   ->  WindowAgg  (cost=19.52..200.30 rows=20 width=8)
+         ->  Foreign Scan on testmulticorn  (cost=10.00..200.00 rows=20 width=10)
+(3 rows)
+
+SELECT max(test1) OVER (ORDER BY test1) FROM testmulticorn LIMIT 1 OFFSET 1;
+NOTICE:  []
+NOTICE:  ['test1']
+NOTICE:  requested sort(s): 
+NOTICE:  SortKey(attname='test1', attnum=1, is_reversed=False, nulls_first=False, collate=None)
+    max     
+------------
+ 01-01-2011
+(1 row)
+
+-- Verify limit and offset are not pushed down with joins 
+-- TODO: optimize by pushing down limit/offset on one side of the join
+EXPLAIN SELECT * FROM testmulticorn m1 JOIN testmulticorn m2 ON m1.test1 = m2.test1 LIMIT 1 OFFSET 1;
+                                        QUERY PLAN                                         
+-------------------------------------------------------------------------------------------
+ Limit  (cost=59.30..98.61 rows=1 width=24)
+   ->  Nested Loop  (cost=20.00..806.05 rows=20 width=24)
+         Join Filter: (m1.test1 = m2.test1)
+         ->  Foreign Scan on testmulticorn m1  (cost=10.00..400.00 rows=20 width=20)
+         ->  Materialize  (cost=10.00..400.10 rows=20 width=20)
+               ->  Foreign Scan on testmulticorn m2  (cost=10.00..400.00 rows=20 width=20)
+(6 rows)
+
+SELECT * FROM testmulticorn m1 JOIN testmulticorn m2 ON m1.test1 = m2.test1 LIMIT 1 OFFSET 1;
+NOTICE:  []
+NOTICE:  ['test1', 'test2']
+NOTICE:  []
+NOTICE:  ['test1', 'test2']
+   test1    |          test2           |   test1    |          test2           
+------------+--------------------------+------------+--------------------------
+ 01-01-2011 | Sun Jan 02 14:30:25 2011 | 01-01-2011 | Sun Jan 02 14:30:25 2011
+(1 row)
+
 DROP USER MAPPING FOR current_user SERVER multicorn_srv;
 DROP EXTENSION multicorn cascade;
 NOTICE:  drop cascades to 4 other objects
diff --git a/test-3.9/sql/multicorn_test_limit.sql b/test-3.9/sql/multicorn_test_limit.sql
@@ -55,5 +55,22 @@ SELECT * FROM testmulticorn ORDER BY test1 LIMIT 1 OFFSET 1;
 EXPLAIN SELECT * FROM testmulticorn_nosort ORDER BY test1 LIMIT 1 OFFSET 1;
 SELECT * FROM testmulticorn_nosort ORDER BY test1 LIMIT 1 OFFSET 1;
 
+-- Verify limit and offset are not pushed down with where (which we may eventually support in the future)
+EXPLAIN SELECT * FROM testmulticorn WHERE test1 = '02-03-2011' LIMIT 1 OFFSET 1;
+SELECT * FROM testmulticorn WHERE test1 = '02-03-2011' LIMIT 1 OFFSET 1;
+
+-- Verify limit and offset are not pushed down with group by
+EXPLAIN SELECT max(test1) FROM testmulticorn GROUP BY test1 LIMIT 1 OFFSET 1;
+SELECT max(test1) FROM testmulticorn GROUP BY test1 LIMIT 1 OFFSET 1;
+
+-- Verify limit and offset are not pushed down with window functions
+EXPLAIN SELECT max(test1) OVER (ORDER BY test1) FROM testmulticorn LIMIT 1 OFFSET 1;
+SELECT max(test1) OVER (ORDER BY test1) FROM testmulticorn LIMIT 1 OFFSET 1;
+
+-- Verify limit and offset are not pushed down with joins 
+-- TODO: optimize by pushing down limit/offset on one side of the join
+EXPLAIN SELECT * FROM testmulticorn m1 JOIN testmulticorn m2 ON m1.test1 = m2.test1 LIMIT 1 OFFSET 1;
+SELECT * FROM testmulticorn m1 JOIN testmulticorn m2 ON m1.test1 = m2.test1 LIMIT 1 OFFSET 1;
+
 DROP USER MAPPING FOR current_user SERVER multicorn_srv;
 DROP EXTENSION multicorn cascade;
