fix test

JoshRosen · Apr 4, 2024 · ce7580f · ce7580f
1 parent 240923c
commit ce7580f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
@@ -589,15 +589,19 @@ def writeStream(self) -> DataStreamWriter:
 
         Examples
         --------
+        >>> import time
         >>> import tempfile
         >>> df = spark.readStream.format("rate").load()
         >>> type(df.writeStream)
         <class '...streaming.readwriter.DataStreamWriter'>
 
         >>> with tempfile.TemporaryDirectory(prefix="writeStream") as d:
         ...     # Create a table with Rate source.
-        ...     df.writeStream.toTable(
+        ...     query = df.writeStream.toTable(
         ...         "my_table", checkpointLocation=d)
+        ...     print(type(query))
+        ...     time.sleep(3)
+        ...     query.stop()
         <...streaming.query.StreamingQuery object at 0x...>
         """
         return DataStreamWriter(self)