[SPARK-58551][PYTHON] Python Data Sources Limit Pushdown API

python/docs/source/reference/pyspark.sql/datasource.rst

@@ -32,6 +32,7 @@ Python Data Source
     DataSource.writer
     DataSourceReader.partitions
     DataSourceReader.pushFilters
+    DataSourceReader.pushLimit
     DataSourceReader.read
     DataSourceRegistration.register
     DataSourceStreamReader.commit

python/docs/source/tutorial/sql/python_data_source.rst

@@ -179,6 +179,50 @@ Define the reader logic to generate synthetic data. Use the `faker` library to p
                     row.append(value)
                 yield tuple(row)
 
+Push Down a Limit to a Batch Reader
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When a query only needs the first few rows, a reader can implement ``pushLimit`` to fetch less
+data, for example by adding a page size parameter to a REST request or a ``LIMIT`` clause to a
+SQL query. ``pushLimit`` is called during planning, before ``partitions`` and ``read``, and
+returns whether the reader will make use of the limit. It runs after ``pushFilters`` when the
+query has filters to push down, so state it depends on belongs in ``__init__``.
+
+Pushing down a limit is only a hint: Spark always applies the limit again after the scan, so a
+reader is free to return more rows than requested. Set
+``spark.sql.python.limitPushdown.enabled`` to ``true`` to enable limit pushdown.
+
+.. code-block:: python
+
+    from typing import Dict
+
+    from pyspark.sql.datasource import DataSourceReader, InputPartition
+    from pyspark.sql.types import StructType
+
+    class FakeDataSourceReader(DataSourceReader):
+
+        def __init__(self, schema: StructType, options: Dict[str, str]):
+            self.schema: StructType = schema
+            self.options = options
+            self.limit = None
+
+        def pushLimit(self, limit: int) -> bool:
+            self.limit = limit
+            return True
+
+        def partitions(self):
+            # A limit makes a single request cheaper than a fan-out, since every
+            # partition opens its own connection to the data source.
+            if self.limit is not None:
+                return [InputPartition(None)]
+            return [InputPartition(i) for i in range(16)]
+
+        def read(self, partition):
+            num_rows = int(self.options.get("numRows", 3))
+            if self.limit is not None:
+                num_rows = min(num_rows, self.limit)
+            ...
+
 Implement a Batch Writer
 ~~~~~~~~~~~~~~~~~~~~~~~~
 

python/pyspark/errors/error-conditions.json

@@ -204,7 +204,7 @@
   },
   "DATA_SOURCE_PUSHDOWN_DISABLED": {
     "message": [
-      "<type> implements pushFilters() but filter pushdown is disabled because configuration '<conf>' is false. Set it to true to enable filter pushdown."
+      "<type> implements <method>() but the corresponding pushdown is disabled because configuration '<conf>' is false. Set it to true to enable it."
     ]
   },
   "DATA_SOURCE_RETURN_SCHEMA_MISMATCH": {

python/pyspark/sql/datasource.py

@@ -531,6 +531,13 @@ def pushFilters(self, filters: List["Filter"]) -> Iterable["Filter"]:
         all filters, indicating that no filters can be pushed down. Subclasses can
         override this method to implement filter pushdown.
 
+        .. note::
+            When the query also has a limit to push down (see :meth:`pushLimit`), planning
+            creates a second reader and calls this method on it again with the same filters, so
+            that the reader reaches the same state before the limit is pushed. Implementations
+            must therefore be deterministic -- returning a different set of supported filters
+            the second time fails the query -- and should avoid side effects outside of `self`.
+
         It's recommended to implement this method only for data sources that natively
         support filtering, such as databases and GraphQL APIs.
 
@@ -578,6 +585,73 @@ def pushFilters(self, filters: List["Filter"]) -> Iterable["Filter"]:
         """
         return filters
 
+    def pushLimit(self, limit: int) -> bool:
+        """
+        Called with the maximum number of rows that the query needs from this data source.
+
+        Limit pushdown allows the data source to fetch less data, for example by adding a
+        `LIMIT` clause to a SQL query or a page size parameter to a REST request.
+
+        This method is called once during query planning, before :meth:`partitions` and
+        :meth:`read`. By default, it returns False, indicating that the limit cannot be
+        pushed down. Subclasses can override this method to implement limit pushdown.
+
+        :meth:`pushFilters` is called before this method only when the query has filters that
+        Spark can push down; for a query without them, :meth:`pushFilters` is not called at
+        all. Any state this method relies on must therefore be initialized in `__init__`
+        rather than in :meth:`pushFilters`.
+
+        A limit is only pushed down when every filter was pushed down, because Spark cannot
+        apply a limit before a filter it still has to evaluate itself. To benefit from limit
+        pushdown alongside filters, :meth:`pushFilters` should return an empty iterable.
+
+        Pushing down a limit is only a hint: Spark always applies the limit again after the
+        scan, so it is safe to return True even if `read()` yields more than `limit` rows.
+        Returning True never causes the query to see fewer rows than it requires.
+
+        .. versionadded:: 5.0.0
+
+        Parameters
+        ----------
+        limit : int
+            The maximum number of rows the query needs. Always positive: `LIMIT 0` is
+            optimized into an empty relation and never reaches the data source.
+
+        Returns
+        -------
+        bool
+            True if the data source will use the limit to reduce the amount of data it
+            reads, False otherwise.
+
+        Side effects
+        ------------
+        This method is allowed to modify `self`. The object must remain picklable.
+        Modifications to `self` are visible to the `partitions()` and `read()` methods.
+
+        Notes
+        -----
+        This method is only called when the configuration
+        `spark.sql.python.limitPushdown.enabled` is set to true.
+
+        Examples
+        --------
+        Implement pushLimit to fetch fewer rows from the data source:
+
+        >>> def pushLimit(self, limit):
+        ...     # Save the limit for handling in partitions() and read()
+        ...     self.limit = limit
+        ...     return True
+
+        A limit can also be used to reduce the number of partitions, since every
+        partition opens its own connection to the data source:
+
+        >>> def partitions(self):
+        ...     if self.limit is not None:
+        ...         return [InputPartition(None)]
+        ...     return [InputPartition(i) for i in range(16)]
+        """
+        return False
+
     def partitions(self) -> Sequence[InputPartition]:
         """
         Returns an iterator of partitions for this data source.