Allow Column type for timezone argument in pyspark.sql.functions

[charlietsai]

In the functions here:

pyspark-stubs/third_party/3/pyspark/sql/functions.pyi

Lines 100 to 101 in 3c4684a

	def from_utc_timestamp(timestamp: ColumnOrName, tz: str) -> Column: ...
	def to_utc_timestamp(timestamp: ColumnOrName, tz: str) -> Column: ...

we currently have tz: str but this can also be specified as a Column

Example:

>>> from pyspark.sql import functions
>>> df = spark.sql("SELECT CAST(0 AS TIMESTAMP) AS timestamp, 'Asia/Tokyo' AS tz")
>>> df.select(functions.from_utc_timestamp(df.timestamp, df.tz)).collect()
[Row(from_utc_timestamp(timestamp, tz)=datetime.datetime(1970, 1, 1, 18, 0))]

I think this could be expanded to tz: ColumnOrName?

[zero323]

Thanks for reporting @charlietsai. Care to open a PR for this?

I think this could be expanded to tz: ColumnOrName?

Indeed, although in this particular case I'd prefer to have explicit Union[Column, str] to stress out that meaning is different from ColumnOrName. It won't make any difference for the type checker, but can make a big one for anyone who reads the annotation.

[charlietsai]

Makes sense to me! Opening a PR

[zero323]

Thank you @charlietsai. Merged into master and backported to branch-3.0 and branch-2.4

[charlietsai]

Thanks @zero323 for this work!