-
Notifications
You must be signed in to change notification settings - Fork 29k
[SPARK-48497][PYTHON][DOCS] Add an example for Python data source writer in user guide #46833
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Closed
allisonwang-db
wants to merge
3
commits into
apache:master
from
allisonwang-db:spark-48497-ds-write-doc
Closed
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -73,6 +73,9 @@ Method that needs to be implemented for a capability: | |
| def reader(self, schema: StructType): | ||
| return FakeDataSourceReader(schema, self.options) | ||
|
|
||
| def writer(self, schema: StructType, overwrite: bool): | ||
| return FakeDataSourceWriter(self.options) | ||
|
|
||
| def streamReader(self, schema: StructType): | ||
| return FakeStreamReader(schema, self.options) | ||
|
|
||
|
|
@@ -83,8 +86,8 @@ Method that needs to be implemented for a capability: | |
| def streamWriter(self, schema: StructType, overwrite: bool): | ||
| return FakeStreamWriter(self.options) | ||
|
|
||
| Implementing Reader for Python Data Source | ||
| ------------------------------------------ | ||
| Implementing Batch Reader and Writer for Python Data Source | ||
| ----------------------------------------------------------- | ||
| **Implement the Reader** | ||
|
|
||
| Define the reader logic to generate synthetic data. Use the `faker` library to populate each field in the schema. | ||
|
|
@@ -109,6 +112,43 @@ Define the reader logic to generate synthetic data. Use the `faker` library to p | |
| row.append(value) | ||
| yield tuple(row) | ||
|
|
||
| **Implement the Writer** | ||
|
|
||
| Create a fake data source writer that processes each partition of data, counts the rows, and either | ||
| prints the total count of rows after a successful write or the number of failed tasks if the writing process fails. | ||
|
|
||
| .. code-block:: python | ||
|
|
||
| from dataclasses import dataclass | ||
| from typing import Iterator, List | ||
|
|
||
| from pyspark.sql.types import Row | ||
| from pyspark.sql.datasource import DataSource, DataSourceWriter, WriterCommitMessage | ||
|
|
||
| @dataclass | ||
| class SimpleCommitMessage(WriterCommitMessage): | ||
| partition_id: int | ||
| count: int | ||
|
|
||
| class FakeDataSourceWriter(DataSourceWriter): | ||
|
|
||
| def write(self, rows: Iterator[Row]) -> SimpleCommitMessage: | ||
| from pyspark import TaskContext | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we can import this on the top together.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This import actually needs to be inside the write method otherwise it will throw a serialization error. |
||
|
|
||
| context = TaskContext.get() | ||
| partition_id = context.partitionId() | ||
| cnt = sum(1 for _ in rows) | ||
| return SimpleCommitMessage(partition_id=partition_id, count=cnt) | ||
|
|
||
| def commit(self, messages: List[SimpleCommitMessage]) -> None: | ||
| total_count = sum(message.count for message in messages) | ||
| print(f"Total number of rows: {total_count}") | ||
|
|
||
| def abort(self, messages: List[SimpleCommitMessage]) -> None: | ||
| failed_count = sum(message is None for message in messages) | ||
| print(f"Number of failed tasks: {failed_count}") | ||
|
|
||
|
|
||
| Implementing Streaming Reader and Writer for Python Data Source | ||
| --------------------------------------------------------------- | ||
| **Implement the Stream Reader** | ||
|
|
@@ -267,7 +307,9 @@ After defining your data source, it must be registered before usage. | |
|
|
||
| spark.dataSource.register(FakeDataSource) | ||
|
|
||
| Use the fake datasource with the default schema and options: | ||
| **Read From a Python Data Source** | ||
|
|
||
| Read from the fake datasource with the default schema and options: | ||
|
|
||
| .. code-block:: python | ||
|
|
||
|
|
@@ -281,7 +323,7 @@ Use the fake datasource with the default schema and options: | |
| # | Amy Martin|1988-10-28| 68076| Oregon| | ||
| # +-----------+----------+-------+-------+ | ||
|
|
||
| Use the fake datasource with a custom schema: | ||
| Read from the fake datasource with a custom schema: | ||
|
|
||
| .. code-block:: python | ||
|
|
||
|
|
@@ -295,7 +337,7 @@ Use the fake datasource with a custom schema: | |
| # |Mrs. Jacqueline Brown|Maynard Inc | | ||
| # +---------------------+--------------+ | ||
|
|
||
| Use the fake datasource with a different number of rows: | ||
| Read from the fake datasource with a different number of rows: | ||
|
|
||
| .. code-block:: python | ||
|
|
||
|
|
@@ -311,6 +353,18 @@ Use the fake datasource with a different number of rows: | |
| # | Douglas James|2007-01-18| 46226| Alabama| | ||
| # +--------------+----------+-------+------------+ | ||
|
|
||
| **Write To a Python Data Source** | ||
|
|
||
| To write data to a custom location, make sure that you specify the `mode()` clause. Supported modes are `append` and `overwrite`. | ||
|
|
||
| .. code-block:: python | ||
|
|
||
| df = spark.range(0, 10, 1, 5) | ||
| df.write.format("fake").mode("append").save() | ||
|
|
||
| # You can check the Spark log (standard error) to see the output of the write operation. | ||
| # Total number of rows: 10 | ||
|
|
||
| **Use a Python Data Source in Streaming Query** | ||
|
|
||
| Once we register the python data source, we can also use it in streaming queries as source of readStream() or sink of writeStream() by passing short name or full name to format(). | ||
|
|
@@ -319,7 +373,7 @@ Start a query that read from fake python data source and write to console | |
|
|
||
| .. code-block:: python | ||
|
|
||
| query = spark.readStream.format("fake").load().writeStream().format("console").start() | ||
| query = spark.readStream.format("fake").load().writeStream.format("console").start() | ||
|
|
||
| # +---+ | ||
| # | id| | ||
|
|
@@ -338,4 +392,4 @@ We can also use the same data source in streaming reader and writer | |
|
|
||
| .. code-block:: python | ||
|
|
||
| query = spark.readStream.format("fake").load().writeStream().format("fake").start("/output_path") | ||
| query = spark.readStream.format("fake").load().writeStream.format("fake").start("/output_path") | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
per PEP 8