This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 28e1c525f734 [SPARK-58551][PYTHON] Python Data Sources Limit Pushdown
API
28e1c525f734 is described below
commit 28e1c525f734cb748168be9ac4e99b1cf65bbb34
Author: Ganesha S <[email protected]>
AuthorDate: Mon Aug 24 14:33:10 2026 +0800
[SPARK-58551][PYTHON] Python Data Sources Limit Pushdown API
**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. [...]
**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)
Closes #57752 from ganeshashree/SPARK-58551.
Lead-authored-by: Ganesha S <[email protected]>
Co-authored-by: Ganesha S <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit b51665875a3b9971713076b815cf7f344278b6c7)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../source/reference/pyspark.sql/datasource.rst | 1 +
.../source/tutorial/sql/python_data_source.rst | 44 ++
python/pyspark/errors/error-conditions.json | 2 +-
python/pyspark/sql/datasource.py | 92 ++-
python/pyspark/sql/tests/test_python_datasource.py | 635 +++++++++++++++++++++
.../sql/worker/data_source_pushdown_filters.py | 118 +++-
python/pyspark/sql/worker/plan_data_source_read.py | 27 +-
python/pyspark/sql/worker/utils.py | 41 +-
.../org/apache/spark/sql/internal/SQLConf.scala | 14 +
.../datasources/v2/python/PythonDataSourceV2.scala | 22 +-
.../datasources/v2/python/PythonScan.scala | 19 +-
.../datasources/v2/python/PythonScanBuilder.scala | 98 +++-
.../v2/python/UserDefinedPythonDataSource.scala | 97 +++-
.../execution/python/PythonDataSourceSuite.scala | 276 ++++++++-
14 files changed, 1420 insertions(+), 66 deletions(-)
diff --git a/python/docs/source/reference/pyspark.sql/datasource.rst
b/python/docs/source/reference/pyspark.sql/datasource.rst
index bb52ef26d94f..65a6544d587c 100644
--- a/python/docs/source/reference/pyspark.sql/datasource.rst
+++ b/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
diff --git a/python/docs/source/tutorial/sql/python_data_source.rst
b/python/docs/source/tutorial/sql/python_data_source.rst
index 6e9b3d9bd63c..63a41c65b1eb 100644
--- a/python/docs/source/tutorial/sql/python_data_source.rst
+++ b/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
~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/python/pyspark/errors/error-conditions.json
b/python/pyspark/errors/error-conditions.json
index 07ce4314763a..ae35e98a5e5c 100644
--- a/python/pyspark/errors/error-conditions.json
+++ b/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": {
diff --git a/python/pyspark/sql/datasource.py b/python/pyspark/sql/datasource.py
index ad28136e8c36..c499a0a01736 100644
--- a/python/pyspark/sql/datasource.py
+++ b/python/pyspark/sql/datasource.py
@@ -527,9 +527,18 @@ class DataSourceReader(ABC):
can improve performance by reducing the amount of data that needs to be
processed by Spark.
- This method is called once during query planning. By default, it
returns
- all filters, indicating that no filters can be pushed down. Subclasses
can
- override this method to implement filter pushdown.
+ This method may be called more than once during query planning, on
separate reader
+ instances (see the note below). By default, it returns all filters,
indicating that no
+ filters can be pushed down. Subclasses can override this method to
implement filter
+ pushdown.
+
+ .. note::
+ When limit pushdown is enabled (see :meth:`pushLimit`), planning
may create
+ additional reader instances and call this method on each with the
same filters, so
+ that a reader reaches the same state before a limit is pushed and
read planning runs
+ only after the limit is known. Implementations must therefore be
deterministic --
+ returning a different set of supported filters across calls fails
the query -- and
+ must not rely on 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 +587,83 @@ class DataSourceReader(ABC):
"""
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. This method may use state that :meth:`pushFilters` set when
filters were pushed,
+ but must not assume :meth:`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, :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:: 4.4.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
+ only when this method returns True. When it returns False, planning
proceeds as if
+ `pushLimit` was never called -- `partitions()` and `read()` run on a
reader that did
+ not observe these modifications -- so any state they rely on must be
initialized in
+ `__init__` instead.
+
+ 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 :meth:`partitions` and :meth:`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)]
+ """
+ return False
+
def partitions(self) -> Sequence[InputPartition]:
"""
Returns an iterator of partitions for this data source.
diff --git a/python/pyspark/sql/tests/test_python_datasource.py
b/python/pyspark/sql/tests/test_python_datasource.py
index 675c4b19fad9..95663c9c0ec5 100644
--- a/python/pyspark/sql/tests/test_python_datasource.py
+++ b/python/pyspark/sql/tests/test_python_datasource.py
@@ -409,6 +409,641 @@ class BasePythonDataSourceTestsMixin:
with self.assertRaisesRegex(Exception,
"DATA_SOURCE_PUSHDOWN_DISABLED"):
df.show()
+ def test_limit_pushdown(self):
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.limit = None
+
+ def pushLimit(self, limit: int) -> bool:
+ self.limit = limit
+ return True
+
+ def partitions(self):
+ assert self.limit == 2, self.limit
+ return super().partitions()
+
+ def read(self, partition):
+ assert self.limit == 2, self.limit
+ # Only produce as many rows as the query asked for.
+ for i in range(self.limit):
+ yield (i,)
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(2)
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ def test_limit_pushdown_not_supported(self):
+ class TestDataSourceReader(DataSourceReader):
+ def pushLimit(self, limit: int) -> bool:
+ # The reader cannot make use of the limit.
+ return False
+
+ def read(self, partition):
+ yield from [(0,), (1,), (2,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(2)
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ def test_limit_pushdown_rejected_does_not_plan_under_mutated_state(self):
+ # A reader may mutate itself while considering a limit and then reject
it. partitions()
+ # and read() must not run under that rejected state; the scan must
behave as if pushLimit
+ # was never called.
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.considering_limit = False
+
+ def pushLimit(self, limit: int) -> bool:
+ # Mutate self, then reject the limit.
+ self.considering_limit = True
+ return False
+
+ def partitions(self):
+ assert not self.considering_limit, "partitions() planned under
rejected limit"
+ return [InputPartition(0)]
+
+ def read(self, partition):
+ assert not self.considering_limit, "read() planned under
rejected limit"
+ yield from [(0,), (1,), (2,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(2)
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ def test_limit_pushdown_partitions_planned_after_push_limit(self):
+ # With both filter and limit pushdown enabled, partitions() must be
planned only after
+ # pushLimit -- never in the filter-pushdown pass before the limit is
known -- so a reader
+ # that shapes partitions from the pushed limit sees it, and does not
do (discarded) full
+ # partition discovery first.
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.limit = None
+
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ return [] # accept all filters
+
+ def pushLimit(self, limit: int) -> bool:
+ self.limit = limit
+ return True
+
+ def partitions(self):
+ # Fails if planned before pushLimit set the limit.
+ assert self.limit == 2, self.limit
+ return [InputPartition(0)]
+
+ def read(self, partition):
+ assert self.limit == 2, self.limit
+ yield from [(1,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x =
1").limit(2)
+ assertDataFrameEqual(df, [Row(x=1), Row(x=1)])
+
+ def test_limit_pushdown_over_delivering_reader(self):
+ # A reader that accepts the limit but ignores it must not change the
query result:
+ # Spark always applies the limit again after the scan.
+ class TestDataSourceReader(DataSourceReader):
+ def pushLimit(self, limit: int) -> bool:
+ return True
+
+ def read(self, partition):
+ yield from [(i,) for i in range(100)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(3)
+ self.assertEqual(df.count(), 3)
+
+ def test_limit_pushdown_with_filter(self):
+ # The limit is pushed after the filters, and the reader sees both. All
filters are
+ # accepted here, so no post-scan filter remains to block limit
pushdown.
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.filters = []
+ self.limit = None
+
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ self.filters = list(filters)
+ return []
+
+ def pushLimit(self, limit: int) -> bool:
+ # pushFilters must have been called before pushLimit.
+ assert EqualTo(("x",), 1) in self.filters, self.filters
+ self.limit = limit
+ return True
+
+ def read(self, partition):
+ assert EqualTo(("x",), 1) in self.filters, self.filters
+ assert self.limit == 5, self.limit
+ yield from [(1,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x =
1").limit(5)
+ # All filters are reported as fully pushed, so Spark does not
re-apply them.
+ assertDataFrameEqual(df, [Row(x=1), Row(x=1)])
+
+ def test_limit_pushdown_blocked_by_post_scan_filter(self):
+ # A filter that the reader does not accept stays as a post-scan
filter, which prevents
+ # LIMIT from being pushed: applying the limit before that filter could
drop rows the
+ # query needs. The result must still be correct.
+ class TestDataSourceReader(DataSourceReader):
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ # Accept nothing.
+ return filters
+
+ def pushLimit(self, limit: int) -> bool:
+ raise AssertionError("pushLimit should not be called")
+
+ def read(self, partition):
+ yield from [(1,), (2,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x =
1").limit(5)
+ assertDataFrameEqual(df, [Row(x=1), Row(x=1)])
+
+ def test_limit_pushdown_rejected_with_accepted_filter(self):
+ # The reader accepts the filter but rejects the limit. Because every
filter was pushed,
+ # no post-scan filter blocks the limit, so pushLimit is called -- and
returns False. The
+ # worker then plans no read info (the rejected reader may be mutated),
so build() re-plans
+ # the filters-only read from a fresh reader. That re-plan must
reproduce the pushed-filter
+ # state, and the reader must return only the matching rows.
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.filters = []
+
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ self.filters = list(filters)
+ return [] # accept all filters
+
+ def pushLimit(self, limit: int) -> bool:
+ # pushLimit runs only after the filters were accepted.
+ assert EqualTo(("x",), 1) in self.filters, self.filters
+ return False # reject the limit
+
+ def read(self, partition):
+ # Spark removed the accepted filter, so the reader must apply
it itself.
+ assert EqualTo(("x",), 1) in self.filters, self.filters
+ yield from [(1,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x =
1").limit(5)
+ assertDataFrameEqual(df, [Row(x=1), Row(x=1)])
+
+ def test_limit_pushdown_nondeterministic_push_filters(self):
+ # Pushing a limit replays pushFilters on a fresh reader. Spark has
already committed to
+ # the first pass's filter decision, so a reader that reports a
different supported set
+ # the second time must fail the query instead of silently returning
wrong rows.
+ with tempfile.TemporaryDirectory(prefix="test_limit_pushdown_nondet")
as d:
+ counter_path = os.path.join(d, "calls")
+
+ class TestDataSourceReader(DataSourceReader):
+ def pushFilters(self, filters: List[Filter]) ->
Iterable[Filter]:
+ # Accept everything on the first call, nothing on the
replay.
+ n = 0
+ if os.path.exists(counter_path):
+ with open(counter_path) as f:
+ n = int(f.read().strip() or 0)
+ with open(counter_path, "w") as f:
+ f.write(str(n + 1))
+ return [] if n == 0 else list(filters)
+
+ def pushLimit(self, limit: int) -> bool:
+ return True
+
+ def read(self, partition):
+ yield from [(1,), (2,), (3,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x
= 1").limit(5)
+ with self.assertRaisesRegex(Exception, "must be
deterministic"):
+ df.collect()
+
+ def test_limit_pushdown_rejected_nondeterministic_push_filters(self):
+ # When a pushed limit is rejected, planning falls back to a fresh
reader. That fallback
+ # replay of pushFilters must also be validated: a reader that agrees
on the earlier passes
+ # but diverges on the fallback must fail the query, not silently read
unfiltered rows.
+ with
tempfile.TemporaryDirectory(prefix="test_limit_pushdown_reject_nondet") as d:
+ counter_path = os.path.join(d, "calls")
+
+ class TestDataSourceReader(DataSourceReader):
+ def pushFilters(self, filters: List[Filter]) ->
Iterable[Filter]:
+ n = 0
+ if os.path.exists(counter_path):
+ with open(counter_path) as f:
+ n = int(f.read().strip() or 0)
+ with open(counter_path, "w") as f:
+ f.write(str(n + 1))
+ # Accept everything on the first two passes, nothing on
the fallback replay.
+ return [] if n < 2 else list(filters)
+
+ def pushLimit(self, limit: int) -> bool:
+ return False # reject the limit, triggering the
fresh-reader fallback
+
+ def read(self, partition):
+ yield from [(1,), (2,), (3,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x
= 1").limit(5)
+ with self.assertRaisesRegex(Exception, "must be
deterministic"):
+ df.collect()
+
+ def test_limit_pushdown_zero(self):
+ # `LIMIT 0` never reaches the data source: EliminateLimits rewrites it
to an empty
+ # relation before operator pushdown runs, so the scan is removed
altogether and neither
+ # pushLimit nor read is called.
+ class TestDataSourceReader(DataSourceReader):
+ def pushLimit(self, limit: int) -> bool:
+ raise AssertionError("pushLimit should not be called for LIMIT
0")
+
+ def read(self, partition):
+ raise AssertionError("read should not be called for LIMIT 0")
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(0)
+ assertDataFrameEqual(df, [])
+
+ def test_limit_pushdown_disabled_with_filter_pushdown_enabled(self):
+ # The reader implements pushLimit while limit pushdown is disabled.
Filter pushdown runs
+ # a different planning worker, which must not silently ignore
pushLimit either.
+ class TestDataSourceReader(DataSourceReader):
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ return []
+
+ def pushLimit(self, limit: int) -> bool:
+ return True
+
+ def read(self, partition):
+ yield from [(1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": False,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x =
1").limit(1)
+ with self.assertRaisesRegex(Exception,
"DATA_SOURCE_PUSHDOWN_DISABLED"):
+ df.show()
+
+ def test_filter_pushdown_disabled_with_limit_pushdown_enabled(self):
+ # The reader implements pushFilters while filter pushdown is disabled.
A limit-only scan
+ # runs the limit-pushdown worker, which caches the read info, so
`plan_data_source_read`
+ # never runs. That worker must not silently ignore pushFilters either.
No filter is used
+ # here on purpose: with filter pushdown disabled a filter would stay
above the scan and
+ # block limit pushdown, so the limit worker would never run.
+ class TestDataSourceReader(DataSourceReader):
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ return []
+
+ def pushLimit(self, limit: int) -> bool:
+ return True
+
+ def read(self, partition):
+ yield from [(1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": False,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(1)
+ with self.assertRaisesRegex(Exception,
"DATA_SOURCE_PUSHDOWN_DISABLED"):
+ df.show()
+
+ def test_limit_pushdown_disabled(self):
+ class TestDataSourceReader(DataSourceReader):
+ def pushLimit(self, limit: int) -> bool:
+ assert False
+
+ def read(self, partition):
+ assert False
+
+ class TestDataSource(DataSource):
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": False}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").schema("x
int").load()
+ with self.assertRaisesRegex(Exception,
"DATA_SOURCE_PUSHDOWN_DISABLED"):
+ df.show()
+
+ def test_pushdown_disabled_check_survives_reused_readinfo_cache(self):
+ # The provider caches the pushdown-free read info across scans of the
same relation. That
+ # cache must not hide the DATA_SOURCE_PUSHDOWN_DISABLED check: warming
it with pushdown
+ # enabled and then rescanning the same relation with pushdown disabled
must still raise,
+ # because the reader implements a pushdown method the disabled config
would silently
+ # ignore. The cache records the pushdown flags it was populated under
and recomputes when
+ # they change, so the check runs for the flags currently in effect.
+ class TestDataSourceReader(DataSourceReader):
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ return []
+
+ def read(self, partition):
+ yield from [(0,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load()
+
+ # Warm the provider's read-info cache with filter pushdown enabled.
The query has no
+ # filters, so pushFilters is never called and the check passes.
+ with self.sql_conf({"spark.sql.python.filterPushdown.enabled": True}):
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ # Rescan the same relation with filter pushdown disabled.
`df.select("*")` re-plans the
+ # scan on the same (provider-scoped) data source, so a stale cache
would skip the check.
+ with self.sql_conf({"spark.sql.python.filterPushdown.enabled": False}):
+ with self.assertRaisesRegex(Exception,
"DATA_SOURCE_PUSHDOWN_DISABLED"):
+ df.select("*").collect()
+
+ def
test_pushdown_disabled_check_survives_reused_readinfo_cache_limit(self):
+ # Symmetric to the filter case above, for the limit flag. The
read-info cache is keyed by
+ # both pushdown flags, so a regression that dropped only the limit
flag from the key must
+ # be caught too. Warm the cache with limit pushdown enabled (no LIMIT
in the query, so
+ # pushLimit is never called and the check passes), then rescan the
same relation with
+ # limit pushdown disabled and assert the reader's pushLimit
implementation is reported
+ # instead of silently ignored.
+ class TestDataSourceReader(DataSourceReader):
+ def pushLimit(self, limit: int) -> bool:
+ return True
+
+ def read(self, partition):
+ yield from [(0,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": False}):
+ with self.assertRaisesRegex(Exception,
"DATA_SOURCE_PUSHDOWN_DISABLED"):
+ df.select("*").collect()
+
+ def test_limit_pushdown_not_implemented(self):
+ # A reader that does not implement pushLimit is unaffected when the
conf is on.
+ class TestDataSourceReader(DataSourceReader):
+ def read(self, partition):
+ yield from [(0,), (1,), (2,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(2)
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ def test_limit_pushdown_only_does_not_call_push_filters(self):
+ # A limit-only query (no filters) must not trigger a spurious
pushFilters([]) call, even
+ # when filter pushdown is also enabled: the worker skips pushFilters
when there is nothing
+ # to push. A reader may therefore implement both and still expect
pushFilters to be called
+ # only when the query has pushable filters.
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.limit = None
+
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ raise AssertionError("pushFilters should not be called without
filters")
+
+ def pushLimit(self, limit: int) -> bool:
+ self.limit = limit
+ return True
+
+ def read(self, partition):
+ assert self.limit == 2, self.limit
+ yield from [(0,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(2)
+ assertDataFrameEqual(df, [Row(x=0), Row(x=1)])
+
+ def test_limit_pushdown_invalid_return_type(self):
+ # pushLimit must return a bool. A non-bool return is rejected during
planning.
+ class TestDataSourceReader(DataSourceReader):
+ def pushLimit(self, limit: int):
+ return "yes"
+
+ def read(self, partition):
+ yield from [(0,), (1,), (2,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().limit(2)
+ with self.assertRaisesRegex(Exception,
"DATA_SOURCE_INVALID_RETURN_TYPE"):
+ df.collect()
+
+ def test_filter_pushdown_no_limit_with_limit_pushdown_enabled(self):
+ # Filter pushdown defers planning while limit pushdown is enabled,
expecting a possible
+ # limit pass. When the query has no limit, that pass never comes, so
build() plans the
+ # read with the filters only. The reader accepts the filter (so Spark
removes it from the
+ # plan), which means the deferred re-plan must reproduce the
pushed-filter state and the
+ # reader itself must return only the matching rows.
+ class TestDataSourceReader(DataSourceReader):
+ def __init__(self):
+ self.filters = []
+
+ def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+ self.filters = list(filters)
+ return [] # accept all filters
+
+ def pushLimit(self, limit: int) -> bool:
+ raise AssertionError("pushLimit should not be called without a
limit")
+
+ def read(self, partition):
+ # Spark removed the accepted filter, so the reader is
responsible for it.
+ assert EqualTo(("x",), 1) in self.filters, self.filters
+ yield from [(1,), (1,)]
+
+ class TestDataSource(DataSource):
+ def schema(self):
+ return "x int"
+
+ def reader(self, schema) -> "DataSourceReader":
+ return TestDataSourceReader()
+
+ with self.sql_conf(
+ {
+ "spark.sql.python.filterPushdown.enabled": True,
+ "spark.sql.python.limitPushdown.enabled": True,
+ }
+ ):
+ self.spark.dataSource.register(TestDataSource)
+ df = self.spark.read.format("TestDataSource").load().filter("x =
1")
+ assertDataFrameEqual(df, [Row(x=1), Row(x=1)])
+
def _check_filters(self, sql_type, sql_filter, python_filters):
"""
Parameters
diff --git a/python/pyspark/sql/worker/data_source_pushdown_filters.py
b/python/pyspark/sql/worker/data_source_pushdown_filters.py
index a64993642299..14562921a2b3 100644
--- a/python/pyspark/sql/worker/data_source_pushdown_filters.py
+++ b/python/pyspark/sql/worker/data_source_pushdown_filters.py
@@ -45,7 +45,11 @@ from pyspark.sql.datasource import (
)
from pyspark.sql.types import StructType, VariantVal,
_parse_datatype_json_string
from pyspark.sql.worker.plan_data_source_read import
write_read_func_and_partitions
-from pyspark.sql.worker.utils import worker_run
+from pyspark.sql.worker.utils import (
+ check_pushdown_not_disabled,
+ is_method_overridden,
+ worker_run,
+)
from pyspark.worker_util import (
get_sock_file_to_executor,
pickleSer,
@@ -113,24 +117,37 @@ def deserializeFilter(jsonDict: dict) -> Filter:
def _main(infile: IO, outfile: IO) -> None:
"""
- Main method for planning a data source read with filter pushdown.
+ Main method for planning a data source read with filter and limit pushdown.
- This process is invoked from the
`UserDefinedPythonDataSourceReadRunner.runInPython`
- method in the optimizer rule `PlanPythonDataSourceScan` in JVM. This
process is responsible
- for creating a `DataSourceReader` object, applying filter pushdown, and
sending the
- information needed back to the JVM.
+ This process is invoked from the
`UserDefinedPythonDataSourceFilterPushdownRunner.runInPython`
+ method, which `PythonScanBuilder` runs on the JVM while pushing filters
and a limit into the
+ scan. This process is responsible for creating a `DataSourceReader`
object, applying filter
+ and limit pushdown, and sending the information needed back to the JVM.
The infile and outfile are connected to the JVM via a socket. The JVM
sends the following
information to this process via the socket:
- a `DataSource` instance representing the data source
- a `StructType` instance representing the output schema of the data source
- a list of filters to be pushed down
+ - the limit to be pushed down, or -1 if there is none
- configuration values
This process then creates a `DataSourceReader` instance by calling the
`reader` method
on the `DataSource` instance. It applies the filters by calling the
`pushFilters` method
- on the reader and determines which filters are supported. The indices of
the supported
- filters are sent back to the JVM, along with the list of partitions and
the read function.
+ on the reader and determines which filters are supported.
+
+ When a limit is sent, it is pushed down by calling the `pushLimit` method
on the reader
+ after `pushFilters`, and whether the reader accepted it is sent back to
the JVM. The JVM
+ replays the same filters when pushing down a limit, so that the reader
reaches the same
+ state as it did during filter pushdown before `pushLimit` is called on it.
+
+ The values sent back to the JVM, in order, are:
+ - a flag (1 or 0) for whether a read function and partitions follow. They
do not when the
+ filter-pushdown pass defers planning (limit pushdown is enabled, so a
limit pass may follow)
+ or when a pushed limit was rejected; the JVM then plans the read itself,
once it knows
+ whether a limit is pushed. When the flag is 1, the read function and
partitions follow.
+ - the indices of the supported filters.
+ - whether the limit was pushed down (1 or 0).
"""
# Receive the data source instance.
data_source = read_command(pickleSer, infile)
@@ -173,9 +190,12 @@ def _main(infile: IO, outfile: IO) -> None:
filter_dicts = json.loads(json_str)
filters = [FilterRef(deserializeFilter(f)) for f in filter_dicts]
- # Push down the filters and get the indices of the unsupported filters.
+ # Push down the filters and get the indices of the unsupported
filters. `pushFilters` is
+ # not called when there is nothing to push, so that a reader planning
a limit-only scan
+ # does not observe a spurious empty pushFilters call.
unsupported_filters = set(
- FilterRef(f) for f in reader.pushFilters([ref.filter for ref in
filters])
+ FilterRef(f)
+ for f in (reader.pushFilters([ref.filter for ref in filters]) if
filters else [])
)
supported_filter_indices = []
for i, filter in enumerate(filters):
@@ -195,30 +215,88 @@ def _main(infile: IO, outfile: IO) -> None:
},
)
+ # Receive the limit to push down. -1 means there is no limit.
+ limit = read_int(infile)
+
# Receive the max arrow batch size.
max_arrow_batch_size = read_int(infile)
assert max_arrow_batch_size > 0, (
"The maximum arrow batch size should be greater than 0, but got "
f"'{max_arrow_batch_size}'"
)
+ enable_filter_pushdown = read_bool(infile)
+ enable_limit_pushdown = read_bool(infile)
binary_as_bytes = read_bool(infile)
+ # Whether this worker should plan (and send) the read function and
partitions. The JVM
+ # sets this to False for the filter-pushdown pass when limit pushdown
is enabled, so that
+ # `partitions()`/`read()` are not planned before a possible
`pushLimit`. A later limit
+ # pass (or build-time planning on the JVM) then plans once all
pushdowns are known,
+ # matching the public contract that `pushLimit` runs before
`partitions()`/`read()`.
+ plan_read_info = read_bool(infile)
- # Return the read function and partitions. Doing this in the same
worker
- # as filter pushdown helps reduce the number of Python worker calls.
- write_read_func_and_partitions(
- outfile,
- reader=reader,
- data_source=data_source,
- schema=schema,
- max_arrow_batch_size=max_arrow_batch_size,
- binary_as_bytes=binary_as_bytes,
- )
+ # Do not silently ignore a pushdown method that the reader implements
while the
+ # corresponding pushdown is disabled. This worker also runs when only
one of filter or
+ # limit pushdown is enabled -- e.g. a limit-only scan caches the read
info here, so
+ # `plan_data_source_read` never runs and cannot validate the other
method -- so both are
+ # validated here as well as in `plan_data_source_read`.
+ check_pushdown_not_disabled(reader, enable_filter_pushdown,
enable_limit_pushdown)
+
+ # Push down the limit, if any. This must happen after pushFilters,
matching the
+ # operator order that DSv2 uses on the JVM side.
+ #
+ # Whether the read function and partitions follow. Start from what the
JVM requested
+ # (`plan_read_info`), then turn it off if the reader rejects a pushed
limit (below).
+ is_limit_pushed = False
+ send_read_info = plan_read_info
+ # Only call `pushLimit` if the reader actually overrides it: the
inherited default is a
+ # no-op that returns False, so calling it would needlessly drive the
rejection path (and
+ # an extra reader construction on the JVM) for readers that do not
implement limit
+ # pushdown.
+ if limit >= 0 and is_method_overridden(reader, "pushLimit"):
+ is_limit_pushed = reader.pushLimit(limit)
+ if not isinstance(is_limit_pushed, bool):
+ raise PySparkValueError(
+ errorClass="DATA_SOURCE_INVALID_RETURN_TYPE",
+ messageParameters={
+ "type": type(is_limit_pushed).__name__,
+ "name": type(reader).__name__ + ".pushLimit",
+ "supported_types": "bool",
+ },
+ )
+ if not is_limit_pushed:
+ # The reader considered the limit and rejected it, possibly
mutating itself. Do
+ # not plan `partitions()`/`read()` from this reader; send no
read info so the JVM
+ # plans the filters-only read from a fresh reader instead.
That path also
+ # re-validates that the replayed `pushFilters` makes the same
decision, keeping
+ # the rejected reader state out of the scan.
+ send_read_info = False
+
+ # Send whether the read function and partitions follow, before them,
so the JVM knows
+ # whether to read them (and so a planning exception is still surfaced
first). None follow
+ # when the filter-pushdown pass defers planning (limit pushdown
enabled) or when a pushed
+ # limit was rejected; the JVM then plans once it knows whether a limit
is pushed, never on
+ # a reader whose plan would be discarded.
+ write_int(int(send_read_info), outfile)
+ if send_read_info:
+ # Planning here in the same worker as filter/limit pushdown avoids
an extra Python
+ # worker call.
+ write_read_func_and_partitions(
+ outfile,
+ reader=reader,
+ data_source=data_source,
+ schema=schema,
+ max_arrow_batch_size=max_arrow_batch_size,
+ binary_as_bytes=binary_as_bytes,
+ )
# Return the supported filter indices.
write_int(len(supported_filter_indices), outfile)
for index in supported_filter_indices:
write_int(index, outfile)
+ # Return whether the limit was pushed down, as 1 or 0.
+ write_int(int(is_limit_pushed), outfile)
+
def main(infile: IO, outfile: IO) -> None:
worker_run(_main, infile, outfile)
diff --git a/python/pyspark/sql/worker/plan_data_source_read.py
b/python/pyspark/sql/worker/plan_data_source_read.py
index 1de63bd74166..5de6ae71334c 100644
--- a/python/pyspark/sql/worker/plan_data_source_read.py
+++ b/python/pyspark/sql/worker/plan_data_source_read.py
@@ -42,7 +42,7 @@ from pyspark.sql.types import (
BinaryType,
StructType,
)
-from pyspark.sql.worker.utils import worker_run
+from pyspark.sql.worker.utils import check_pushdown_not_disabled, worker_run
from pyspark.worker_util import (
get_sock_file_to_executor,
read_command,
@@ -283,9 +283,15 @@ def _main(infile: IO, outfile: IO) -> None:
for creating a `DataSourceReader` object and send the information needed
back to the JVM.
The infile and outfile are connected to the JVM via a socket. The JVM
sends the following
- information to this process via the socket:
+ information to this process via the socket, in this order (the protocol is
positional):
- a `DataSource` instance representing the data source
+ - a `StructType` instance representing the input schema from the child plan
- a `StructType` instance representing the output schema of the data source
+ - the max Arrow batch size (int)
+ - whether filter pushdown is enabled (bool)
+ - whether limit pushdown is enabled (bool)
+ - whether this is a streaming read (bool)
+ - whether binary values are returned as `bytes` (bool)
This process then creates a `DataSourceReader` instance by calling the
`reader` method
on the `DataSource` instance. Then it calls the `partitions()` method of
the reader and
@@ -339,6 +345,7 @@ def _main(infile: IO, outfile: IO) -> None:
f"The maximum arrow batch size should be greater than 0, but got
'{max_arrow_batch_size}'"
)
enable_pushdown = read_bool(infile)
+ enable_limit_pushdown = read_bool(infile)
is_streaming = read_bool(infile)
binary_as_bytes = read_bool(infile)
@@ -360,19 +367,9 @@ def _main(infile: IO, outfile: IO) -> None:
"actual": f"'{type(reader).__name__}'",
},
)
- is_pushdown_implemented = (
- getattr(reader.pushFilters, "__func__", None) is not
DataSourceReader.pushFilters
- )
- if is_pushdown_implemented and not enable_pushdown:
- # Do not silently ignore pushFilters when pushdown is disabled.
- # Raise an error to ask the user to enable pushdown.
- raise PySparkAssertionError(
- errorClass="DATA_SOURCE_PUSHDOWN_DISABLED",
- messageParameters={
- "type": type(reader).__name__,
- "conf": "spark.sql.python.filterPushdown.enabled",
- },
- )
+ # Do not silently ignore a pushdown method that the reader
implements while the
+ # corresponding pushdown is disabled. Raise an error to ask the
user to enable it.
+ check_pushdown_not_disabled(reader, enable_pushdown,
enable_limit_pushdown)
# Send the read function and partitions to the JVM.
write_read_func_and_partitions(
diff --git a/python/pyspark/sql/worker/utils.py
b/python/pyspark/sql/worker/utils.py
index 12bdb25e6252..1a25918a3a9c 100644
--- a/python/pyspark/sql/worker/utils.py
+++ b/python/pyspark/sql/worker/utils.py
@@ -17,7 +17,7 @@
import os
import sys
-from typing import Callable, IO, Optional
+from typing import Any, Callable, IO, Optional
from pyspark.accumulators import (
_accumulatorRegistry,
@@ -56,6 +56,45 @@ class RunnerConf(Conf):
return self.get("spark.sql.pyspark.dataSource.profiler", None)
+def is_method_overridden(reader: Any, name: str) -> bool:
+ """
+ Whether `reader` overrides the `DataSourceReader` method `name`, rather
than inheriting the
+ default implementation. Used to detect pushdown methods that a reader
implements while the
+ corresponding pushdown configuration is disabled, so that they are not
silently ignored.
+ """
+ from pyspark.sql.datasource import DataSourceReader
+
+ return getattr(getattr(reader, name), "__func__", None) is not
getattr(DataSourceReader, name)
+
+
+def check_pushdown_not_disabled(
+ reader: Any, enable_filter_pushdown: bool, enable_limit_pushdown: bool
+) -> None:
+ """
+ Raise `DATA_SOURCE_PUSHDOWN_DISABLED` if `reader` implements a pushdown
method while the
+ corresponding pushdown configuration is disabled, so that the method is
not silently ignored.
+
+ This is shared by both planning workers: `plan_data_source_read` runs it
for a plain read,
+ and `data_source_pushdown_filters` runs it too, because a filter- or
limit-only scan caches
+ the read info in that worker and `plan_data_source_read` never runs for
such a scan.
+ """
+ from pyspark.errors import PySparkAssertionError
+
+ for method, conf, enabled in (
+ ("pushFilters", "spark.sql.python.filterPushdown.enabled",
enable_filter_pushdown),
+ ("pushLimit", "spark.sql.python.limitPushdown.enabled",
enable_limit_pushdown),
+ ):
+ if not enabled and is_method_overridden(reader, method):
+ raise PySparkAssertionError(
+ errorClass="DATA_SOURCE_PUSHDOWN_DISABLED",
+ messageParameters={
+ "type": type(reader).__name__,
+ "method": method,
+ "conf": conf,
+ },
+ )
+
+
@with_faulthandler
def worker_run(main: Callable, infile: IO, outfile: IO) -> None:
try:
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index b284bf03c3c2..1db7b42a01ff 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -6948,6 +6948,18 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val PYTHON_LIMIT_PUSHDOWN_ENABLED =
buildConf("spark.sql.python.limitPushdown.enabled")
+ .internal()
+ .doc("When true, enable limit pushdown to Python datasource. Pushing a
limit runs a Python " +
+ "worker during planning; for a limit-only scan this replaces the worker
that plans a " +
+ "plain read, while for a scan that also pushes down filters it runs in
addition to " +
+ "filter pushdown. Spark always applies the limit again after the scan,
so a pushed limit " +
+ "only lets the data source read less data.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(false)
+
val CSV_FILTER_PUSHDOWN_ENABLED =
buildConf("spark.sql.csv.filterPushdown.enabled")
.doc("When true, enable filter pushdown to CSV datasource.")
.version("3.0.0")
@@ -9654,6 +9666,8 @@ class SQLConf extends Serializable with Logging with
SqlApiConf {
def pythonFilterPushDown: Boolean = getConf(PYTHON_FILTER_PUSHDOWN_ENABLED)
+ def pythonLimitPushDown: Boolean = getConf(PYTHON_LIMIT_PUSHDOWN_ENABLED)
+
def csvFilterPushDown: Boolean = getConf(CSV_FILTER_PUSHDOWN_ENABLED)
def jsonFilterPushDown: Boolean = getConf(JSON_FILTER_PUSHDOWN_ENABLED)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala
index 7c113c1cb03a..9fcf81746fe3 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.python
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.connector.catalog._
import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -53,24 +54,35 @@ class PythonDataSourceV2 extends TableProvider {
}
private var readInfo: PythonDataSourceReadInfo = _
+ // The pushdown flags in effect when `readInfo` was cached, as (filter,
limit). The cached read
+ // info is identical regardless of these flags, but planning it also
validates -- via
+ // DATA_SOURCE_PUSHDOWN_DISABLED -- that the reader does not implement a
pushdown method whose
+ // config is off. That validation runs only when the planning worker runs,
so reusing the cache
+ // across a change in these flags (e.g. a reused relation scanned first with
pushdown on, then
+ // off) would skip it. Track the flags and recompute when they differ so the
check always runs
+ // for the flags currently in effect.
+ private var readInfoPushdownFlags: Option[(Boolean, Boolean)] = None
+ // Caches the pushdown-free read info for reads that push nothing down. This
is safe to keep on
+ // the (provider-scoped) data source because every such scan produces the
same full read info.
+ // Pushdown-specific read info is never stored here; it is carried by the
`PythonScan` instead,
+ // so that it cannot leak across scans that share this data source (e.g. a
base DataFrame and
+ // its `.limit(n)`).
def getOrCreateReadInfo(
shortName: String,
options: CaseInsensitiveStringMap,
outputSchema: StructType,
isStreaming: Boolean
): PythonDataSourceReadInfo = {
- if (readInfo == null) {
+ val flags = (SQLConf.get.pythonFilterPushDown,
SQLConf.get.pythonLimitPushDown)
+ if (readInfo == null || !readInfoPushdownFlags.contains(flags)) {
val creationResult = getOrCreateDataSourceInPython(shortName, options,
Some(outputSchema))
readInfo = source.createReadInfoInPython(creationResult, outputSchema,
isStreaming)
+ readInfoPushdownFlags = Some(flags)
}
readInfo
}
- def setReadInfo(readInfo: PythonDataSourceReadInfo): Unit = {
- this.readInfo = readInfo
- }
-
override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
getOrCreateDataSourceInPython(shortName, options, None).schema
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala
index 9e3effe7d441..b91165943227 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala
@@ -31,9 +31,13 @@ class PythonScan(
shortName: String,
outputSchema: StructType,
options: CaseInsensitiveStringMap,
- supportedFilters: Array[Filter]
+ supportedFilters: Array[Filter],
+ pushedLimit: Option[Int] = None,
+ // Read info computed during filter/limit pushdown, if any. Scoped to this
scan so it does not
+ // leak across scans that share the same `PythonDataSourceV2` (see
`PythonScanBuilder`).
+ readInfo: Option[PythonDataSourceReadInfo] = None
) extends Scan with SupportsMetadata {
- override def toBatch: Batch = new PythonBatch(ds, shortName, outputSchema,
options)
+ override def toBatch: Batch = new PythonBatch(ds, shortName, outputSchema,
options, readInfo)
override def toMicroBatchStream(checkpointLocation: String):
MicroBatchStream = {
val runner = PythonMicroBatchStream.createPythonStreamingSourceRunner(
@@ -67,7 +71,7 @@ class PythonScan(
Map(
"PushedFilters" -> supportedFilters.mkString("[", ", ", "]"),
"ReadSchema" -> outputSchema.simpleString
- )
+ ) ++ pushedLimit.map(limit => "PushedLimit" -> s"LIMIT $limit")
}
}
@@ -75,7 +79,11 @@ class PythonBatch(
ds: PythonDataSourceV2,
shortName: String,
outputSchema: StructType,
- options: CaseInsensitiveStringMap) extends Batch {
+ options: CaseInsensitiveStringMap,
+ // Read info already computed during pushdown, if any. When empty (no
pushdown), it is
+ // computed lazily below via the provider, which is safe because that path
always produces
+ // the full, pushdown-free read info.
+ readInfo: Option[PythonDataSourceReadInfo] = None) extends Batch {
private val jobArtifactUUID =
JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
private val sessionUUID = {
SparkSession.getActiveSession.collect {
@@ -85,7 +93,8 @@ class PythonBatch(
}
private lazy val infoInPython: PythonDataSourceReadInfo = {
- ds.getOrCreateReadInfo(shortName, options, outputSchema, isStreaming =
false)
+ readInfo.getOrElse(
+ ds.getOrCreateReadInfo(shortName, options, outputSchema, isStreaming =
false))
}
override def planInputPartitions(): Array[InputPartition] =
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala
index 3dabbcb8af05..2ea1bec70cd1 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala
@@ -16,7 +16,8 @@
*/
package org.apache.spark.sql.execution.datasources.v2.python
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownLimit}
+import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.sources.Filter
import org.apache.spark.sql.types.StructType
@@ -29,25 +30,59 @@ class PythonScanBuilder(
outputSchema: StructType,
options: CaseInsensitiveStringMap)
extends ScanBuilder
- with SupportsPushDownFilters {
+ with SupportsPushDownFilters
+ with SupportsPushDownLimit {
private var supportedFilters: Array[Filter] = Array.empty
+ // All filters handed to `pushFilters`, kept so that `pushLimit` can replay
them and bring the
+ // Python reader back to the same state before calling `pushLimit` on it.
+ private var allFilters: Array[Filter] = Array.empty
+ private var pushedLimit: Option[Int] = None
+ // Read info (partitions + read function) produced as a side effect of
filter/limit pushdown,
+ // carried into the `PythonScan` so it stays scoped to this scan. It must
NOT be stored on the
+ // provider-scoped `PythonDataSourceV2`: a single data source instance is
shared by every scan
+ // built from it -- e.g. a base DataFrame and its `.limit(n)` reuse the same
relation -- so a
+ // pushdown-specific read function stored there would leak into an unrelated
scan and make it
+ // read too few rows.
+ private var readInfo: Option[PythonDataSourceReadInfo] = None
+ // True when a pushdown pass ran but produced no read info, so build() must
plan the read
+ // itself. Two paths set it: (1) the filter-pushdown pass defers planning
while limit pushdown
+ // is enabled and a limit pass might follow; (2) a pushed limit was
rejected, so planning falls
+ // back to a fresh reader. In both, build() plans the read with the filters
only. This is
+ // distinct from `readInfo.isEmpty` when both pushdowns are disabled -- in
that case, no
+ // pushdown pass ran and `PythonScan` plans the read later instead, so
build() must not plan.
+ private var deferredPlanning: Boolean = false
- override def build(): Scan =
- new PythonScan(ds, shortName, outputSchema, options, supportedFilters)
+ override def build(): Scan = {
+ if (deferredPlanning && readInfo.isEmpty) {
+ // Reached when a pushdown pass ran but planned no read info: either the
filter pass
+ // deferred planning and no limit followed, or a pushed limit was
rejected (including a
+ // limit-only scan with no filters). Plan now, with the filters only and
no limit, so that
+ // partitions()/read() run exactly once and only after all pushdowns are
known -- never in
+ // the filter pass before a possible pushLimit.
+ val dataSource = ds.getOrCreateDataSourceInPython(shortName, options,
Some(outputSchema))
+ val result = ds.source.pushdownLimitInPython(dataSource, outputSchema,
allFilters, None)
+ checkFiltersReplayedDeterministically(result)
+ readInfo = result.readInfo
+ }
+ new PythonScan(ds, shortName, outputSchema, options, supportedFilters,
pushedLimit, readInfo)
+ }
// Optionally called by DSv2 once to push down filters before the scan is
built.
override def pushFilters(filters: Array[Filter]): Array[Filter] = {
if (!SQLConf.get.pythonFilterPushDown) {
return filters
}
+ allFilters = filters
val dataSource = ds.getOrCreateDataSourceInPython(shortName, options,
Some(outputSchema))
ds.source.pushdownFiltersInPython(dataSource, outputSchema, filters) match
{
case None => filters // No filters are supported.
case Some(result) =>
- // Filter pushdown also returns partitions and the read function.
- // This helps reduce the number of Python worker calls.
- ds.setReadInfo(result.readInfo)
+ // Filter pushdown may also return the read function and partitions.
It defers that
+ // planning when limit pushdown is enabled -- so partitions()/read()
are not planned
+ // before a possible pushLimit -- in which case the limit pass or
build() plans instead.
+ readInfo = result.readInfo
+ deferredPlanning = result.readInfo.isEmpty
// Partition the filters into supported and unsupported ones.
val isPushed = result.isFilterPushed.zip(filters)
@@ -58,4 +93,53 @@ class PythonScanBuilder(
}
override def pushedFilters(): Array[Filter] = supportedFilters
+
+ // Optionally called by DSv2 once to push down a LIMIT before the scan is
built. DSv2 calls this
+ // after `pushFilters`, so the filters that were pushed there (none, for a
query without
+ // pushable filters) are replayed here to rebuild the same reader state
before `pushLimit` is
+ // invoked on it in Python.
+ override def pushLimit(limit: Int): Boolean = {
+ if (!SQLConf.get.pythonLimitPushDown) {
+ return false
+ }
+
+ val dataSource = ds.getOrCreateDataSourceInPython(shortName, options,
Some(outputSchema))
+ val result = ds.source.pushdownLimitInPython(dataSource, outputSchema,
allFilters, Some(limit))
+ checkFiltersReplayedDeterministically(result)
+
+ // Adopt the read info the worker planned. It is empty when the reader
rejected the limit --
+ // the worker plans nothing then, to keep the rejected (possibly mutated)
reader state out of
+ // the scan -- in which case build() plans the filters-only read from a
fresh reader and
+ // re-validates the replayed filters. Record the pushed limit only when
the reader accepted it.
+ readInfo = result.readInfo
+ deferredPlanning = result.readInfo.isEmpty
+ if (result.isLimitPushed) {
+ pushedLimit = Some(limit)
+ }
+ result.isLimitPushed
+ }
+
+ // The read is replayed on a fresh reader (to push a limit, or to plan after
filter pushdown
+ // deferred). That replay runs `pushFilters` again, which must reach the
same decision as the
+ // first pass. Spark has already committed to that first decision --
`pushedFilters()` was read
+ // by the optimizer and the filters it reported were dropped from the plan
-- so a reader whose
+ // `pushFilters` is not deterministic would leave Spark applying the first
pass's filters while
+ // reading with the replayed reader, silently returning wrong rows. Fail
fast instead.
+ private def checkFiltersReplayedDeterministically(result:
PythonFilterPushdownResult): Unit = {
+ val replayedFilters = result.isFilterPushed.zip(allFilters).collect {
+ case (true, filter) => filter
+ }.toArray
+ if (!replayedFilters.sameElements(supportedFilters)) {
+ throw QueryCompilationErrors.pythonDataSourceError(
+ action = "plan",
+ tpe = "read",
+ msg = "pushFilters() returned a different set of supported filters
when it was replayed " +
+ s"during planning: [${supportedFilters.mkString(", ")}] then " +
+ s"[${replayedFilters.mkString(", ")}]. pushFilters() must be
deterministic.")
+ }
+ }
+
+ // Spark always applies the LIMIT again after the scan: a Python data source
is not trusted to
+ // return at most `limit` rows, and it is free to over-deliver.
+ override def isPartiallyPushed(): Boolean = true
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala
index 7611cf676499..7a36118a1879 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala
@@ -81,7 +81,11 @@ case class UserDefinedPythonDataSource(dataSourceCls:
PythonFunction) {
val runner = new UserDefinedPythonDataSourceFilterPushdownRunner(
createPythonFunction(pythonResult.dataSource),
outputSchema,
- filters
+ filters,
+ limit = None,
+ // Defer planning while limit pushdown is enabled: a later limit pass,
or build(), plans
+ // once it is known whether a limit is pushed. Plan here only when limit
pushdown is off.
+ planReadInfo = !SQLConf.get.pythonLimitPushDown
)
if (runner.isAnyFilterSupported) {
Some(runner.runInPython())
@@ -90,6 +94,36 @@ case class UserDefinedPythonDataSource(dataSourceCls:
PythonFunction) {
}
}
+ /**
+ * (Driver-side) Run Python process to plan the read while optionally
pushing down a limit, and
+ * return a [[PythonFilterPushdownResult]] carrying which filters were
pushed, whether the limit
+ * was pushed, and the planned read info. The worker plans the read function
and partitions in
+ * this pass, except when a pushed `limit` is rejected: the reader may have
mutated itself while
+ * rejecting it, so the result's `readInfo` is then `None` and the caller
plans the filters-only
+ * read from a fresh reader instead.
+ *
+ * `limit` is the limit to push, or `None` to plan the read from the filters
only -- the
+ * build-time path uses `None` when the filter-pushdown pass deferred
planning and no limit was
+ * ultimately pushed.
+ *
+ * `filters` must be the filters that were previously pushed via
[[pushdownFiltersInPython]] --
+ * empty for a limit-only scan, where no filter pushdown ran -- so that
replaying them brings
+ * the reader to the same state it had during filter pushdown before
`pushLimit` is called.
+ */
+ def pushdownLimitInPython(
+ pythonResult: PythonDataSourceCreationResult,
+ outputSchema: StructType,
+ filters: Array[Filter],
+ limit: Option[Int]): PythonFilterPushdownResult = {
+ new UserDefinedPythonDataSourceFilterPushdownRunner(
+ createPythonFunction(pythonResult.dataSource),
+ outputSchema,
+ filters,
+ limit = limit,
+ planReadInfo = true
+ ).runInPython()
+ }
+
/**
* (Driver-side) Run Python process, and get the partition read functions,
and
* partition information.
@@ -345,15 +379,22 @@ private class UserDefinedPythonDataSourceRunner(
}
/**
+ * @param readInfo The read function and partitions, when the worker planned
them. Empty in two
+ * cases: the filter-pushdown pass deferred planning (limit
pushdown enabled), or a
+ * pushed limit was rejected (the rejecting reader may be
mutated). Planning then
+ * happens on a later limit pass or at build time, from a
fresh reader.
* @param isFilterPushed A sequence of bools indicating whether each filter is
pushed down.
+ * @param isLimitPushed Whether the limit is pushed down. False when no limit
was sent, and also
+ * when a limit was sent but the reader rejected it.
*/
case class PythonFilterPushdownResult(
- readInfo: PythonDataSourceReadInfo,
- isFilterPushed: collection.Seq[Boolean]
+ readInfo: Option[PythonDataSourceReadInfo],
+ isFilterPushed: collection.Seq[Boolean],
+ isLimitPushed: Boolean = false
)
/**
- * Push down filters to a Python data source.
+ * Push down filters and optionally a limit to a Python data source.
*
* @param dataSource
* a Python data source instance
@@ -361,11 +402,19 @@ case class PythonFilterPushdownResult(
* output schema of the Python data source
* @param filters
* all filters to be pushed down
+ * @param limit
+ * the limit to be pushed down, if any
*/
private class UserDefinedPythonDataSourceFilterPushdownRunner(
dataSource: PythonFunction,
schema: StructType,
- filters: collection.Seq[Filter])
+ filters: collection.Seq[Filter],
+ limit: Option[Int],
+ // Whether the worker should plan (and send back) the read function and
partitions in this
+ // pass. Kept separate from `limit` so a build-time, filter-only planning
pass can request
+ // planning without a limit -- deriving it from `limit.isDefined` would
force such a pass to
+ // smuggle a dummy limit value in just to trigger planning.
+ planReadInfo: Boolean)
extends PythonPlannerRunner[PythonFilterPushdownResult](dataSource) {
private case class SerializedFilter(
@@ -476,14 +525,41 @@ private class
UserDefinedPythonDataSourceFilterPushdownRunner(
// Send the filters
PythonWorkerUtils.writeUTF(mapper.writeValueAsString(serializedFilters),
dataOut)
+ // Send the limit, if any. -1 means no limit is pushed down.
+ dataOut.writeInt(limit.getOrElse(-1))
+
// Send configurations
dataOut.writeInt(SQLConf.get.arrowMaxRecordsPerBatch)
+ dataOut.writeBoolean(SQLConf.get.pythonFilterPushDown)
+ dataOut.writeBoolean(SQLConf.get.pythonLimitPushDown)
dataOut.writeBoolean(SQLConf.get.pysparkBinaryAsBytes)
+ // Whether the worker should plan the read function and partitions. The
filter-pushdown pass
+ // sets this false while limit pushdown is enabled, because a later limit
pass -- or
+ // build-time planning -- will plan once it is known whether a limit is
pushed, so
+ // `partitions()`/`read()` never run on a reader whose plan would be
discarded.
+ dataOut.writeBoolean(planReadInfo)
}
override protected def receiveFromPython(dataIn: DataInputStream):
PythonFilterPushdownResult = {
- // Receive the read function and the partitions. Also check for exceptions.
- val readInfo = PythonDataSourceReadInfo.receive(dataIn)
+ // Whether a read function + partitions follow. Read first so that an
exception raised in the
+ // worker (sent as PYTHON_EXCEPTION_THROWN) is surfaced here instead of
being misread as data.
+ val hasReadInfo = dataIn.readInt()
+ if (hasReadInfo == SpecialLengths.PYTHON_EXCEPTION_THROWN) {
+ val msg = PythonWorkerUtils.readUTF(dataIn)
+ throw QueryCompilationErrors.pythonDataSourceError(
+ action = "plan",
+ tpe = "read",
+ msg = msg
+ )
+ }
+ // Receive the read function and partitions, when the worker planned them.
None follow when the
+ // filter-pushdown pass defers planning (limit pushdown enabled) or when a
pushed limit was
+ // rejected; build() (or a later limit pass) then plans them from a fresh
reader instead.
+ val readInfo = if (hasReadInfo != 0) {
+ Some(PythonDataSourceReadInfo.receive(dataIn))
+ } else {
+ None
+ }
// Receive the pushed filters as a list of indices.
val numFiltersPushed = dataIn.readInt()
@@ -493,9 +569,13 @@ private class
UserDefinedPythonDataSourceFilterPushdownRunner(
isFilterPushed(serializedFilters(i).index) = true
}
+ // Receive whether the limit was pushed down, sent as 1 or 0.
+ val isLimitPushed = dataIn.readInt() != 0
+
PythonFilterPushdownResult(
readInfo = readInfo,
- isFilterPushed = isFilterPushed
+ isFilterPushed = isFilterPushed,
+ isLimitPushed = isLimitPushed
)
}
}
@@ -575,6 +655,7 @@ private class UserDefinedPythonDataSourceReadRunner(
// Send configurations
dataOut.writeInt(SQLConf.get.arrowMaxRecordsPerBatch)
dataOut.writeBoolean(SQLConf.get.pythonFilterPushDown)
+ dataOut.writeBoolean(SQLConf.get.pythonLimitPushDown)
dataOut.writeBoolean(isStreaming)
dataOut.writeBoolean(SQLConf.get.pysparkBinaryAsBytes)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
index 9e47473b42d7..de96f23e519b 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
@@ -22,7 +22,7 @@ import java.io.{File, FileWriter}
import org.apache.spark.api.python.PythonException
import org.apache.spark.api.python.PythonUtils
import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, Row}
-import org.apache.spark.sql.execution.FilterExec
+import org.apache.spark.sql.execution.{FilterExec, LimitExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.datasources.DataSourceManager
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation}
@@ -308,6 +308,280 @@ class PythonDataSourceSuite extends
PythonDataSourceSuiteBase {
}
}
+ test("data source reader with limit pushdown") {
+ assume(shouldTestPandasUDFs)
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader,
InputPartition
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def __init__(self):
+ | self.limit = None
+ |
+ | def pushLimit(self, limit):
+ | self.limit = limit
+ | return True
+ |
+ | def partitions(self):
+ | # A pushed limit lets the reader plan a single partition.
+ | assert self.limit == 2, self.limit
+ | return [InputPartition(0)]
+ |
+ | def read(self, partition):
+ | for i in range(self.limit):
+ | yield (i,)
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") {
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df = spark.read.format(dataSourceName).schema(schema).load().limit(2)
+ val plan = df.queryExecution.executedPlan
+
+ collectFirst(plan) {
+ case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] =>
+ val p = s.scan.asInstanceOf[PythonScan]
+ assert(p.getMetaData().get("PushedLimit").contains("LIMIT 2"))
+ }.getOrElse(
+ fail(s"PythonScan not found in the plan. Actual plan:\n$plan")
+ )
+
+ // Spark keeps its own limit: a Python data source is not trusted to
honor the pushed limit.
+ assert(collectFirst(plan) { case l: LimitExec => l }.isDefined,
+ s"A limit operator should be retained. Actual plan:\n$plan")
+
+ checkAnswer(df, Seq(Row(0), Row(1)))
+ }
+ }
+
+ test("data source reader limit pushdown not supported by the reader") {
+ assume(shouldTestPandasUDFs)
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def pushLimit(self, limit):
+ | return False
+ |
+ | def read(self, partition):
+ | yield (0,)
+ | yield (1,)
+ | yield (2,)
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") {
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df = spark.read.format(dataSourceName).schema(schema).load().limit(2)
+ val plan = df.queryExecution.executedPlan
+
+ collectFirst(plan) {
+ case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] =>
+ val p = s.scan.asInstanceOf[PythonScan]
+ assert(!p.getMetaData().contains("PushedLimit"))
+ }.getOrElse(
+ fail(s"PythonScan not found in the plan. Actual plan:\n$plan")
+ )
+
+ checkAnswer(df, Seq(Row(0), Row(1)))
+ }
+ }
+
+ test("limit pushdown is not used for top-N (orderBy then limit)") {
+ assume(shouldTestPandasUDFs)
+ // A LIMIT after an ORDER BY is a top-N: the limit cannot be pushed to the
scan, because
+ // pushing it before the sort would drop rows the sort needs. The reader
must not see it.
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader,
InputPartition
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def pushLimit(self, limit):
+ | raise AssertionError("pushLimit must not be called for a
top-N query")
+ |
+ | def partitions(self):
+ | return [InputPartition(0)]
+ |
+ | def read(self, partition):
+ | yield from [(2,), (1,), (3,)]
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") {
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df =
spark.read.format(dataSourceName).schema(schema).load().orderBy("id").limit(2)
+ val plan = df.queryExecution.executedPlan
+
+ collectFirst(plan) {
+ case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] =>
+ val p = s.scan.asInstanceOf[PythonScan]
+ assert(!p.getMetaData().contains("PushedLimit"))
+ }.getOrElse(
+ fail(s"PythonScan not found in the plan. Actual plan:\n$plan")
+ )
+
+ checkAnswer(df, Seq(Row(1), Row(2)))
+ }
+ }
+
+ test("no limit is pushed for a filter-only query with both pushdowns
enabled") {
+ assume(shouldTestPandasUDFs)
+ // The filter-pushdown worker also handles limit pushdown, but a query
without a LIMIT sends
+ // the -1 no-limit sentinel, so `pushLimit` must not be called even when
the reader implements
+ // it and limit pushdown is enabled.
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader,
InputPartition
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def pushFilters(self, filters):
+ | # Reject the filter so Spark re-applies it; this test is
only about the limit.
+ | return filters
+ |
+ | def pushLimit(self, limit):
+ | raise AssertionError("pushLimit must not be called when the
query has no limit")
+ |
+ | def partitions(self):
+ | return [InputPartition(0)]
+ |
+ | def read(self, partition):
+ | yield from [(1,), (2,)]
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ withSQLConf(
+ SQLConf.PYTHON_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") {
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df =
spark.read.format(dataSourceName).schema(schema).load().filter("id = 1")
+ val plan = df.queryExecution.executedPlan
+
+ collectFirst(plan) {
+ case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] =>
+ val p = s.scan.asInstanceOf[PythonScan]
+ assert(!p.getMetaData().contains("PushedLimit"))
+ }.getOrElse(
+ fail(s"PythonScan not found in the plan. Actual plan:\n$plan")
+ )
+
+ checkAnswer(df, Seq(Row(1)))
+ }
+ }
+
+ test("limit pushdown does not leak into a reused base scan") {
+ assume(shouldTestPandasUDFs)
+ // A base DataFrame and its `.limit(n)` share the same Python data source
instance. The read
+ // function produced while pushing the limit must stay scoped to the
limited scan and not
+ // leak into the base scan, which would make the base DataFrame read too
few rows.
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader,
InputPartition
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def pushLimit(self, limit):
+ | self._limit = limit
+ | return True
+ |
+ | def partitions(self):
+ | return [InputPartition(0)]
+ |
+ | def read(self, partition):
+ | n = getattr(self, "_limit", 10)
+ | for i in range(n):
+ | yield (i,)
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") {
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df = spark.read.format(dataSourceName).schema(schema).load()
+ checkAnswer(df.limit(2), Seq(Row(0), Row(1)))
+ checkAnswer(df, (0 until 10).map(Row(_)))
+ }
+ }
+
+ test("filter pushdown does not leak into a reused base scan") {
+ assume(shouldTestPandasUDFs)
+ // Same as above for filter pushdown: the read function bound to the
pushed filters must not
+ // leak into the base scan, which would make the base DataFrame return
only filtered rows.
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader,
InputPartition
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def pushFilters(self, filters):
+ | self._filtered = True
+ | return []
+ |
+ | def partitions(self):
+ | return [InputPartition(0)]
+ |
+ | def read(self, partition):
+ | if getattr(self, "_filtered", False):
+ | yield (1,)
+ | else:
+ | for i in range(10):
+ | yield (i,)
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ withSQLConf(SQLConf.PYTHON_FILTER_PUSHDOWN_ENABLED.key -> "true") {
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df = spark.read.format(dataSourceName).schema(schema).load()
+ checkAnswer(df.filter("id = 1"), Seq(Row(1)))
+ checkAnswer(df, (0 until 10).map(Row(_)))
+ }
+ }
+
test("register data source") {
assume(shouldTestPandasUDFs)
val dataSourceScript =
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
