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

[ganeshashree]

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)