This is an automated email from the ASF dual-hosted git repository.

gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 15e99cf676d9 [MINOR][CONNECT][PYTHON] Fix minor test issue and 
properly yield extension from result
15e99cf676d9 is described below

commit 15e99cf676d9de02c54ca5ebe9a2bc6a3ce014e5
Author: Martin Grund <[email protected]>
AuthorDate: Fri Dec 15 08:55:55 2023 -0800

    [MINOR][CONNECT][PYTHON] Fix minor test issue and properly yield extension 
from result
    
    ### What changes were proposed in this pull request?
    This patch is a very minor cleanup to fix an issue where a testcase was not 
properly calling it's superclass and we missed properly yielding the right type 
in the response processing.
    
    ### Why are the changes needed?
    Stability
    
    ### Does this PR introduce _any_ user-facing change?
    No
    
    ### How was this patch tested?
    Existing coverage.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    No
    
    Closes #44372 from grundprinzip/minor_cleanup.
    
    Authored-by: Martin Grund <[email protected]>
    Signed-off-by: Hyukjin Kwon <[email protected]>
---
 python/pyspark/ml/tests/connect/test_connect_function.py | 2 +-
 python/pyspark/sql/connect/client/core.py                | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/python/pyspark/ml/tests/connect/test_connect_function.py 
b/python/pyspark/ml/tests/connect/test_connect_function.py
index 2b2cd3bc3a89..b38d415e2bb2 100644
--- a/python/pyspark/ml/tests/connect/test_connect_function.py
+++ b/python/pyspark/ml/tests/connect/test_connect_function.py
@@ -49,7 +49,7 @@ class SparkConnectMLFunctionTests(ReusedConnectTestCase, 
PandasOnSparkTestUtils,
     @classmethod
     def tearDownClass(cls):
         cls.spark = cls.connect  # Stopping Spark Connect closes the session 
in JVM at the server.
-        super(SparkConnectMLFunctionTests, cls).setUpClass()
+        super(SparkConnectMLFunctionTests, cls).tearDownClass()
         del os.environ["PYSPARK_NO_NAMESPACE_SHARE"]
 
     def compare_by_show(self, df1, df2, n: int = 20, truncate: int = 20):
diff --git a/python/pyspark/sql/connect/client/core.py 
b/python/pyspark/sql/connect/client/core.py
index 0b502494f781..85b2b98fef13 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -51,7 +51,7 @@ import pyarrow as pa
 import google.protobuf.message
 from grpc_status import rpc_status
 import grpc
-from google.protobuf import text_format
+from google.protobuf import text_format, any_pb2
 from google.rpc import error_details_pb2
 
 from pyspark.loose_version import LooseVersion
@@ -1152,6 +1152,7 @@ class SparkConnectClient(object):
                 PlanMetrics,
                 PlanObservedMetrics,
                 Dict[str, Any],
+                any_pb2.Any,
             ]
         ]:
             nonlocal num_records
@@ -1198,6 +1199,8 @@ class SparkConnectClient(object):
                     addresses = [address for address in resource.addresses]
                     resources[key] = ResourceInformation(name, addresses)
                 yield {"get_resources_command_result": resources}
+            if b.HasField("extension"):
+                yield b.extension
             if b.HasField("arrow_batch"):
                 logger.debug(
                     f"Received arrow batch rows={b.arrow_batch.row_count} "


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to