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]
