You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2023/11/30 23:41:26 UTC
(spark) branch master updated: [SPARK-45940][PYTHON] Add InputPartition to DataSourceReader interface
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 2fad7b7ba28e [SPARK-45940][PYTHON] Add InputPartition to DataSourceReader interface
2fad7b7ba28e is described below
commit 2fad7b7ba28ef29e06b2f6a664ec54fe561e3b41
Author: allisonwang-db <al...@databricks.com>
AuthorDate: Fri Dec 1 08:41:13 2023 +0900
[SPARK-45940][PYTHON] Add InputPartition to DataSourceReader interface
### What changes were proposed in this pull request?
This PR introduces a new Python class `InputPartition` that represents the partition value returned by the `partitions` method in DataSourceReader.
### Why are the changes needed?
Before this PR, the `partitions` method can return anything, and it can be confusing to infer what is the `partition` argument in the `read(self, partition)` method.
Adding `InputPartition` can make the Python data source API more intuitive and user-friendly.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
New unit tests.
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #44085 from allisonwang-db/spark-45940-input-partition.
Authored-by: allisonwang-db <al...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/sql/datasource.py | 86 ++++++++++----
python/pyspark/sql/tests/test_python_datasource.py | 7 +-
python/pyspark/sql/worker/plan_data_source_read.py | 41 ++++++-
.../execution/python/PythonDataSourceSuite.scala | 126 +++++++++++++++++++--
4 files changed, 221 insertions(+), 39 deletions(-)
diff --git a/python/pyspark/sql/datasource.py b/python/pyspark/sql/datasource.py
index 032e3f48a82e..1c5b6d663285 100644
--- a/python/pyspark/sql/datasource.py
+++ b/python/pyspark/sql/datasource.py
@@ -15,7 +15,7 @@
# limitations under the License.
#
from abc import ABC, abstractmethod
-from typing import final, Any, Dict, Iterator, List, Tuple, Type, Union, TYPE_CHECKING
+from typing import final, Any, Dict, Iterator, List, Sequence, Tuple, Type, Union, TYPE_CHECKING
from pyspark.sql import Row
from pyspark.sql.types import StructType
@@ -145,6 +145,44 @@ class DataSource(ABC):
raise NotImplementedError
+class InputPartition:
+ """
+ A base class representing an input partition returned by the `partitions()`
+ method of `DataSourceReader`.
+
+ .. versionadded: 4.0.0
+
+ Notes
+ -----
+ This class must be picklable.
+
+ Examples
+ --------
+ Use the default input partition implementation:
+
+ >>> def partitions(self):
+ ... return [InputPartition(1)]
+
+ Subclass the input partition class:
+
+ >>> from dataclasses import dataclass
+ >>> @dataclass
+ ... class RangeInputPartition(InputPartition):
+ ... start: int
+ ... end: int
+
+ >>> def partitions(self):
+ ... return [RangeInputPartition(1, 3), RangeInputPartition(4, 6)]
+ """
+
+ def __init__(self, value: Any) -> None:
+ self.value = value
+
+ def __repr__(self) -> str:
+ attributes = ", ".join([f"{k}={v!r}" for k, v in self.__dict__.items()])
+ return f"{self.__class__.__name__}({attributes})"
+
+
class DataSourceReader(ABC):
"""
A base class for data source readers. Data source readers are responsible for
@@ -153,7 +191,7 @@ class DataSourceReader(ABC):
.. versionadded: 4.0.0
"""
- def partitions(self) -> Iterator[Any]:
+ def partitions(self) -> Sequence[InputPartition]:
"""
Returns an iterator of partitions for this data source.
@@ -171,40 +209,40 @@ class DataSourceReader(ABC):
Returns
-------
- Iterator[Any]
- An iterator of partitions for this data source. The partition value can be
- any serializable objects.
+ Sequence[InputPartition]
+ A sequence of partitions for this data source. Each partition value
+ must be an instance of `InputPartition` or a subclass of it.
Notes
-----
- This method should not return any un-picklable objects.
+ All partition values must be picklable objects.
Examples
--------
Returns a list of integers:
>>> def partitions(self):
- ... return [1, 2, 3]
+ ... return [InputPartition(1), InputPartition(2), InputPartition(3)]
Returns a list of string:
>>> def partitions(self):
- ... return ["a", "b", "c"]
+ ... return [InputPartition("a"), InputPartition("b"), InputPartition("c")]
- Returns a list of tuples:
+ Returns a list of ranges:
- >>> def partitions(self):
- ... return [("a", 1), ("b", 2), ("c", 3)]
-
- Returns a list of dictionaries:
+ >>> class RangeInputPartition(InputPartition):
+ ... def __init__(self, start, end):
+ ... self.start = start
+ ... self.end = end
>>> def partitions(self):
- ... return [{"a": 1}, {"b": 2}, {"c": 3}]
+ ... return [RangeInputPartition(1, 3), RangeInputPartition(5, 10)]
"""
- yield None
+ raise NotImplementedError
@abstractmethod
- def read(self, partition: Any) -> Iterator[Union[Tuple, Row]]:
+ def read(self, partition: InputPartition) -> Iterator[Union[Tuple, Row]]:
"""
Generates data for a given partition and returns an iterator of tuples or rows.
@@ -229,15 +267,15 @@ class DataSourceReader(ABC):
--------
Yields a list of tuples:
- >>> def read(self, partition):
- ... yield (partition, 0)
- ... yield (partition, 1)
+ >>> def read(self, partition: InputPartition):
+ ... yield (partition.value, 0)
+ ... yield (partition.value, 1)
Yields a list of rows:
- >>> def read(self, partition):
- ... yield Row(partition=partition, value=0)
- ... yield Row(partition=partition, value=1)
+ >>> def read(self, partition: InputPartition):
+ ... yield Row(partition=partition.value, value=0)
+ ... yield Row(partition=partition.value, value=1)
"""
...
@@ -313,6 +351,10 @@ class WriterCommitMessage:
sent back to the driver side as input parameter of ``commit`` or ``abort`` method.
.. versionadded: 4.0.0
+
+ Notes
+ -----
+ This class must be picklable.
"""
...
diff --git a/python/pyspark/sql/tests/test_python_datasource.py b/python/pyspark/sql/tests/test_python_datasource.py
index bab062c48215..8c7074c72a64 100644
--- a/python/pyspark/sql/tests/test_python_datasource.py
+++ b/python/pyspark/sql/tests/test_python_datasource.py
@@ -17,7 +17,7 @@
import os
import unittest
-from pyspark.sql.datasource import DataSource, DataSourceReader
+from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
from pyspark.sql.types import Row
from pyspark.testing import assertDataFrameEqual
from pyspark.testing.sqlutils import ReusedSQLTestCase
@@ -46,7 +46,6 @@ class BasePythonDataSourceTestsMixin:
yield None,
reader = MyDataSourceReader()
- self.assertEqual(list(reader.partitions()), [None])
self.assertEqual(list(reader.read(None)), [(None,)])
def test_data_source_register(self):
@@ -91,10 +90,10 @@ class BasePythonDataSourceTestsMixin:
num_partitions = int(self.options["num_partitions"])
else:
num_partitions = self.DEFAULT_NUM_PARTITIONS
- return range(num_partitions)
+ return [InputPartition(i) for i in range(num_partitions)]
def read(self, partition):
- yield partition, str(partition)
+ yield partition.value, str(partition.value)
class InMemoryDataSource(DataSource):
@classmethod
diff --git a/python/pyspark/sql/worker/plan_data_source_read.py b/python/pyspark/sql/worker/plan_data_source_read.py
index 9d3efbb1b34d..c6abef4509d5 100644
--- a/python/pyspark/sql/worker/plan_data_source_read.py
+++ b/python/pyspark/sql/worker/plan_data_source_read.py
@@ -28,7 +28,7 @@ from pyspark.serializers import (
SpecialLengths,
CloudPickleSerializer,
)
-from pyspark.sql.datasource import DataSource
+from pyspark.sql.datasource import DataSource, InputPartition
from pyspark.sql.types import _parse_datatype_json_string, StructType
from pyspark.util import handle_worker_exception
from pyspark.worker_util import (
@@ -110,10 +110,41 @@ def main(infile: IO, outfile: IO) -> None:
message_parameters={"type": "reader", "error": str(e)},
)
- # Generate all partitions.
- partitions = list(reader.partitions() or [])
- if len(partitions) == 0:
- partitions = [None]
+ # Get the partitions if any.
+ try:
+ partitions = reader.partitions()
+ if not isinstance(partitions, list):
+ raise PySparkRuntimeError(
+ error_class="PYTHON_DATA_SOURCE_CREATE_ERROR",
+ message_parameters={
+ "type": "reader",
+ "error": (
+ "Expect 'partitions' to return a list, but got "
+ f"'{type(partitions).__name__}'"
+ ),
+ },
+ )
+ if not all(isinstance(p, InputPartition) for p in partitions):
+ partition_types = ", ".join([f"'{type(p).__name__}'" for p in partitions])
+ raise PySparkRuntimeError(
+ error_class="PYTHON_DATA_SOURCE_CREATE_ERROR",
+ message_parameters={
+ "type": "reader",
+ "error": (
+ "All elements in 'partitions' should be of type "
+ f"'InputPartition', but got {partition_types}"
+ ),
+ },
+ )
+ if len(partitions) == 0:
+ partitions = [None] # type: ignore
+ except NotImplementedError:
+ partitions = [None] # type: ignore
+ except Exception as e:
+ raise PySparkRuntimeError(
+ error_class="PYTHON_DATA_SOURCE_CREATE_ERROR",
+ message_parameters={"type": "reader", "error": str(e)},
+ )
# Construct a UDTF.
class PythonDataSourceReaderUDTF:
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 33b34b39ab2e..02cdc63e8d3b 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
@@ -29,20 +29,21 @@ class PythonDataSourceSuite extends QueryTest with SharedSparkSession {
private def dataSourceName = "SimpleDataSource"
private def simpleDataSourceReaderScript: String =
"""
+ |from pyspark.sql.datasource import DataSourceReader, InputPartition
|class SimpleDataSourceReader(DataSourceReader):
| def partitions(self):
- | return range(0, 2)
+ | return [InputPartition(i) for i in range(2)]
| def read(self, partition):
- | yield (0, partition)
- | yield (1, partition)
- | yield (2, partition)
+ | yield (0, partition.value)
+ | yield (1, partition.value)
+ | yield (2, partition.value)
|""".stripMargin
test("simple data source") {
assume(shouldTestPythonUDFs)
val dataSourceScript =
s"""
- |from pyspark.sql.datasource import DataSource, DataSourceReader
+ |from pyspark.sql.datasource import DataSource
|$simpleDataSourceReaderScript
|
|class $dataSourceName(DataSource):
@@ -179,7 +180,7 @@ class PythonDataSourceSuite extends QueryTest with SharedSparkSession {
assume(shouldTestPythonUDFs)
val dataSourceScript =
s"""
- |from pyspark.sql.datasource import DataSource, DataSourceReader
+ |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
|import json
|
|class SimpleDataSourceReader(DataSourceReader):
@@ -193,10 +194,14 @@ class PythonDataSourceSuite extends QueryTest with SharedSparkSession {
| paths = [self.options["path"]]
| else:
| paths = []
- | return paths
+ | return [InputPartition(p) for p in paths]
|
| def read(self, path):
- | yield (path, 1)
+ | if path is not None:
+ | assert isinstance(path, InputPartition)
+ | yield (path.value, 1)
+ | else:
+ | yield (path, 1)
|
|class $dataSourceName(DataSource):
| @classmethod
@@ -272,4 +277,109 @@ class PythonDataSourceSuite extends QueryTest with SharedSparkSession {
assert(err.getMessage.contains("PYTHON_DATA_SOURCE_TYPE_MISMATCH"))
assert(err.getMessage.contains("PySparkAssertionError"))
}
+
+ test("data source read with custom partitions") {
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
+ |class RangePartition(InputPartition):
+ | def __init__(self, start, end):
+ | self.start = start
+ | self.end = end
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def partitions(self):
+ | return [RangePartition(1, 2), RangePartition(3, 4)]
+ |
+ | def read(self, partition: RangePartition):
+ | start, end = partition.start, partition.end
+ | for i in range(start, end):
+ | yield (i, )
+ |
+ |class $dataSourceName(DataSource):
+ | def schema(self) -> str:
+ | return "id INT"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val dataSource = createUserDefinedPythonDataSource(dataSourceName, dataSourceScript)
+ val df = dataSource(spark, provider = dataSourceName)
+ checkAnswer(df, Seq(Row(1), Row(3)))
+ }
+
+ test("data source read with empty partitions") {
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def partitions(self):
+ | return []
+ |
+ | def read(self, partition):
+ | if partition is None:
+ | yield ("success", )
+ | else:
+ | yield ("failed", )
+ |
+ |class $dataSourceName(DataSource):
+ | def schema(self) -> str:
+ | return "status STRING"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val dataSource = createUserDefinedPythonDataSource(dataSourceName, dataSourceScript)
+ val df = dataSource(spark, provider = dataSourceName)
+ checkAnswer(df, Row("success"))
+ }
+
+ test("data source read with invalid partitions") {
+ val reader1 =
+ s"""
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def partitions(self):
+ | return 1
+ | def read(self, partition):
+ | ...
+ |""".stripMargin
+
+ val reader2 =
+ s"""
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def partitions(self):
+ | return [1, 2]
+ | def read(self, partition):
+ | ...
+ |""".stripMargin
+
+ val reader3 =
+ s"""
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def partitions(self):
+ | raise Exception("error")
+ | def read(self, partition):
+ | ...
+ |""".stripMargin
+
+ Seq(reader1, reader2, reader3).foreach { readerScript =>
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import DataSource, DataSourceReader
+ |$readerScript
+ |
+ |class $dataSourceName(DataSource):
+ | def schema(self) -> str:
+ | return "id INT"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val dataSource = createUserDefinedPythonDataSource(dataSourceName, dataSourceScript)
+ val err = intercept[AnalysisException](
+ dataSource(spark, provider = dataSourceName).collect())
+ assert(err.getErrorClass == "PYTHON_DATA_SOURCE_FAILED_TO_PLAN_IN_PYTHON")
+ assert(err.getMessage.contains("PYTHON_DATA_SOURCE_CREATE_ERROR"))
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org