@@ -31,6 +31,11 @@ class PartitionStrategy(ABC):
3131 :meth:`AnnotatedDataFrame.export_partitioned`.
3232 """
3333
34+ #: Time zone this strategy deliberately buckets in, if any. Setting it states that a
35+ #: non-UTC bucketing is intended and suppresses the warning from
36+ #: :func:`warn_on_local_time_buckets`.
37+ time_zone : str | None = None
38+
3439 @abstractmethod
3540 def key_expr (self ) -> polars .Expr :
3641 """A polars expression computing this strategy's per-row partition key."""
@@ -72,24 +77,38 @@ class ByTime(PartitionStrategy):
7277 for one file per hour, ``"1d"`` for one file per day, ``"1w"`` per week.
7378
7479 .. note::
75- Truncation happens in the timestamp column's own time zone, so a column in e.g.
76- ``Asia/Tokyo`` yields Tokyo-day buckets, not UTC-day buckets - the same instant then
77- lands in a differently named file. :meth:`AnnotatedDataFrame.export_partitioned` warns
78- about this; convert the column with ``dt.convert_time_zone("UTC")`` for UTC buckets.
79- A naive (zone-less) column is truncated as-is.
80+ By default truncation happens in the timestamp column's own time zone, so a column in
81+ e.g. ``Asia/Tokyo`` yields Tokyo-day buckets, not UTC-day buckets - the same instant
82+ then lands in a differently named file, and
83+ :meth:`AnnotatedDataFrame.export_partitioned` warns about it. Pass ``time_zone="UTC"``
84+ to bucket by UTC regardless of how the column is zoned.
85+
86+ Example:
87+
88+ .. code-block:: python
89+
90+ # one file per UTC day, whatever zone the column carries
91+ ByTime("timestamp", every="1d", time_zone="UTC")
8092 """
8193
8294 def __init__ (
8395 self ,
8496 timestamp_column : str ,
8597 every : str ,
8698 * ,
99+ time_zone : str | None = None ,
87100 format : str | None = None ,
88101 prefix : str | None = None ,
89102 ):
90103 """
91104 :param timestamp_column: Name of the datetime column to partition by
92105 :param every: Truncation interval, e.g. ``"1h"``, ``"1d"``, ``"1w"``
106+ :param time_zone: Convert the column to this time zone before truncating, e.g.
107+ ``"UTC"`` for UTC buckets or ``"Europe/Oslo"`` to bucket by local days on purpose.
108+ The default (``None``) truncates in whatever zone the column carries, which is
109+ polars' own behaviour. A naive (zone-less) column is read as UTC, so
110+ ``time_zone="UTC"`` leaves it unchanged. Setting this states the intent, so the
111+ local-time warning of :meth:`AnnotatedDataFrame.export_partitioned` is suppressed.
93112 :param format: Optional ``strftime`` format for the filename stem (default:
94113 automatically pick the coarsest representation matching the truncated value, e.g.
95114 ``"2026-09-14"`` for a day, ``"2026-09-14T13"`` for an hour). Use e.g.
@@ -100,11 +119,15 @@ def __init__(
100119 """
101120 self .timestamp_column = timestamp_column
102121 self .every = every
122+ self .time_zone = time_zone
103123 self .format = format
104124 self .prefix = prefix
105125
106126 def key_expr (self ) -> polars .Expr :
107- return polars .col (self .timestamp_column ).dt .truncate (self .every )
127+ column = polars .col (self .timestamp_column )
128+ if self .time_zone is not None :
129+ column = column .dt .convert_time_zone (self .time_zone )
130+ return column .dt .truncate (self .every )
108131
109132 def filename (self , key : Any ) -> str :
110133 stem = (
@@ -170,8 +193,8 @@ def warn_on_local_time_buckets(keys: polars.Series) -> None:
170193 column carrying e.g. ``Asia/Tokyo`` produces local-day buckets - the same instant lands
171194 in a different file than it would in UTC, and the resulting filenames cannot be read
172195 back without knowing that zone. Partitioning by local days is legitimate, so this only
173- warns; convert the column with ``dt.convert_time_zone( "UTC")`` beforehand to get UTC
174- buckets .
196+ warns; pass ``ByTime(..., time_zone= "UTC")`` for UTC buckets, or name the zone explicitly
197+ (:attr:`PartitionStrategy.time_zone`) to state the intent and silence this .
175198
176199 :param keys: The materialized partition key column, as produced by
177200 :meth:`PartitionStrategy.key_expr`
@@ -186,8 +209,9 @@ def warn_on_local_time_buckets(keys: polars.Series) -> None:
186209 if time_zone is not None and time_zone != "UTC" :
187210 logger .warning (
188211 f"Partition key '{ name } ' is a datetime in time zone '{ time_zone } ', so its buckets"
189- f" (and the resulting filenames) follow { time_zone } days, not UTC days. Convert the"
190- f" column with dt.convert_time_zone('UTC') first if UTC buckets were intended."
212+ f" (and the resulting filenames) follow { time_zone } days, not UTC days. Pass"
213+ f" time_zone='UTC' if UTC buckets were intended, or time_zone='{ time_zone } ' to"
214+ f" state that they were not."
191215 )
192216
193217
0 commit comments