jorisvandenbossche commented on code in PR #46529:
URL: https://github.com/apache/spark/pull/46529#discussion_r1610585401
##########
python/pyspark/sql/pandas/types.py:
##########
@@ -232,6 +312,124 @@ def _get_local_timezone() -> str:
return os.environ.get("TZ", "dateutil/:")
+def _check_arrow_array_timestamps_localize(
+ a: Union["pa.Array", "pa.ChunkedArray"],
+ dt: DataType,
+ truncate: bool = True,
+ timezone: Optional[str] = None,
+) -> Union["pa.Array", "pa.ChunkedArray"]:
+ """
+ Convert Arrow timestamps to timezone-naive in the specified timezone if
the specified Spark
+ data type is TimestampType, and optionally truncate nanosecond timestamps
to microseconds.
+
+ This function works on Arrow Arrays and ChunkedArrays, and it recurses to
convert nested
+ timestamps.
+
+ Parameters
+ ----------
+ a : :class:`pyarrow.Array` or :class:`pyarrow.ChunkedArray`
+ dt : :class:`DataType`
+ The Spark data type corresponding to the Arrow Array to be converted.
+ truncate : bool, default True
+ Whether to truncate nanosecond timestamps to microseconds. (default
``True``)
+ timezone : str, optional
+ The timezone to convert from. If there is a timestamp type, it's
required.
+
+ Returns
+ -------
+ :class:`pyarrow.Array` or :class:`pyarrow.ChunkedArray`
+ """
+ import pyarrow.types as types
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ if isinstance(a, pa.ChunkedArray) and (types.is_nested(a.type) or
types.is_dictionary(a.type)):
+ return pa.chunked_array(
+ [
+ _check_arrow_array_timestamps_localize(chunk, dt, truncate,
timezone)
+ for chunk in a.iterchunks()
+ ]
+ )
+
+ if types.is_timestamp(a.type) and truncate and a.type.unit == "ns":
+ a = pc.floor_temporal(a, unit="microsecond")
+
+ if types.is_timestamp(a.type) and a.type.tz is None and type(dt) ==
TimestampType:
+ assert timezone is not None
+
+ # Only localize timestamps that will become Spark TimestampType
columns.
+ # Do not localize timestamps that will become Spark TimestampNTZType
columns.
+ return pc.assume_timezone(a, timezone)
+ if types.is_list(a.type):
+ at: ArrayType = cast(ArrayType, dt)
+ return pa.ListArray.from_arrays(
+ a.offsets,
+ _check_arrow_array_timestamps_localize(a.values, at.elementType,
truncate, timezone),
+ )
Review Comment:
I don't know how much more complex you want to make this part of the code,
but one disadvantage with the current `mask` keyword is that it is not zero
copy (since it inverts the validity bitmap twice), so you might want to avoid
having to do this as much as possible:
- you could check if the list's `value_type` is either timestamp or nested,
and only in that case call `_check_arrow_array_timestamps_localize` on the
value, and otherwise just return `a` (we have `pa.types.is_nested` that could
help with that). That should already make any simple list of numeric type fully
zero-copy
- if you do have to recreated the ListArray, you can probably do
`mask=a.is_null() if a.null_count else None` to avoid allocating a full bitmap
in case there are no missing values
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
