Skip to content

[SPARK-43031] [SS] [Connect] Enable unit test and doctest for streaming #40691

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.

Sign up for GitHub

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
WweiL wants to merge 13 commits into from
Closed

[SPARK-43031] [SS] [Connect] Enable unit test and doctest for streaming #40691

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
4d0fcdd
done
WweiL Apr 6, 2023
0ae7e33
add versionchanged to query and readwriter
WweiL Apr 7, 2023
b14fe81
fix conflict, put TODO to correct position
WweiL Apr 7, 2023
17720b7
style
WweiL Apr 10, 2023
1e68a3c
comments
WweiL Apr 10, 2023
dc05be8
address comments, add a new foreachBatch test class, remove all ELLIP…
WweiL Apr 10, 2023
c1674eb
minor
WweiL Apr 10, 2023
23b9c93
minor
WweiL Apr 10, 2023
60ddd01
lint
WweiL Apr 11, 2023
e576821
remove empty line
WweiL Apr 11, 2023
26e2488
remove several docs in connect readwriter.py and query.py to pass doc…
WweiL Apr 12, 2023
e25f7e6
minor, add back doc tests in module.py
WweiL Apr 12, 2023
aa1d4c2
style
WweiL Apr 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions dev/sparktestsupport/modules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -495,6 +495,8 @@ def __hash__(self):
"pyspark.sql.tests.test_serde",
"pyspark.sql.tests.test_session",
"pyspark.sql.tests.streaming.test_streaming",
"pyspark.sql.tests.streaming.test_streaming_foreach",
"pyspark.sql.tests.streaming.test_streaming_foreachBatch",
"pyspark.sql.tests.streaming.test_streaming_listener",
"pyspark.sql.tests.test_types",
"pyspark.sql.tests.test_udf",
Expand Down Expand Up @@ -749,6 +751,8 @@ def __hash__(self):
"pyspark.sql.connect.dataframe",
"pyspark.sql.connect.functions",
"pyspark.sql.connect.avro.functions",
"pyspark.sql.connect.streaming.readwriter",
"pyspark.sql.connect.streaming.query",
# sql unittests
"pyspark.sql.tests.connect.test_client",
"pyspark.sql.tests.connect.test_connect_plan",
Expand All @@ -773,6 +777,7 @@ def __hash__(self):
"pyspark.sql.tests.connect.test_parity_arrow_map",
"pyspark.sql.tests.connect.test_parity_pandas_grouped_map",
"pyspark.sql.tests.connect.test_parity_pandas_cogrouped_map",
"pyspark.sql.tests.connect.streaming.test_parity_streaming",
# ml doctests
"pyspark.ml.connect.functions",
# ml unittests
Expand Down
28 changes: 25 additions & 3 deletions python/pyspark/sql/connect/streaming/query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#

import json
import sys
from typing import TYPE_CHECKING, Any, cast, Dict, List, Optional

from pyspark.errors import StreamingQueryException
Expand Down Expand Up @@ -149,10 +150,31 @@ def _execute_streaming_query_cmd(


def _test() -> None:
# TODO(SPARK-43031): port _test() from legacy query.py.
pass
import doctest
import os
from pyspark.sql import SparkSession as PySparkSession
import pyspark.sql.connect.streaming.query

os.chdir(os.environ["SPARK_HOME"])

globs = pyspark.sql.connect.streaming.query.__dict__.copy()

globs["spark"] = (
PySparkSession.builder.appName("sql.connect.streaming.query tests")
.remote("local[4]")
.getOrCreate()
)

(failure_count, test_count) = doctest.testmod(
pyspark.sql.connect.streaming.query,
globs=globs,
optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF,
)
globs["spark"].stop()

if failure_count:
sys.exit(-1)


if __name__ == "__main__":
# TODO(SPARK-43031): Add this file dev/sparktestsupport/modules.py to enable testing in CI.
_test()
102 changes: 95 additions & 7 deletions python/pyspark/sql/connect/streaming/readwriter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,9 +168,75 @@ def json(

@WweiL WweiL Apr 6, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please ignore the change in this file. They are added and to be reviewed in #40689.

json.__doc__ = PySparkDataStreamReader.json.__doc__

# def orc() TODO
# def parquet() TODO
# def text() TODO
def orc(
self,
path: str,
mergeSchema: Optional[bool] = None,
pathGlobFilter: Optional[Union[bool, str]] = None,
recursiveFileLookup: Optional[Union[bool, str]] = None,
) -> "DataFrame":
self._set_opts(
mergeSchema=mergeSchema,
pathGlobFilter=pathGlobFilter,
recursiveFileLookup=recursiveFileLookup,
)
if isinstance(path, str):
return self.load(path=path, format="orc")
else:
raise TypeError("path can be only a single string")

orc.__doc__ = PySparkDataStreamReader.orc.__doc__

def parquet(
self,
path: str,
mergeSchema: Optional[bool] = None,
pathGlobFilter: Optional[Union[bool, str]] = None,
recursiveFileLookup: Optional[Union[bool, str]] = None,
datetimeRebaseMode: Optional[Union[bool, str]] = None,
int96RebaseMode: Optional[Union[bool, str]] = None,
) -> "DataFrame":
self._set_opts(
mergeSchema=mergeSchema,
pathGlobFilter=pathGlobFilter,
recursiveFileLookup=recursiveFileLookup,
datetimeRebaseMode=datetimeRebaseMode,
int96RebaseMode=int96RebaseMode,
)
self._set_opts(
mergeSchema=mergeSchema,
pathGlobFilter=pathGlobFilter,
recursiveFileLookup=recursiveFileLookup,
datetimeRebaseMode=datetimeRebaseMode,
int96RebaseMode=int96RebaseMode,
)
if isinstance(path, str):
return self.load(path=path, format="parquet")
else:
raise TypeError("path can be only a single string")

parquet.__doc__ = PySparkDataStreamReader.parquet.__doc__

def text(
self,
path: str,
wholetext: bool = False,
lineSep: Optional[str] = None,
pathGlobFilter: Optional[Union[bool, str]] = None,
recursiveFileLookup: Optional[Union[bool, str]] = None,
) -> "DataFrame":
self._set_opts(
wholetext=wholetext,
lineSep=lineSep,
pathGlobFilter=pathGlobFilter,
recursiveFileLookup=recursiveFileLookup,
)
if isinstance(path, str):
return self.load(path=path, format="text")
else:
raise TypeError("path can be only a single string")

text.__doc__ = PySparkDataStreamReader.text.__doc__

def csv(
self,
Expand Down Expand Up @@ -245,7 +311,7 @@ def csv(

csv.__doc__ = PySparkDataStreamReader.csv.__doc__

# def table() TODO. Use Read(table_name) relation.
# def table() TODO(SPARK-43042). Use Read(table_name) relation.


DataStreamReader.__doc__ = PySparkDataStreamReader.__doc__
Expand Down Expand Up @@ -460,10 +526,32 @@ def toTable(


def _test() -> None:
# TODO(SPARK-43031): port _test() from legacy query.py.
pass
import sys
import doctest
from pyspark.sql import SparkSession as PySparkSession
import pyspark.sql.connect.streaming.readwriter

globs = pyspark.sql.connect.readwriter.__dict__.copy()

globs["spark"] = (
PySparkSession.builder.appName("sql.connect.streaming.readwriter tests")
.remote("local[4]")
.getOrCreate()
)

(failure_count, test_count) = doctest.testmod(
pyspark.sql.connect.streaming.readwriter,
globs=globs,
optionflags=doctest.ELLIPSIS
| doctest.NORMALIZE_WHITESPACE
| doctest.IGNORE_EXCEPTION_DETAIL,
)

globs["spark"].stop()

if failure_count:
sys.exit(-1)


if __name__ == "__main__":
# TODO(SPARK-43031): Add this file dev/sparktestsupport/modules.py to enable testing in CI.
_test()
2 changes: 1 addition & 1 deletion python/pyspark/sql/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -535,7 +535,7 @@ def writeStream(self) -> DataStreamWriter:
>>> with tempfile.TemporaryDirectory() as d:
... # Create a table with Rate source.
... df.writeStream.toTable(
... "my_table", checkpointLocation=d) # doctest: +ELLIPSIS

@pengzhon-db pengzhon-db Apr 11, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

curious what does # doctest: +ELLIPSIS mean?

@WweiL WweiL Apr 11, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My understanding is that in doctest the result will be kind of regex checked if this flag is set. Like the line below right now is <...streaming.query.StreamingQuery object at 0x...>, but before it was <pyspark.sql.streaming.query.StreamingQuery object at 0x...>, which would conflict with connect's test, which returns <pyspark.sql.connect.streaming.query.StreamingQuery object at 0x...>

So to make this test works for both connect and non-connect, we enable this regex-like check and replace below with ...
But it doesn't matter anyway, as we enabled the flag in test options in the __main__ method below

... "my_table", checkpointLocation=d)
<...streaming.query.StreamingQuery object at 0x...>
"""
return DataStreamWriter(self)
Expand Down
20 changes: 13 additions & 7 deletions python/pyspark/sql/streaming/query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
#

import json
import sys
from typing import Any, Dict, List, Optional

from py4j.java_gateway import JavaObject, java_import
Expand All @@ -37,6 +36,9 @@ class StreamingQuery:

.. versionadded:: 2.0.0

.. versionchanged:: 3.5.0
Supports Spark Connect.

Notes
-----
This API is evolving.
Expand Down Expand Up @@ -68,7 +70,7 @@ def id(self) -> str:

Get the unique id of this query that persists across restarts from checkpoint data

>>> sq.id # doctest: +ELLIPSIS
>>> sq.id
'...'

>>> sq.stop()
Expand All @@ -95,7 +97,7 @@ def runId(self) -> str:

Get the unique id of this query that does not persist across restarts

>>> sq.runId # doctest: +ELLIPSIS
>>> sq.runId
'...'

>>> sq.stop()
Expand Down Expand Up @@ -154,6 +156,7 @@ def isActive(self) -> bool:
"""
return self._jsq.isActive()

# TODO(SPARK-42960): remove the doctest: +SKIP flag below
def awaitTermination(self, timeout: Optional[int] = None) -> Optional[bool]:
"""
Waits for the termination of `this` query, either by :func:`query.stop()` or by an
Expand Down Expand Up @@ -188,7 +191,7 @@ def awaitTermination(self, timeout: Optional[int] = None) -> Optional[bool]:

Return whether the query has terminated or not within 5 seconds

>>> sq.awaitTermination(5)
>>> sq.awaitTermination(5) # doctest: +SKIP

@rangadi rangadi Apr 10, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are these needed? This won't be tested in connect yet, right?

@WweiL WweiL Apr 10, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right.. This actually silents this line in the doctest for both connect and non-connect I believe (@HyukjinKwon please correct me if I'm wrong) , just some temporary pain we have to bear with as connect doesn't support this yet

@rangadi rangadi Apr 10, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the pain? What fails if you remove this diff?

@WweiL WweiL Apr 10, 2023 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing fails but we skip some tests for also non-connect scenario

@rangadi rangadi Apr 10, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lets not skip any tests. Any skip should have a TODO comment.

@WweiL WweiL Apr 10, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure! I've moved the TODO to the top of these methods to avoid them showing up in the PySpark doc

@HyukjinKwon HyukjinKwon Apr 10, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we're absolutely sure that we will enable this back very soon, I am fine. As @rangadi said, I tried my best to avoid skipping the tests, but I did few times when I am about to enable it back very soon.

@WweiL WweiL Apr 12, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see. I think I'll just remove all of the SKIP flags and disable the doctest for now. Some of the functionalities might not be supported that soon I think.

@rangadi rangadi Apr 12, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Better solution is to comment out setting __doc__ for awaitTermination() in connect's query.py.

@WweiL WweiL Apr 12, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Done

False

>>> sq.stop()
Expand Down Expand Up @@ -219,7 +222,7 @@ def status(self) -> Dict[str, Any]:

Get the current status of the query

>>> sq.status # doctest: +ELLIPSIS
>>> sq.status
{'message': '...', 'isDataAvailable': ..., 'isTriggerActive': ...}

>>> sq.stop()
Expand Down Expand Up @@ -248,7 +251,7 @@ def recentProgress(self) -> List[Dict[str, Any]]:

Get an array of the most recent query progress updates for this query

>>> sq.recentProgress # doctest: +ELLIPSIS
>>> sq.recentProgress
[...]

>>> sq.stop()
Expand Down Expand Up @@ -314,6 +317,7 @@ def processAllAvailable(self) -> None:
"""
return self._jsq.processAllAvailable()

# TODO(SPARK-42940): remove the doctest: +SKIP flag below
def stop(self) -> None:
"""
Stop this streaming query.
Expand All @@ -330,7 +334,8 @@ def stop(self) -> None:
Stop streaming query

>>> sq.stop()
>>> sq.isActive

>>> sq.isActive # doctest: +SKIP

@rangadi rangadi Apr 10, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this required? Why is it skipped?

@WweiL WweiL Apr 10, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Calling isActive after the query is stopped will throw error right now before the better session management is implemented

@rangadi rangadi Apr 10, 2023

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is see, please add a TODO here with a brief comment.

False
"""
self._jsq.stop()
Expand Down Expand Up @@ -632,6 +637,7 @@ def removeListener(self, listener: StreamingQueryListener) -> None:
def _test() -> None:
import doctest
import os
import sys
from pyspark.sql import SparkSession
import pyspark.sql.streaming.query
from py4j.protocol import Py4JError
Expand Down
Loading