fixup

Author: HyukjinKwon

python/pyspark/sql/tests/connect/test_connect_plan.py

@@ -333,6 +333,11 @@ def test_observe(self):
         from pyspark.sql.connect.observation import Observation
 
         class MockDF(DataFrame):
+            def __new__(cls, df: DataFrame) -> "DataFrame":
+                self = object.__new__(cls)
+                self.__init__(df)  # type: ignore[misc]
+                return self
+
             def __init__(self, df: DataFrame):
                 super().__init__(df._plan, df._session)
 

python/pyspark/testing/mlutils.py

@@ -99,6 +99,11 @@ def tearDownClass(cls):
 
 
 class MockDataset(DataFrame):
+    def __new__(cls) -> "DataFrame":
+        self = object.__new__(cls)
+        self.__init__()
+        return self
+
     def __init__(self):
         self.index = 0