ganeshashree opened a new pull request, #57752:
URL: https://github.com/apache/spark/pull/57752
**What changes were proposed in this pull request?**
This PR adds LIMIT pushdown to Python Data Sources, completing the pushdown
family alongside filter pushdown (SPARK-51271).
A new optional DataSourceReader.pushLimit method:
```
def pushLimit(self, limit: int) -> bool:
"""Return True if the reader will use the limit to read less data."""
return False
```
It is called once during planning, before partitions() and read(), so a
reader can use the limit to plan cheaper work:
```
class RestReader(DataSourceReader):
def __init__(self):
self.limit = None
def pushLimit(self, limit):
self.limit = limit
return True
def partitions(self):
# With a limit, one narrow request beats a fan-out.
return [InputPartition(None)] if self.limit else [InputPartition(i)
for i in range(16)]
def read(self, partition):
params = {"max": self.limit} if self.limit else {}
```
Supporting changes:
- PythonScanBuilder mixes in SupportsPushDownLimit. isPartiallyPushed
returns true, so Spark always re-applies the LIMIT after the scan — arbitrary
user Python is not trusted to return at most limit rows, making a pushed limit
purely a hint to read less data.
- DSv2 calls pushLimit after pushFilters, but the planning worker exits
between the two calls, so the filters are replayed on a fresh reader to bring
it to the same state before the limit is pushed. The replayed decision is
validated against the first pass and the query fails fast on mismatch (see the
user-facing note below).
- A pushed limit is reported as PushedLimit in the scan metadata, visible in
df.explain(mode="formatted").
- New internal config spark.sql.python.limitPushdown.enabled (default
false), mirroring spark.sql.python.filterPushdown.enabled, since it costs one
extra Python worker invocation during planning.
- A reader implementing pushLimit while the config is disabled raises
DATA_SOURCE_PUSHDOWN_DISABLED instead of having the method silently ignored,
matching pushFilters. That message template is now parameterized by method name
so it reads correctly for both pushdowns.
**Why are the changes needed?**
Python Data Sources cannot use a query's LIMIT to reduce the work they do. A
DataSourceReader plans its full set of partitions with no knowledge of the
limit, and reads at Arrow batch granularity (10,000 rows by default), so LIMIT
5 over a REST or database-backed source can still cost many paged requests or a
full extract. The reader has no way to add a LIMIT clause, set a page size
parameter, or open fewer connections.
JVM DSv2 sources have had this via SupportsPushDownLimit since 3.3.0. This
is the Python counterpart, alongside SPARK-51713 for column pruning.
**Does this PR introduce any user-facing change?**
Yes, new API. DataSourceReader.pushLimit is only called when the new
spark.sql.python.limitPushdown.enabled config is true (default false), and
readers that do not implement it are unaffected. This lands in an unreleased
branch, so there is no change relative to any released version.
Two semantics deserve reviewer attention:
1. A limit is only pushed when every filter was pushed. Spark cannot apply a
limit before a filter it must still evaluate itself, so any residual post-scan
filter blocks limit pushdown. This is pre-existing shared DSv2 behavior,
asserted for JDBC in JDBCV2Suite ("LIMIT is pushed down only if all the filters
are pushed down"). Practically, a reader must also accept the IsNotNull filters
Spark generates in order to combine filter and limit pushdown.
2. Implementing pushFilters alongside pushLimit requires pushFilters to be
deterministic. The limit push replays the filters on a fresh reader, and by
then Spark has already dropped the first pass's pushed filters from the plan. A
reader reporting a different supported set the second time would leave Spark
reading via the second reader while trusting the first decision — silently
returning wrong rows. This is now validated, failing with a clear error rather
than producing bad results. Both points are documented in the API docs.
**How was this patch tested?**
New tests in PythonDataSourceSuite: the pushed and not-pushed cases, that
Spark retains its own limit operator, and PushedLimit scan metadata.
New tests in test_python_datasource.py:
- a pushed limit reaching partitions() and read()
- a reader declining the limit
- a reader that accepts but ignores the limit — the query still returns
exactly n rows
- limits combined with filters
- a post-scan filter blocking limit pushdown
- non-deterministic pushFilters failing the query
- LIMIT 0, which EliminateLimits turns into an empty relation so it never
reaches the source
- the config disabled, both alone and with filter pushdown enabled
- a reader that does not implement pushLimit
Full suites pass: test_python_datasource (101),
test_python_streaming_datasource (16), PythonDataSourceSuite (25).
**Was this patch authored or co-authored using generative AI tooling?**
Generated-by: Claude Code (Opus 5)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
