Initial implementation

Author: abramsh

Makefile

@@ -123,6 +123,7 @@ TESTS        = test-$(PYTHON_TEST_VERSION)/sql/multicorn_cache_invalidation.sql
   test-$(PYTHON_TEST_VERSION)/sql/multicorn_sequence_test.sql \
   test-$(PYTHON_TEST_VERSION)/sql/multicorn_test_date.sql \
   test-$(PYTHON_TEST_VERSION)/sql/multicorn_test_dict.sql \
+  test-$(PYTHON_TEST_VERSION)/sql/multicorn_test_limit.sql \
   test-$(PYTHON_TEST_VERSION)/sql/multicorn_test_list.sql \
   test-$(PYTHON_TEST_VERSION)/sql/multicorn_test_sort.sql \
   test-$(PYTHON_TEST_VERSION)/sql/write_savepoints.sql \

flake.nix

@@ -156,7 +156,7 @@
         pythonVersion = pkgs.lib.versions.majorMinor test_python.version;
         isPython312OrHigher = pkgs.lib.versionAtLeast pythonVersion "3.12";
 
-        baseTestCount = if pkgs.lib.versionOlder pgMajorVersion "14" then 18 else 19;
+        baseTestCount = if pkgs.lib.versionOlder pgMajorVersion "14" then 19 else 20;
         expectedTestCount = toString (baseTestCount - (if isPython312OrHigher then 1 else 0));
       in pkgs.stdenv.mkDerivation {
         name = "multicorn2-python-test-pg${test_postgresql.version}-py${test_python.version}";

python/multicorn/__init__.py

@@ -212,6 +212,21 @@ def can_sort(self, sortkeys):
         """
         return []
 
+    def can_limit(self, limit, offset):
+        """
+        Method called from the planner to ask the FDW whether it supports LIMIT pushdown.
+        This method is only called if the rest of the query can be pushed down including the sort and quals. For example,
+        if the query has a GROUP BY clause, this method will not be called.
+
+        Args:
+            limit (int or None): The limit to apply to the query.
+            offset (int or None): The offset to apply to the query.
+
+        Return:
+            True if the FDW can support both LIMIT and OFFSET pushdown, Falseotherwise.
+        """
+        return False
+    
     def get_path_keys(self):
         u"""
         Method called from the planner to add additional Path to the planner.
@@ -269,7 +284,7 @@ def get_path_keys(self):
         """
         return []
 
-    def explain(self, quals, columns, sortkeys=None, verbose=False):
+    def explain(self, quals, columns, sortkeys=None, verbose=False, limit=None, offset=None):
         """Hook called on explain.
 
         The arguments are the same as the :meth:`execute`, with the addition of
@@ -280,7 +295,7 @@ def explain(self, quals, columns, sortkeys=None, verbose=False):
         """
         return []
 
-    def execute(self, quals, columns, sortkeys=None):
+    def execute(self, quals, columns, sortkeys=None, limit=None, offset=None):
         """Execute a query in the foreign data wrapper.
 
         This method is called at the first iteration.
@@ -313,6 +328,8 @@ def execute(self, quals, columns, sortkeys=None):
                 should be in the sequence.
             sortkeys (list): A list of :class:`SortKey`
                 that the FDW said it can enforce.
+            limit (int): The limit to apply to the query.
+            offset (int): The offset to apply to the query.
 
         Returns:
             An iterable of python objects which can be converted back to PostgreSQL.

python/multicorn/testfdw.py

@@ -5,7 +5,7 @@
 from itertools import cycle
 from datetime import datetime
 from operator import itemgetter
-
+from itertools import islice
 
 class TestForeignDataWrapper(ForeignDataWrapper):
 
@@ -15,6 +15,12 @@ def __init__(self, options, columns):
         super(TestForeignDataWrapper, self).__init__(options, columns)
         self.columns = columns
         self.test_type = options.get('test_type', None)
+        self.canlimit = options.get('canlimit', False)
+        if isinstance(self.canlimit, str):
+            self.canlimit = self.canlimit.lower() == 'true'
+        self.cansort = options.get('cansort', True)
+        if isinstance(self.cansort, str):
+            self.cansort = self.cansort.lower() == 'true'
         self.test_subtype = options.get('test_subtype', None)
         self.tx_hook = options.get('tx_hook', False)
         self._modify_batch_size = int(options.get('modify_batch_size', 1))
@@ -79,7 +85,7 @@ def _as_generator(self, quals, columns):
                                                           index)
             yield line
 
-    def execute(self, quals, columns, sortkeys=None):
+    def execute(self, quals, columns, sortkeys=None, limit=None, offset=None):
         sortkeys = sortkeys or []
         log_to_postgres(str(sorted(quals)))
         log_to_postgres(str(sorted(columns)))
@@ -99,14 +105,15 @@ def execute(self, quals, columns, sortkeys=None):
                 k = sortkeys[0];
                 res = self._as_generator(quals, columns)
                 if (self.test_type == 'sequence'):
-                    return sorted(res, key=itemgetter(k.attnum - 1),
+                    res = sorted(res, key=itemgetter(k.attnum - 1),
                                   reverse=k.is_reversed)
                 else:
-                    return sorted(res, key=itemgetter(k.attname),
+                    res = sorted(res, key=itemgetter(k.attname),
                                   reverse=k.is_reversed)
-            return self._as_generator(quals, columns)
+                return res[offset:offset + limit] if offset else res[:limit]
+            return islice(self._as_generator(quals, columns), offset, (offset or 0) + limit if limit else None)
 
-    def explain(self, quals, columns, sortkeys=None, verbose=False):
+    def explain(self, quals, columns, sortkeys=None, verbose=False, limit=None, offset=None):
         if self.noisy_explain:
             log_to_postgres("EXPLAIN quals=%r" % (sorted(quals),))
             log_to_postgres("EXPLAIN columns=%r" % (sorted(columns),))
@@ -127,8 +134,13 @@ def get_path_keys(self):
 
     def can_sort(self, sortkeys):
         # assume sort pushdown ok for all cols, in any order, any collation
+        if not self.cansort:
+            return []
         return sortkeys
 
+    def can_limit(self, limit, offset):
+        return self.canlimit
+
     def update(self, rowid, newvalues):
         if self.test_type == 'nowrite':
             super(TestForeignDataWrapper, self).update(rowid, newvalues)

src/multicorn.c

@@ -56,6 +56,12 @@ static void multicornGetForeignRelSize(PlannerInfo *root,
 static void multicornGetForeignPaths(PlannerInfo *root,
 						 RelOptInfo *baserel,
 						 Oid foreigntableid);
+static void multicornGetForeignUpperPaths(PlannerInfo *root, 
+                            UpperRelationKind stage,
+                            RelOptInfo *input_rel,
+                            RelOptInfo *output_rel, 
+                            void *extra);
+
 static ForeignScan *multicornGetForeignPlan(PlannerInfo *root,
 						RelOptInfo *baserel,
 						Oid foreigntableid,
@@ -119,6 +125,11 @@ static void multicorn_xact_callback(XactEvent event, void *arg);
 void	   *serializePlanState(MulticornPlanState * planstate);
 MulticornExecState *initializeExecState(void *internal_plan_state);
 
+static void add_foreign_final_paths(PlannerInfo *root,
+                                    RelOptInfo *input_rel,
+                                    RelOptInfo *final_rel,
+                                    FinalPathExtraData *extra);
+                                    
 /* Hash table mapping oid to fdw instances */
 HTAB	   *InstancesHash;
 
@@ -174,6 +185,7 @@ multicorn_handler(PG_FUNCTION_ARGS)
 	/* Plan phase */
 	fdw_routine->GetForeignRelSize = multicornGetForeignRelSize;
 	fdw_routine->GetForeignPaths = multicornGetForeignPaths;
+	fdw_routine->GetForeignUpperPaths = multicornGetForeignUpperPaths;
 	fdw_routine->GetForeignPlan = multicornGetForeignPlan;
 	fdw_routine->ExplainForeignScan = multicornExplainForeignScan;
 
@@ -359,7 +371,7 @@ multicornGetForeignPaths(PlannerInfo *root,
 	/* Try to find parameterized paths */
 	pathes = findPaths(root, baserel, possiblePaths, planstate->startupCost,
 			planstate, apply_pathkeys, deparsed_pathkeys);
-
+    
 	/* Add a simple default path */
 	pathes = lappend(pathes, create_foreignscan_path(root, baserel,
 		 	NULL,  /* default pathtarget */
@@ -403,6 +415,11 @@ multicornGetForeignPaths(PlannerInfo *root,
 		{
 			ForeignPath *newpath;
 
+            MulticornPathState *pathstate = (MulticornPathState *)calloc(1, sizeof(MulticornPathState));
+            pathstate->pathkeys = deparsed_pathkeys;
+            pathstate->limit = -1;
+            pathstate->offset = -1;
+        
 			newpath = create_foreignscan_path(root, baserel,
 					NULL,  /* default pathtarget */
 					path->path.rows,
@@ -415,7 +432,7 @@ multicornGetForeignPaths(PlannerInfo *root,
 #if PG_VERSION_NUM >= 170000
 					NULL,
 #endif
-					(void *) deparsed_pathkeys);
+					(void *)pathstate);
 
 			newpath->path.param_info = path->path.param_info;
 			add_path(baserel, (Path *) newpath);
@@ -424,6 +441,150 @@ multicornGetForeignPaths(PlannerInfo *root,
 	errorCheck();
 }
 
+/*
+ * multicornGetForeignUpperPaths
+ *		Add paths for post-join operations like aggregation, grouping etc. if
+ *		corresponding operations are safe to push down.
+ *
+ * Right now, we only support limit/offset pushdown.  We'll add others later.
+ */
+static void multicornGetForeignUpperPaths(PlannerInfo *root, 
+                                          UpperRelationKind stage,
+                                          RelOptInfo *input_rel,
+                                          RelOptInfo *output_rel, 
+                                          void *extra)
+{
+    switch (stage)
+	{
+		case UPPERREL_FINAL:
+            add_foreign_final_paths(root, input_rel, output_rel, (FinalPathExtraData *)extra);
+			break;
+		default:
+			break;
+	}
+}
+
+/*
+ * add_foreign_final_paths
+ *		Add foreign paths for performing the final processing remotely.
+ *
+ * Given input_rel contains the source-data Paths.  The paths are added to the
+ * given final_rel.
+ */
+ static void
+ add_foreign_final_paths(PlannerInfo *root, RelOptInfo *input_rel,
+                         RelOptInfo *final_rel,
+                         FinalPathExtraData *extra)
+ {
+    Query *parse = root->parse;
+    MulticornPathState *pathstate;
+    MulticornPlanState *planstate;
+    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)
+        return;
+    
+    /* We only support limits for SELECT commands */
+    if (parse->commandType != CMD_SELECT)
+        return;
+
+    /* We do not support pushing down FETCH FIRST .. WITH TIES */
+    if (parse->limitOption == LIMIT_OPTION_WITH_TIES)
+        return;
+
+    /* only push down constant LIMITs... */
+    if ((parse->limitCount && !IsA(parse->limitCount, Const)) || (parse->limitOffset && !IsA(parse->limitOffset, Const)))
+        return;
+
+    /* ... which are not NULL */
+    if((parse->limitCount && ((Const *)parse->limitCount)->constisnull) || (parse->limitOffset && ((Const *)parse->limitOffset)->constisnull))
+        return;
+
+    /* Extract the limit and offset */
+    if (parse->limitCount)
+        limitCount = DatumGetInt32(((Const *)parse->limitCount)->constvalue);
+    
+    if (parse->limitOffset)
+        limitOffset = DatumGetInt32(((Const *)parse->limitOffset)->constvalue);
+
+    /* 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;
+   
+    /* 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->limit = limitCount;
+    pathstate->offset = limitOffset;
+    final_path = create_foreign_upper_path(root,
+                                        input_rel,
+                                        root->upper_targets[UPPERREL_FINAL],
+                                        limitCount,
+                                        planstate->startupCost,
+                                        limitCount * planstate->width,
+                                        NULL, /* pathkeys will be applied in the input_rel */
+                                        NULL, /* no extra plan */
+#if PG_VERSION_NUM >= 170000
+                                        NULL, /* no fdw_restrictinfo list */
+#endif
+                                        (void*)pathstate);
+    /* and add it to the final_rel */
+    add_path(final_rel, (Path *) final_path);
+}
+
 /*
  * multicornGetForeignPlan
  *		Create a ForeignScan plan node for scanning the foreign table
@@ -458,7 +619,21 @@ multicornGetForeignPlan(PlannerInfo *root,
 				&planstate->qual_list);
 		}
 	}
-	planstate->pathkeys = (List *) best_path->fdw_private;
+
+    if (best_path->fdw_private)
+    {
+        MulticornPathState *pathstate = (MulticornPathState *) best_path->fdw_private;
+        planstate->pathkeys = pathstate->pathkeys;
+        planstate->limit = pathstate->limit;
+        planstate->offset = pathstate->offset;
+    }
+    else
+    {
+        planstate->pathkeys = NIL;
+        planstate->limit = -1;
+        planstate->offset = -1;
+    }
+
 	return make_foreignscan(tlist,
 							scan_clauses,
 							scan_relid,
@@ -1162,13 +1337,19 @@ serializePlanState(MulticornPlanState * state)
 	List	   *result = NULL;
 
 	result = lappend(result, makeConst(INT4OID,
-						  -1, InvalidOid, 4, Int32GetDatum(state->numattrs), false, true));
+					-1, InvalidOid, 4, Int32GetDatum(state->numattrs), false, true));
 	result = lappend(result, makeConst(INT4OID,
 					-1, InvalidOid, 4, Int32GetDatum(state->foreigntableid), false, true));
 	result = lappend(result, state->target_list);
 
 	result = lappend(result, serializeDeparsedSortGroup(state->pathkeys));
 
+    result = lappend(result, makeConst(INT4OID,
+                    -1, InvalidOid, 4, Int32GetDatum(state->limit), false, true));
+
+    result = lappend(result, makeConst(INT4OID,
+                    -1, InvalidOid, 4, Int32GetDatum(state->offset), false, true));
+
 	return result;
 }
 
@@ -1195,5 +1376,7 @@ initializeExecState(void *internalstate)
 	execstate->cinfos = palloc0(sizeof(ConversionInfo *) * attnum);
 	execstate->values = palloc(attnum * sizeof(Datum));
 	execstate->nulls = palloc(attnum * sizeof(bool));
+    execstate->limit = DatumGetInt32(((Const*)list_nth(values,4))->constvalue);
+    execstate->offset = DatumGetInt32(((Const*)list_nth(values,5))->constvalue);
 	return execstate;
 }