[SPARK-54247][PYTHON][FOLLOWUP] Ensure socket is not leaked and add regression tests

### What changes were proposed in this pull request?

* #53200 did not fix the socket leak issue properly. This PR fixed it by closing the socket as well
* In order to prevent such mistakes, regression tests are added
* `_local_iterator_from_socket` is leaking socket too, fixed in this PR as well.

### Why are the changes needed?

To avoid leaking socket (technically it's still closed at GC, it was explained in #53200).

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

A new regression test is added and I confirmed the test fails before fix and passes after.

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #53203 from gaogaotiantian/really-fix-socket-leak.

Authored-by: Tian Gao <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>

Author: gaogaotiantian

python/pyspark/sql/tests/test_dataframe.py

@@ -20,6 +20,7 @@
 import pydoc
 import shutil
 import tempfile
+import warnings
 import unittest
 from typing import cast
 import io
@@ -1044,6 +1045,18 @@ def test_local_checkpoint_dataframe_with_storage_level(self):
         df = self.spark.range(10).localCheckpoint(eager=True, storageLevel=StorageLevel.DISK_ONLY)
         df.collect()
 
+    def test_socket_leak(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always", ResourceWarning)
+            df = self.spark.range(10)
+            df.collect()
+
+            df = self.spark.range(10)
+            for _ in df.toLocalIterator():
+                pass
+
+        self.assertEqual(w, [])
+
     def test_transpose(self):
         df = self.spark.createDataFrame([{"a": "x", "b": "y", "c": "z"}])
 

python/pyspark/util.py

@@ -668,7 +668,7 @@ class PythonEvalType:
     SQL_ARROW_UDTF: "SQLArrowUDTFType" = 302
 
 
-def _create_local_socket(sock_info: "JavaArray") -> "io.BufferedRWPair":
+def _create_local_socket(sock_info: "JavaArray") -> Tuple["io.BufferedRWPair", "socket.socket"]:
     """
     Create a local socket that can be used to load deserialized data from the JVM
 
@@ -689,7 +689,7 @@ def _create_local_socket(sock_info: "JavaArray") -> "io.BufferedRWPair":
     # The RDD materialization time is unpredictable, if we set a timeout for socket reading
     # operation, it will very possibly fail. See SPARK-18281.
     sock.settimeout(None)
-    return sockfile
+    return sockfile, sock
 
 
 @contextmanager
@@ -710,10 +710,11 @@ def _load_from_socket(sock_info: "JavaArray", serializer: "Serializer") -> Itera
     usually a generator that yields deserialized data
     """
     try:
-        sockfile = _create_local_socket(sock_info)
+        sockfile, sock = _create_local_socket(sock_info)
         yield serializer.load_stream(sockfile)
     finally:
         sockfile.close()
+        sock.close()
 
 
 def _local_iterator_from_socket(sock_info: "JavaArray", serializer: "Serializer") -> Iterator[Any]:
@@ -725,7 +726,7 @@ def __init__(self, _sock_info: "JavaArray", _serializer: "Serializer"):
             auth_secret: str
             jsocket_auth_server: "JavaObject"
             port, auth_secret, self.jsocket_auth_server = _sock_info
-            self._sockfile = _create_local_socket((port, auth_secret))
+            self._sockfile, self._sock = _create_local_socket((port, auth_secret))
             self._serializer = _serializer
             self._read_iter: Iterator[Any] = iter([])  # Initialize as empty iterator
             self._read_status = 1
@@ -761,6 +762,8 @@ def __del__(self) -> None:
                 except Exception:
                     # Ignore any errors, socket is automatically closed when garbage-collected
                     pass
+            self._sockfile.close()
+            self._sock.close()
 
     return iter(PyLocalIterable(sock_info, serializer))