pyspark.sql.datasource.DataSourceReader.pushLimit

DataSourceReader.pushLimit(limit)

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 partitions() and read(). By default, it returns False, indicating that the limit cannot be pushed down. Subclasses can override this method to implement limit pushdown.

pushFilters() is called before this method only when the query has filters that Spark can push down; for a query without them, pushFilters() is not called at all. This method may use state that pushFilters() set when filters were pushed, but must not assume pushFilters() ran: initialize defaults in __init__ so this method works whether or not it did.

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, 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.

Added in version 4.4.0.

Parameters:

limitint

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.

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. Initialize the limit in __init__, because partitions() and read() may run even when pushLimit was not called – for a query without a limit, or when this method returned False:

>>> class MyReader(DataSourceReader):
...     def __init__(self):
...         self.limit = None
...
...     def pushLimit(self, limit):
...         # Save the limit for handling in partitions() and read().
...         self.limit = limit
...         return True
...
...     def partitions(self):
...         # A limit can reduce the number of partitions, 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)]