Skip to content

Commit 6739e4f

Browse files
committed
[MINOR][PYTHON][TESTS] Update test_arrow_udf_output_timestamps_ltz to check the timezone conversion
### What changes were proposed in this pull request? Update `test_arrow_udf_output_timestamps_ltz` to check the timezone conversion ### Why are the changes needed? to make sure that arrow udf need to properly handle the timezone by itself ### Does this PR introduce _any_ user-facing change? no, test-only ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #52675 from zhengruifeng/refine_ts_ltz_test. Authored-by: Ruifeng Zheng <[email protected]> Signed-off-by: Ruifeng Zheng <[email protected]>
1 parent 8430dbf commit 6739e4f

File tree

1 file changed

+8
-5
lines changed

1 file changed

+8
-5
lines changed

python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -348,8 +348,11 @@ def extract_second(d):
348348
)
349349

350350
def test_arrow_udf_output_timestamps_ltz(self):
351+
from zoneinfo import ZoneInfo
351352
import pyarrow as pa
352353

354+
tz = self.spark.conf.get("spark.sql.session.timeZone")
355+
353356
df = self.spark.sql(
354357
"""
355358
SELECT * FROM VALUES
@@ -371,18 +374,18 @@ def build_ts(y, m, d, h, mi, s):
371374
int(h[i].as_py()),
372375
int(mi[i].as_py()),
373376
int(s[i].as_py()),
374-
tzinfo=datetime.timezone.utc,
375-
)
377+
tzinfo=ZoneInfo(tz),
378+
).astimezone(datetime.timezone.utc)
376379
for i in range(len(y))
377380
]
378381
return pa.array(dates, pa.timestamp("us", "UTC"))
379382

380383
result = df.select(build_ts("y", "m", "d", "h", "mi", "s").alias("ts"))
381384
self.assertEqual(
382385
[
383-
Row(ts=datetime.datetime(2022, 1, 5, 7, 0, 1)),
384-
Row(ts=datetime.datetime(2023, 2, 6, 8, 1, 2)),
385-
Row(ts=datetime.datetime(2024, 3, 7, 9, 2, 3)),
386+
Row(ts=datetime.datetime(2022, 1, 5, 15, 0, 1)),
387+
Row(ts=datetime.datetime(2023, 2, 6, 16, 1, 2)),
388+
Row(ts=datetime.datetime(2024, 3, 7, 17, 2, 3)),
386389
],
387390
result.collect(),
388391
)

0 commit comments

Comments
 (0)