diff --git a/python/docs/source/reference/pyspark.sql/datasource.rst b/python/docs/source/reference/pyspark.sql/datasource.rst index bb52ef26d94f7..65a6544d587c5 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 6e9b3d9bd63cb..63a41c65b1eb4 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 07ce4314763a4..ae35e98a5e5c0 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -204,7 +204,7 @@ }, "DATA_SOURCE_PUSHDOWN_DISABLED": { "message": [ - " implements pushFilters() but filter pushdown is disabled because configuration '' is false. Set it to true to enable filter pushdown." + " implements () but the corresponding pushdown is disabled because configuration '' 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 ad28136e8c361..2d5d7699250da 100644 --- a/python/pyspark/sql/datasource.py +++ b/python/pyspark/sql/datasource.py @@ -531,6 +531,13 @@ def pushFilters(self, filters: List["Filter"]) -> Iterable["Filter"]: all filters, indicating that no filters can be pushed down. Subclasses can override this method to implement filter pushdown. + .. note:: + When the query also has a limit to push down (see :meth:`pushLimit`), planning + creates a second reader and calls this method on it again with the same filters, so + that the reader reaches the same state before the limit is pushed. Implementations + must therefore be deterministic -- returning a different set of supported filters + the second time fails the query -- and should avoid side effects outside of `self`. + It's recommended to implement this method only for data sources that natively support filtering, such as databases and GraphQL APIs. @@ -578,6 +585,73 @@ def pushFilters(self, filters: List["Filter"]) -> Iterable["Filter"]: """ return filters + def pushLimit(self, limit: int) -> bool: + """ + Called with the maximum number of rows that the query needs from this data source. + + Limit pushdown allows the data source to fetch less data, for example by adding a + `LIMIT` clause to a SQL query or a page size parameter to a REST request. + + This method is called once during query planning, before :meth:`partitions` and + :meth:`read`. By default, it returns False, indicating that the limit cannot be + pushed down. Subclasses can override this method to implement limit pushdown. + + :meth:`pushFilters` is called before this method only when the query has filters that + Spark can push down; for a query without them, :meth:`pushFilters` is not called at + all. Any state this method relies on must therefore be initialized in `__init__` + rather than in :meth:`pushFilters`. + + A limit is only pushed down when every filter was pushed down, because Spark cannot + apply a limit before a filter it still has to evaluate itself. To benefit from limit + pushdown alongside filters, :meth:`pushFilters` should return an empty iterable. + + Pushing down a limit is only a hint: Spark always applies the limit again after the + scan, so it is safe to return True even if `read()` yields more than `limit` rows. + Returning True never causes the query to see fewer rows than it requires. + + .. versionadded:: 5.0.0 + + Parameters + ---------- + limit : int + The maximum number of rows the query needs. Always positive: `LIMIT 0` is + optimized into an empty relation and never reaches the data source. + + Returns + ------- + bool + True if the data source will use the limit to reduce the amount of data it + reads, False otherwise. + + Side effects + ------------ + This method is allowed to modify `self`. The object must remain picklable. + Modifications to `self` are visible to the `partitions()` and `read()` methods. + + Notes + ----- + This method is only called when the configuration + `spark.sql.python.limitPushdown.enabled` is set to true. + + Examples + -------- + Implement pushLimit to fetch fewer rows from the data source: + + >>> def pushLimit(self, limit): + ... # Save the limit for handling in partitions() and read() + ... self.limit = limit + ... return True + + A limit can also be used to reduce the number of partitions, since every + partition opens its own connection to the data source: + + >>> def partitions(self): + ... if self.limit is not None: + ... return [InputPartition(None)] + ... return [InputPartition(i) for i in range(16)] + """ + return False + def partitions(self) -> Sequence[InputPartition]: """ Returns an iterator of partitions for this data source. diff --git a/python/pyspark/sql/tests/test_python_datasource.py b/python/pyspark/sql/tests/test_python_datasource.py index 675c4b19fad94..284f802e25de1 100644 --- a/python/pyspark/sql/tests/test_python_datasource.py +++ b/python/pyspark/sql/tests/test_python_datasource.py @@ -409,6 +409,285 @@ def reader(self, schema) -> "DataSourceReader": 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_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_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_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_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_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 _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 a649936422996..074c6b723c8cd 100644 --- a/python/pyspark/sql/worker/data_source_pushdown_filters.py +++ b/python/pyspark/sql/worker/data_source_pushdown_filters.py @@ -45,7 +45,7 @@ ) 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 is_method_overridden, worker_run from pyspark.worker_util import ( get_sock_file_to_executor, pickleSer, @@ -113,24 +113,30 @@ 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. + 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. + + 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. """ # Receive the data source instance. data_source = read_command(pickleSer, infile) @@ -173,9 +179,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,14 +204,46 @@ 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_limit_pushdown = read_bool(infile) binary_as_bytes = read_bool(infile) + if not enable_limit_pushdown and is_method_overridden(reader, "pushLimit"): + # Do not silently ignore pushLimit when limit pushdown is disabled. This worker also + # runs for filter-only pushdown, where no limit is ever sent, so the check has to + # happen here as well as in `plan_data_source_read`. + raise PySparkAssertionError( + errorClass="DATA_SOURCE_PUSHDOWN_DISABLED", + messageParameters={ + "type": type(reader).__name__, + "method": "pushLimit", + "conf": "spark.sql.python.limitPushdown.enabled", + }, + ) + + # Push down the limit, if any. This must happen after pushFilters, matching the + # operator order that DSv2 uses on the JVM side. + is_limit_pushed = False + if limit >= 0: + 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", + }, + ) + # 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( @@ -219,6 +260,9 @@ def _main(infile: IO, outfile: IO) -> None: 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 1de63bd74166d..8236cec7ae370 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 @@ BinaryType, StructType, ) -from pyspark.sql.worker.utils import worker_run +from pyspark.sql.worker.utils import is_method_overridden, worker_run from pyspark.worker_util import ( get_sock_file_to_executor, read_command, @@ -339,6 +339,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 +361,21 @@ 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. + for method, conf, enabled in ( + ("pushFilters", "spark.sql.python.filterPushdown.enabled", enable_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, + }, + ) # 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 12bdb25e62529..6b06af28e1c38 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,17 @@ def profiler(self) -> Optional[str]: 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) + + @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 403de67e56e7a..25306724b30ef 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 @@ -6774,6 +6774,16 @@ 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, at the cost of running " + + "Python worker one additional time during planning. Spark always applies the limit " + + "again after the scan, so a pushed limit only lets the data source read less data.") + .version("5.0.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") @@ -9396,6 +9406,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/PythonScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala index 9e3effe7d441d..2eb56d711fce1 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,7 +31,8 @@ class PythonScan( shortName: String, outputSchema: StructType, options: CaseInsensitiveStringMap, - supportedFilters: Array[Filter] + supportedFilters: Array[Filter], + pushedLimit: Option[Int] = None ) extends Scan with SupportsMetadata { override def toBatch: Batch = new PythonBatch(ds, shortName, outputSchema, options) @@ -67,7 +68,7 @@ class PythonScan( Map( "PushedFilters" -> supportedFilters.mkString("[", ", ", "]"), "ReadSchema" -> outputSchema.simpleString - ) + ) ++ pushedLimit.map(limit => "PushedLimit" -> s"LIMIT $limit") } } 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 3dabbcb8af05b..0abf2b4f2ef91 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,17 +30,23 @@ 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 override def build(): Scan = - new PythonScan(ds, shortName, outputSchema, options, supportedFilters) + new PythonScan(ds, shortName, outputSchema, options, supportedFilters, pushedLimit) // 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 { @@ -58,4 +65,49 @@ 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, limit) + + // The replay runs `pushFilters` on a fresh reader, 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 second pass's reader, silently returning wrong rows. + // Fail fast instead. + val replayedFilters = result.isFilterPushed.zip(allFilters).collect { + case (true, filter) => filter + }.toArray + if (!replayedFilters.sameElements(supportedFilters)) { + throw QueryCompilationErrors.pythonDataSourceError( + action = "plan", + tpe = "limit", + msg = "pushFilters() returned a different set of supported filters when it was called " + + "again to push down a limit: " + + s"[${supportedFilters.mkString(", ")}] then [${replayedFilters.mkString(", ")}]. " + + "pushFilters() must be deterministic.") + } + + // Limit pushdown also returns partitions and the read function, which reflect the pushed + // limit. They are valid whether or not the limit was pushed, because the filters were + // replayed identically, so always replace the read info computed by filter pushdown. + ds.setReadInfo(result.readInfo) + if (result.isLimitPushed) { + pushedLimit = Some(limit) + } + result.isLimitPushed + } + + // 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 7611cf6764995..fb59a2c9c489d 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,8 @@ case class UserDefinedPythonDataSource(dataSourceCls: PythonFunction) { val runner = new UserDefinedPythonDataSourceFilterPushdownRunner( createPythonFunction(pythonResult.dataSource), outputSchema, - filters + filters, + limit = None ) if (runner.isAnyFilterSupported) { Some(runner.runInPython()) @@ -90,6 +91,27 @@ case class UserDefinedPythonDataSource(dataSourceCls: PythonFunction) { } } + /** + * (Driver-side) Run Python process to push down a limit, get the updated data source + * instance and the pushdown result. + * + * `filters` must be the same filters that were previously passed to + * [[pushdownFiltersInPython]], so that the reader reaches the same state before + * `pushLimit` is called on it. + */ + def pushdownLimitInPython( + pythonResult: PythonDataSourceCreationResult, + outputSchema: StructType, + filters: Array[Filter], + limit: Int): PythonFilterPushdownResult = { + new UserDefinedPythonDataSourceFilterPushdownRunner( + createPythonFunction(pythonResult.dataSource), + outputSchema, + filters, + limit = Some(limit) + ).runInPython() + } + /** * (Driver-side) Run Python process, and get the partition read functions, and * partition information. @@ -346,14 +368,16 @@ private class UserDefinedPythonDataSourceRunner( /** * @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. */ case class PythonFilterPushdownResult( readInfo: PythonDataSourceReadInfo, - isFilterPushed: collection.Seq[Boolean] + 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 +385,14 @@ 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]) extends PythonPlannerRunner[PythonFilterPushdownResult](dataSource) { private case class SerializedFilter( @@ -476,8 +503,12 @@ 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.pythonLimitPushDown) dataOut.writeBoolean(SQLConf.get.pysparkBinaryAsBytes) } @@ -493,9 +524,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 +610,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 9e47473b42d75..bc096b5b8949c 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,102 @@ 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("register data source") { assume(shouldTestPandasUDFs) val dataSourceScript =