zhengruifeng commented on code in PR #52321:
URL: https://github.com/apache/spark/pull/52321#discussion_r2377491457
##########
python/pyspark/sql/tests/test_observation.py:
##########
@@ -190,10 +186,46 @@ def test_observe_on_commands(self):
]:
with self.subTest(command=command):
observation = Observation()
- observed_df = df.observe(observation,
count(lit(1)).alias("cnt"))
+ observed_df = df.observe(observation,
F.count(F.lit(1)).alias("cnt"))
action(observed_df)
self.assertEqual(observation.get, dict(cnt=50))
+ def test_observe_with_struct_type(self):
+ observation = Observation("struct")
+
+ df = self.spark.range(10).observe(
+ observation,
+ F.struct(F.count(F.lit(1)).alias("rows"),
F.max("id").alias("maxid")).alias("struct"),
+ )
+
+ assertDataFrameEqual(df, [Row(id=id) for id in range(10)])
+
+ self.assertEqual(observation.get, {"struct": Row(rows=10, maxid=9)})
+
+ def test_observe_with_array_type(self):
+ observation = Observation("array")
+
+ df = self.spark.range(10).observe(
+ observation,
+ F.array(F.count(F.lit(1))).alias("array"),
+ )
+
+ assertDataFrameEqual(df, [Row(id=id) for id in range(10)])
+
+ self.assertEqual(observation.get, {"array": [10]})
+
+ def test_observe_with_map_type(self):
Review Comment:
@ueshin this test occasionally fails (e.g.
https://github.com/apache/spark/actions/runs/17988107433/job/51171506921)
```
======================================================================
FAIL [0.159s]: test_observe_with_map_type
(pyspark.sql.tests.connect.test_parity_observation.DataFrameObservationParityTests.test_observe_with_map_type)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/__w/spark/spark/python/pyspark/sql/tests/test_observation.py", line
227, in test_observe_with_map_type
self.assertEqual(observation.get, {"map": {"count": 10}})
~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: {'map': {'': 10}} != {'map': {'count': 10}}
- {'map': {'': 10}}
+ {'map': {'count': 10}}
? +++++
----------------------------------------------------------------------
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
