Skip to content

Commit e3d8c2f

Browse files
committed
core: allow timezone to be explicitely selected for export_partitioned
1 parent 2b3bb44 commit e3d8c2f

3 files changed

Lines changed: 70 additions & 13 deletions

File tree

‎src/damast/core/dataframe.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,9 @@ def export_partitioned(
289289

290290
key_col = "__damast_partition_key__"
291291
collected = self.lazyframe.with_columns(strategy.key_expr().alias(key_col)).collect()
292-
warn_on_local_time_buckets(collected[key_col])
292+
if strategy.time_zone is None:
293+
# only an accidental local-time bucketing is worth a warning
294+
warn_on_local_time_buckets(collected[key_col])
293295

294296
written: list[Path] = []
295297
# as_dict=False + re-deriving the key from the partition itself (rather than

‎src/damast/core/partitioning.py‎

Lines changed: 34 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -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

‎tests/damast/core/test_partitioning.py‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -375,8 +375,8 @@ def test_expected_paths_names_do_not_match_a_non_utc_archive(tmp_path):
375375
written = _zoned_adf("Asia/Tokyo").export_partitioned(tmp_path, ByTime("timestamp", every="1d"))
376376
predicted = SaveAs.expected_paths(
377377
"time:timestamp+daily:%Y-%m-%d",
378-
start=dt.datetime(2026, 1, 1, tzinfo=dt.UTC),
379-
end=dt.datetime(2026, 1, 1, 23, 59, tzinfo=dt.UTC),
378+
start=dt.datetime(2026, 1, 1, tzinfo=dt.timezone.utc),
379+
end=dt.datetime(2026, 1, 1, 23, 59, tzinfo=dt.timezone.utc),
380380
)
381381

382382
assert [p.name for p in predicted] == ["2026-01-01.parquet"]
@@ -432,3 +432,34 @@ def test_export_partitioned_suffix_is_configurable(timeseries_adf, tmp_path):
432432
# the default is unchanged, and the files stay readable whatever they are called
433433
written = timeseries_adf.export_partitioned(tmp_path / "default", ByColumn("mmsi"))
434434
assert sorted(p.name for p in written) == ["mmsi_1.parquet", "mmsi_2.parquet", "mmsi_3.parquet"]
435+
436+
437+
def test_by_time_time_zone_converts_before_truncating(tmp_path, caplog):
438+
"""`time_zone` decides which day an instant belongs to, whatever the column carries."""
439+
tokyo = _zoned_adf("Asia/Tokyo")
440+
441+
# the column's own zone (default): Tokyo days
442+
assert [p.name for p in tokyo.export_partitioned(tmp_path / "own", ByTime("timestamp", every="1d"))] == [
443+
"2026-01-02.parquet"
444+
]
445+
# converted to UTC: the same instants are a UTC 1 Jan
446+
caplog.clear() # the default export above warned - that is its own test
447+
with caplog.at_level("WARNING", logger="damast.core.partitioning"):
448+
written = tokyo.export_partitioned(tmp_path / "utc", ByTime("timestamp", every="1d", time_zone="UTC"))
449+
assert [p.name for p in written] == ["2026-01-01.parquet"]
450+
assert caplog.records == []
451+
452+
# a naive column is read as UTC, so converting it to UTC changes nothing
453+
assert [p.name for p in _zoned_adf(None).export_partitioned(
454+
tmp_path / "naive", ByTime("timestamp", every="1d", time_zone="UTC"))] == ["2026-01-01.parquet"]
455+
456+
457+
def test_by_time_explicit_local_time_zone_is_not_warned_about(tmp_path, caplog):
458+
"""Deliberately bucketing by local days is legitimate - only an accidental one warns."""
459+
with caplog.at_level("WARNING", logger="damast.core.partitioning"):
460+
written = _zoned_adf("UTC").export_partitioned(
461+
tmp_path, ByTime("timestamp", every="1d", time_zone="Asia/Tokyo")
462+
)
463+
464+
assert [p.name for p in written] == ["2026-01-02.parquet"]
465+
assert caplog.records == []

0 commit comments

Comments
 (0)