diff --git a/sdm/ensemble.py b/sdm/ensemble.py index 83fbbe14e9..6de66fdec5 100644 --- a/sdm/ensemble.py +++ b/sdm/ensemble.py @@ -171,6 +171,7 @@ def _pack_tables( # must match for torch.stack to preserve member semantics. compatibility_key = ( table.column_groups, + table.datetime_units, tuple( ( stype, diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 78734c8892..7b5e6d9878 100644 --- a/sdm/models/base/icl.py +++ b/sdm/models/base/icl.py @@ -775,7 +775,11 @@ def _stack(tables: Sequence[TableTensor]) -> TableTensor: renamed: list[Tensor] = [ ref, *( - table.__class__(columns=columns, **dict(TableTensor.items(table))) + table.__class__( + columns=columns, + datetime_units=table.datetime_units, + **dict(TableTensor.items(table)), + ) for table in tables[1:] ), ] diff --git a/sdm/models/kumo/relational/model.py b/sdm/models/kumo/relational/model.py index 7be9720c99..64fc917be5 100644 --- a/sdm/models/kumo/relational/model.py +++ b/sdm/models/kumo/relational/model.py @@ -30,6 +30,7 @@ from sdm.models.tabiclv2.icl import ICLBlock from sdm.models.tabiclv2.row_embedding import RowEmbedding from sdm.processing import Recipe, Standardize +from sdm.tensor.datetime import DatetimeUnit, relative_days class KumoRelational(ICLModel): @@ -399,6 +400,10 @@ def forward( rel_time_i = self._get_rel_time( # Inject relative time: datetime=context.related_tables.tables[name].datetime, seed_datetime=x_context.datetime, + datetime_units=context.related_tables.tables[ + name + ].datetime_units, + seed_datetime_units=x_context.datetime_units, task_row=context_task_row_i, standardizer=standardizer, ) @@ -420,6 +425,10 @@ def forward( rel_time_i = self._get_rel_time( # Inject relative time: datetime=query.related_tables.tables[name].datetime, seed_datetime=x_query.datetime, + datetime_units=query.related_tables.tables[ + name + ].datetime_units, + seed_datetime_units=x_query.datetime_units, task_row=query_task_row_i, standardizer=cast( Standardize, cache[f"table_{name}.standardizer"] @@ -571,6 +580,8 @@ def _get_rel_time( seed_datetime: Tensor, task_row: Tensor, standardizer: Standardize, + datetime_units: Sequence[DatetimeUnit], + seed_datetime_units: Sequence[DatetimeUnit], ) -> Tensor | None: if datetime.size(-1) == 0 or seed_datetime.size(-1) == 0: @@ -582,8 +593,10 @@ def _get_rel_time( na_mask = na_mask.unsqueeze(-2) | (datetime == NaT).unsqueeze(-1) na_mask = na_mask.flatten(-2) - rel_time = seed_datetime.unsqueeze(-2) - datetime.unsqueeze(-1) - rel_time = rel_time.flatten(-2) / (24 * 60 * 60 * 1_000_000) + rel_time = relative_days( + datetime, datetime_units, seed_datetime, seed_datetime_units + ) + rel_time = rel_time.flatten(-2) rel_time = rel_time.sign() * rel_time.abs().log1p() if not standardizer.is_fitted: diff --git a/sdm/processing/common/select.py b/sdm/processing/common/select.py index 90475e0a6d..7a6255b873 100644 --- a/sdm/processing/common/select.py +++ b/sdm/processing/common/select.py @@ -58,7 +58,9 @@ def _transform_ensemble( stype: block[..., : self.max_columns] for stype, block in TableTensor.items(group) } - groups.append(group.__class__(columns=columns, **blocks)) + groups.append( + group._new_with_blocks(columns=columns, **blocks) + ) return ensemble_table.replace_groups(groups) assert self.method == "round_robin" @@ -97,7 +99,7 @@ def _transform_ensemble( dim=-1, ) tables.append( - table.__class__( + table._new_with_blocks( columns=columns, **blocks, ) diff --git a/sdm/processing/datetime/calendar.py b/sdm/processing/datetime/calendar.py index 611c55439c..274c9dd551 100644 --- a/sdm/processing/datetime/calendar.py +++ b/sdm/processing/datetime/calendar.py @@ -9,10 +9,7 @@ from sdm import NaT, Stype, TableTensor from sdm.processing import Processor - -US_PER_MINUTE = 60 * 1_000_000 -US_PER_HOUR = 60 * US_PER_MINUTE -US_PER_DAY = 24 * US_PER_HOUR +from sdm.tensor.datetime import units_per_day class AddCalendarFields(Processor): @@ -53,11 +50,12 @@ def _transform(self, table: TableTensor) -> TableTensor: return table datetime = table.datetime + per_day = units_per_day(table.datetime_units, datetime.device) na_mask = datetime == NaT days: Tensor | None = None if len({"weekday", "day_of_month", "month"} & set(self.fields)) > 0: - days = datetime.div(US_PER_DAY, rounding_mode="floor") + days = datetime.div(per_day, rounding_mode="floor") year = month = day = None if len({"day_of_month", "month"} & set(self.fields)) > 0: @@ -66,19 +64,21 @@ def _transform(self, table: TableTensor) -> TableTensor: time_of_day: Tensor | None = None if len({"minute", "hour"} & set(self.fields)) > 0: - time_of_day = datetime.remainder(US_PER_DAY) + time_of_day = datetime.remainder(per_day) outs: dict[str, Tensor] = {} for field in self.fields: if field == "minute": assert time_of_day is not None - out = time_of_day.div(US_PER_MINUTE, rounding_mode="floor") + out = time_of_day.div( + per_day // (24 * 60), rounding_mode="floor" + ) outs[field] = out.remainder(60) continue if field == "hour": assert time_of_day is not None - out = time_of_day.div(US_PER_HOUR, rounding_mode="floor") + out = time_of_day.div(per_day // 24, rounding_mode="floor") outs[field] = out continue diff --git a/sdm/relational/backend/_cugraph.py b/sdm/relational/backend/_cugraph.py index 41627ab67a..c2b47d7e52 100644 --- a/sdm/relational/backend/_cugraph.py +++ b/sdm/relational/backend/_cugraph.py @@ -11,6 +11,7 @@ from sdm import NaT, TableTensor from sdm._warnings import warn_once from sdm.relational import RelationalData, TaskLink +from sdm.relational.backend._time import TemporalClock from sdm.relational.join import join_index _INTEGER_DTYPES = { @@ -46,6 +47,15 @@ def __init__( self.data = data self.time_columns = time_columns + self._clock = TemporalClock( + { + name: ( + data.tables[name][column].datetime.squeeze(-1), + data.tables[name][column].datetime_units[0], + ) + for name, column in time_columns.items() + } + ) self._cp = cp self._pylibcugraph = pylibcugraph @@ -115,7 +125,10 @@ def sample( torch.iinfo(torch.int64).max, ) else: - seed_time = task_table[task_time_column].datetime.squeeze(-1) + table = task_table[task_time_column] + seed_time = self._clock.query_time( + table.datetime.squeeze(-1), table.datetime_units[0] + ) if self._num_edges == 0 or len(num_neighbors) == 0: return self._seed_nodes(seed=seed, task_link=task_link) @@ -153,12 +166,7 @@ def _build_graph(self) -> None: edge_types: list[Tensor] = [] edge_times: list[Tensor] = [] - times = { - table_name: self.data.tables[table_name][ - column_name - ].datetime.squeeze(-1) - for table_name, column_name in self.time_columns.items() - } + times = self._clock.times for i, (relationship, edge_index) in enumerate( zip(self.data.relationships, edge_indices) diff --git a/sdm/relational/backend/_pyg_lib.py b/sdm/relational/backend/_pyg_lib.py index 17876ddc22..7440d7cdf7 100644 --- a/sdm/relational/backend/_pyg_lib.py +++ b/sdm/relational/backend/_pyg_lib.py @@ -9,6 +9,7 @@ from sdm import TableTensor from sdm.relational import RelationalData, TaskLink +from sdm.relational.backend._time import TemporalClock from sdm.relational.join import join_index @@ -53,10 +54,15 @@ def __init__( self._row_dict: dict[tuple[str, str, str], Tensor] = {} self._colptr_dict: dict[tuple[str, str, str], Tensor] = {} - self._time_dict: dict[str, Tensor] = {} + times = {} for table_name, time_column in self.time_columns.items(): - time = self.data.tables[table_name][time_column].datetime - self._time_dict[table_name] = time.squeeze(-1).contiguous() + table = self.data.tables[table_name][time_column] + times[table_name] = ( + table.datetime.squeeze(-1).contiguous(), + table.datetime_units[0], + ) + self._clock = TemporalClock(times) + self._time_dict = self._clock.times for i, (rel, edge_index) in enumerate( zip(self.data.relationships, self.data.edge_indices()) @@ -112,7 +118,10 @@ def sample( ) if task_time_column is not None: - seed_time = task_table[task_time_column].datetime.squeeze(-1) + table = task_table[task_time_column] + seed_time = self._clock.query_time( + table.datetime.squeeze(-1), table.datetime_units[0] + ) else: fill_value = torch.iinfo(torch.int64).max seed_time = torch.full_like(seed, fill_value) diff --git a/sdm/relational/backend/_time.py b/sdm/relational/backend/_time.py new file mode 100644 index 0000000000..2f18364fd3 --- /dev/null +++ b/sdm/relational/backend/_time.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Mapping + +from torch import Tensor + +from sdm.tensor.datetime import DatetimeUnit, convert_datetime, datetime_parts + + +class TemporalClock: + """Normalize sampling clocks without changing stored feature timestamps. + + All-microsecond graphs retain their full range and raw clock values. + Graphs with nanosecond event timestamps use nanoseconds; mixed clocks + must fit the int64 nanosecond timestamp range. + """ + + def __init__( + self, times: Mapping[str, tuple[Tensor, DatetimeUnit]] + ) -> None: + self.unit: DatetimeUnit = ( + "ns" if any(unit == "ns" for _, unit in times.values()) else "us" + ) + self.times = { + name: convert_datetime(value, unit, self.unit) + for name, (value, unit) in times.items() + } + + def query_time(self, value: Tensor, unit: DatetimeUnit) -> Tensor: + if unit == "ns" and self.unit == "us": + # For an inclusive comparison against microsecond event clocks, + # t_us <= t_ns is exactly t_us <= floor(t_ns / 1000), including + # negative epochs. Only this comparison bound is derived; task + # tensors and their calendar/relative-time features stay lossless. + return datetime_parts(value, unit)[0] + return convert_datetime(value, unit, self.unit) diff --git a/sdm/relational/sampler.py b/sdm/relational/sampler.py index c3b18391ea..8f8406ba16 100644 --- a/sdm/relational/sampler.py +++ b/sdm/relational/sampler.py @@ -58,7 +58,11 @@ class RelationalSampler: data: The collection of named tables and their relationships. time_columns: Mapping from table name to the datetime column used for temporal sampling. A row in a time-aware table can only be sampled - if its timestamp does not exceed the query timestamp. + if its timestamp does not exceed the query timestamp. Nanosecond + query anchors remain exact. Graphs with nanosecond event clocks + require all configured event clocks and queries to fit the int64 + nanosecond timestamp range; all-microsecond graphs retain their + original supported range. """ def __init__( diff --git a/sdm/tensor/datetime.py b/sdm/tensor/datetime.py new file mode 100644 index 0000000000..dcfe5e1aa0 --- /dev/null +++ b/sdm/tensor/datetime.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence +from typing import Literal + +import torch +from torch import Tensor + +from sdm import NaT + +DatetimeUnit = Literal["us", "ns"] +US_PER_DAY = 24 * 60 * 60 * 1_000_000 +NS_PER_DAY = US_PER_DAY * 1_000 + + +def datetime_parts(value: Tensor, unit: DatetimeUnit) -> tuple[Tensor, Tensor]: + """Return exact microseconds and a nanosecond remainder for ordering.""" + if unit == "us": + return value, torch.zeros_like(value) + coarse = value.div(1_000, rounding_mode="floor") + coarse = coarse.masked_fill(value == NaT, NaT) + return coarse, value.remainder(1_000).masked_fill(value == NaT, 0) + + +def convert_datetime( + value: Tensor, + source: DatetimeUnit, + target: DatetimeUnit, +) -> Tensor: + """Convert a datetime block without dropping precision or overflowing. + + Cross-unit conversion validates user data at the conversion boundary, + synchronizing CUDA inputs so invalid data raises without a device assert. + Equal-unit operations retain the original tensor without validation. + """ + if source == target: + return value + valid = value != NaT + if source == "us": + low = (torch.iinfo(torch.int64).min + 1 + 999) // 1_000 + high = torch.iinfo(torch.int64).max // 1_000 + if not ((~valid) | ((value >= low) & (value <= high))).all(): + raise ValueError( + "Datetime value is outside the nanosecond timestamp range" + ) + return (value.masked_fill(~valid, 0) * 1_000).masked_fill(~valid, NaT) + if not ((~valid) | (value.remainder(1_000) == 0)).all(): + raise ValueError("Datetime conversion would lose nanosecond precision") + return value.div(1_000, rounding_mode="floor").masked_fill(~valid, NaT) + + +def units_per_day( + units: Sequence[DatetimeUnit], device: torch.device +) -> int | Tensor: + """Return scalar or column-wise clock ticks per calendar day.""" + if all(unit == "us" for unit in units): + return US_PER_DAY + if all(unit == "ns" for unit in units): + return NS_PER_DAY + return torch.tensor( + [US_PER_DAY if unit == "us" else NS_PER_DAY for unit in units], + dtype=torch.int64, + device=device, + ) + + +def relative_days( + value: Tensor, + units: Sequence[DatetimeUnit], + seed: Tensor, + seed_units: Sequence[DatetimeUnit], +) -> Tensor: + """Compute relative day features while preserving subsecond differences.""" + if all(unit == "us" for unit in (*units, *seed_units)): + return (seed.unsqueeze(-2) - value.unsqueeze(-1)) / US_PER_DAY + value_scale = units_per_day(units, value.device) // (24 * 60 * 60) + seed_scale = units_per_day(seed_units, seed.device) // (24 * 60 * 60) + # Split before subtracting: absolute epoch nanoseconds lose precision in + # float64, and a direct int64 subtraction can overflow. + seconds = seed.div(seed_scale, rounding_mode="floor").unsqueeze(-2) + seconds = seconds - value.div( + value_scale, rounding_mode="floor" + ).unsqueeze(-1) + seed_ns = seed.remainder(seed_scale) * (1_000_000_000 // seed_scale) + value_ns = value.remainder(value_scale) * (1_000_000_000 // value_scale) + nanos = seed_ns.unsqueeze(-2) - value_ns.unsqueeze(-1) + seconds = seconds + nanos.div(1_000_000_000, rounding_mode="floor") + nanos = nanos.remainder(1_000_000_000) + # Match signs to preserve +/-1ns across a second boundary. + borrow = (seconds < 0) & (nanos != 0) + seconds = seconds + borrow + nanos = nanos - borrow * 1_000_000_000 + result = ( + seconds.to(torch.float64) + nanos.to(torch.float64) / 1_000_000_000 + ) / (24 * 60 * 60) + return result.to(torch.get_default_dtype()) diff --git a/sdm/tensor/table.py b/sdm/tensor/table.py index 996fc8b0e6..f39c083ffc 100644 --- a/sdm/tensor/table.py +++ b/sdm/tensor/table.py @@ -20,6 +20,7 @@ from sdm import NaT, Stype, StypeLike from sdm.tensor import CategoricalTensor, ColumnarTensor, StringTensor +from sdm.tensor.datetime import DatetimeUnit, convert_datetime, datetime_parts from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf from sdm.tensor.mixin import _resolve_device @@ -197,10 +198,19 @@ class TableTensor(Tensor): columns: Column names grouped by semantic type. numerical: The numerical column block of shape ``[..., C_num]``. categorical: The categorical column block of shape ``[..., C_cat]``. - datetime: The ``datetime64[us]`` column block of shape ``[..., C_dt]``. + datetime: Integer timestamps of shape ``[..., C_dt]``. Units are given + by ``datetime_units``; existing raw inputs default to microseconds. text: The text column block of shape ``[..., C_text]``. id: The identifier column block of shape ``[..., C_id]``. device: The device. + datetime_units: Units of each datetime column, defaulting to ``"us"``. + Arrow/pandas ingress retains nanoseconds when conversion to + microseconds would lose precision. Combining rows with different + units promotes the affected columns to nanoseconds, requiring + values to fit the int64 nanosecond timestamp range. In-place + assignment requires matching units to preserve views. Convert mixed + units eagerly before compiled execution; cross-unit range + validation synchronizes CUDA, while unchanged units do not. """ HANDLED_FUNCTIONS: ClassVar[ @@ -210,6 +220,7 @@ class TableTensor(Tensor): _numerical: Tensor _categorical: CategoricalTensor _datetime: Tensor + _datetime_units: tuple[DatetimeUnit, ...] _text: StringTensor _id: ColumnarTensor _columns: dict[Stype, tuple[str, ...]] @@ -230,6 +241,7 @@ def __init__( text: StringTensor | None = None, id: ColumnarTensor | None = None, device: torch.device | str | None = None, + datetime_units: Sequence[DatetimeUnit] | None = None, ) -> None: pass @@ -243,6 +255,7 @@ def __new__( text: StringTensor | None = None, id: ColumnarTensor | None = None, device: torch.device | str | None = None, + datetime_units: Sequence[DatetimeUnit] | None = None, ) -> Self: r"""Create a tensor wrapper.""" if size is not None and len(size) == 0: @@ -310,6 +323,14 @@ def __new__( if id is None: id = ColumnarTensor((), size=size, device=device) + datetime_units = ( + ("us",) * datetime.size(-1) + if datetime_units is None + else tuple(datetime_units) + ) + if len(datetime_units) != datetime.size(-1): + raise ValueError("Expected one datetime unit per datetime column") + columns = { Stype(stype): tuple(names) for stype, names in (columns or {}).items() @@ -344,6 +365,7 @@ def __new__( numerical=numerical, categorical=categorical, datetime=datetime, + datetime_units=datetime_units, text=text, id=id, cls=None if cls is TableTensor else cls, @@ -357,6 +379,7 @@ def _make_table_tensor( numerical: Tensor, categorical: CategoricalTensor, datetime: Tensor, + datetime_units: tuple[DatetimeUnit, ...], text: StringTensor, id: ColumnarTensor, cls: type[TableTensor] | None = None, @@ -380,6 +403,7 @@ def _make_table_tensor( out._numerical = numerical out._categorical = categorical out._datetime = datetime + out._datetime_units = datetime_units out._text = text out._id = id out._columns = stype_columns @@ -420,6 +444,7 @@ def from_arrow( for column, stype in stypes.items(): columns[Stype(stype)].append(column) + datetime_units: list[DatetimeUnit] = [] blocks: dict[Stype, Tensor] = {} for stype in columns: tensors: list[Tensor] = [] @@ -433,7 +458,19 @@ def from_arrow( elif stype == Stype.categorical: tensor = CategoricalTensor.from_arrow(array) elif stype == Stype.datetime: - array = array.cast(pa.timestamp("us")) + try: + converted = array.cast(pa.timestamp("us")) + except pa.ArrowInvalid: + if ( + not pa.types.is_timestamp(array.type) + or array.type.unit != "ns" + ): + raise + converted = array.cast(pa.timestamp("ns")) + datetime_units.append("ns") + else: + datetime_units.append("us") + array = converted values = array.to_numpy(zero_copy_only=False) values = values.astype("int64") tensor = torch.from_numpy(values).unsqueeze(-1) @@ -449,6 +486,7 @@ def from_arrow( return cls( columns=cast(Mapping[StypeLike, Sequence[str]], columns), + datetime_units=datetime_units, **blocks, ) @@ -473,8 +511,10 @@ def to_arrow(self) -> pa.Table: elif stype == Stype.datetime: tensor = tensor.movedim(-1, 0).contiguous().cpu() array = tensor.numpy().reshape(tensor.size(0), -1) - array = array.view("datetime64[us]") - arrays.extend(pa.array(a) for a in array) + arrays.extend( + pa.array(a.view(f"datetime64[{unit}]")) + for a, unit in zip(array, self.datetime_units) + ) else: tensor = tensor.detach().movedim(-1, 0).contiguous().cpu() arrays.extend(to_arrow(t) for t in tensor) @@ -767,6 +807,7 @@ def from_cudf( for column, stype in stypes.items(): columns[Stype(stype)].append(column) + datetime_units: list[DatetimeUnit] = [] blocks: dict[Stype, Tensor] = {} for stype in columns: tensors: list[Tensor] = [] @@ -781,7 +822,14 @@ def from_cudf( elif stype == Stype.categorical: tensor = CategoricalTensor.from_cudf(ser, device=device) elif stype == Stype.datetime: - ser = ser.astype("datetime64[us]", copy=False) + unit: DatetimeUnit = "us" + if ( + ser.dtype == np.dtype("datetime64[ns]") + and (ser.astype("int64").fillna(0) % 1_000 != 0).any() + ): + unit = "ns" + datetime_units.append(unit) + ser = ser.astype(f"datetime64[{unit}]", copy=False) ser = ser.astype("int64", copy=False) if ser.null_count > 0: ser = ser.fillna(NaT) @@ -800,6 +848,7 @@ def from_cudf( return cls( columns=cast(Mapping[StypeLike, Sequence[str]], columns), + datetime_units=datetime_units, **blocks, ) @@ -828,12 +877,13 @@ def to_cudf(self) -> cudf.DataFrame: df = cudf.DataFrame( { name: to_cudf(column, mask).astype( - "datetime64[us]", copy=False + f"datetime64[{unit}]", copy=False ) - for name, column, mask in zip( + for name, column, mask, unit in zip( self._columns[stype], tensor, tensor != NaT, + self.datetime_units, ) } ) @@ -979,6 +1029,48 @@ def is_same_schema(self, other: TableTensor) -> bool: """ return self.schema == other.schema + @property + def datetime_units(self) -> tuple[DatetimeUnit, ...]: + """Units of the integer timestamps, in datetime-column order.""" + return getattr( + self, "_datetime_units", ("us",) * self.datetime.size(-1) + ) + + def _new_with_blocks( + self, + *, + columns: Mapping[StypeLike, Sequence[str]] | None = None, + datetime_units: Sequence[DatetimeUnit] | None = None, + size: Sequence[int] | None = None, + device: torch.device | str | None = None, + **blocks: Any, + ) -> Self: + columns = ( + cast(Mapping[StypeLike, Sequence[str]], self._columns) + if columns is None + else columns + ) + if datetime_units is None: + names = next( + ( + names + for stype, names in columns.items() + if Stype(stype) == Stype.datetime + ), + (), + ) + datetime_units = tuple( + self.datetime_units[self._column_to_loc[name][1]] + for name in names + ) + return self.__class__( + columns=columns, + datetime_units=datetime_units, + size=size, + device=device, + **blocks, + ) + def replace_blocks( self, *, @@ -1004,7 +1096,7 @@ def replace_blocks( text: Replacement text block with shape ``[..., C_text]``. id: Replacement identifier block with shape ``[..., C_id]``. """ - return self.__class__( + return self._new_with_blocks( columns=cast(Mapping[StypeLike, Sequence[str]], self._columns), numerical=self.numerical if numerical is None else numerical, categorical=( @@ -1028,7 +1120,7 @@ def select_stypes( stypes = (stypes,) stypes = tuple(Stype(stype) for stype in stypes) - return self.__class__( + return self._new_with_blocks( size=self.size()[:-1], columns={stype: self._columns[stype] for stype in stypes}, device=self.device, @@ -1056,7 +1148,7 @@ def drop_stypes( stypes = {Stype(stype) for stype in stypes} keep = tuple(stype for stype in self._columns if stype not in stypes) - return self.__class__( + return self._new_with_blocks( size=self.size()[:-1], columns={stype: self._columns[stype] for stype in keep}, device=self.device, @@ -1104,7 +1196,7 @@ def select_columns(self, columns: str | Iterable[str]) -> Self: dim=-1, ) - return self.__class__(columns=columns_dict, **blocks) + return self._new_with_blocks(columns=columns_dict, **blocks) def drop_columns(self, columns: str | Iterable[str]) -> Self: r"""Return a table with ``columns`` removed. @@ -1157,7 +1249,7 @@ def __tensor_flatten__(self) -> tuple[list[str], tuple[Any, ...]]: columns = tuple( (str(stype), names) for stype, names in self._columns.items() ) - ctx = (self.__class__, columns) + ctx = (self.__class__, columns, self.datetime_units) return attrs, ctx @staticmethod @@ -1167,13 +1259,15 @@ def __tensor_unflatten__( outer_size: tuple[int, ...], outer_stride: tuple[int, ...], ) -> TableTensor: - cls, columns = ctx + cls, columns, *precision = ctx + datetime_units = precision[0] if precision else None return cls( size=outer_size[:-1], columns=dict(columns), numerical=inner_tensors["numerical"], categorical=inner_tensors["categorical"], datetime=inner_tensors["datetime"], + datetime_units=datetime_units, text=inner_tensors["text"], id=inner_tensors["id"], ) @@ -1187,6 +1281,8 @@ def __reduce_ex__(self, proto: SupportsIndex) -> Any: self._datetime, self._text, self._id, + None, + self.datetime_units, ) return (self.__class__, args) @@ -1198,6 +1294,16 @@ def __torch_dispatch__( # type: ignore args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> Any: + if func in (aten.copy_.default, aten.index_put_.default): + tables = [arg for arg in args if isinstance(arg, TableTensor)] + if len(tables) > 1 and any( + table.datetime_units != tables[0].datetime_units + for table in tables[1:] + ): + raise ValueError( + "In-place table assignment requires matching " + "datetime units" + ) if func in cls.HANDLED_FUNCTIONS: return cls.HANDLED_FUNCTIONS[func](*args, **(kwargs or {})) @@ -1332,7 +1438,7 @@ def _alias(inp: TableTensor) -> TableTensor: blocks = { stype: aten.alias.default(tensor) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1393,7 +1499,7 @@ def _to_dtype_layout( ) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1496,7 +1602,7 @@ def _contiguous( stype: tensor.contiguous(memory_format=memory_format) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1515,7 +1621,7 @@ def _pin_memory(inp: TableTensor) -> TableTensor: stype: tensor.pin_memory() if tensor.numel() > 0 else tensor for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1540,8 +1646,12 @@ def _equal(inp: TableTensor, other: Tensor) -> bool: if inp.stypes != other.stypes: return False - for stype, block in _align_like(inp, other).items(): - if not block.equal(other.blocks[stype]): + aligned = _align_like(inp, other) + for stype, block in aligned.items(): + if stype == Stype.datetime: + if not _datetime_equal(aligned, other): + return False + elif not block.equal(other.blocks[stype]): return False return True @@ -1562,8 +1672,12 @@ def _allclose( if inp.stypes != other.stypes: return False - for stype, block in _align_like(inp, other).items(): - if stype != Stype.numerical: # Tolerance only applies for numerical: + aligned = _align_like(inp, other) + for stype, block in aligned.items(): + if stype == Stype.datetime: + if not _datetime_equal(aligned, other): + return False + elif stype != Stype.numerical: # Tolerance only applies for numerical: if not block.equal(other.blocks[stype]): return False elif not block.allclose(other.blocks[stype], rtol, atol, equal_nan): @@ -1612,7 +1726,7 @@ def _view(inp: TableTensor, size: Sequence[int]) -> TableTensor: stype: tensor.view((*size[:-1], tensor.size(-1))) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1679,7 +1793,7 @@ def _squeeze_dims(inp: TableTensor, dim: Sequence[int]) -> TableTensor: f"Can't squeeze the column dimension of {inp.__class__.__name__!r}" ) - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1696,7 +1810,7 @@ def _unsqueeze(inp: TableTensor, dim: int) -> TableTensor: f"{inp.__class__.__name__!r}" ) - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1726,7 +1840,7 @@ def _expand( ) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1747,7 +1861,7 @@ def _transpose(inp: TableTensor, dim0: int, dim1: int) -> TableTensor: f"{inp.__class__.__name__!r}" ) - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1765,7 +1879,7 @@ def _permute(inp: TableTensor, dims: Sequence[int]) -> TableTensor: f"Can't permute the column dimension of {inp.__class__.__name__!r}" ) - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1807,7 +1921,7 @@ def _select(inp: TableTensor, dim: int, index: int) -> TableTensor: stype: tensor.select(dim, index) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1844,7 +1958,7 @@ def _slice( for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1872,7 +1986,7 @@ def _narrow( for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -1890,7 +2004,7 @@ def _unbind(inp: TableTensor, dim: int = 0) -> tuple[TableTensor, ...]: stypes = tuple(tensors_dict.keys()) return tuple( - inp.__class__( + inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **dict(zip(stypes, blocks)), ) @@ -1912,7 +2026,7 @@ def _split( f"{inp.__class__.__name__!r} with split size 1" ) return tuple( - inp.__class__( + inp._new_with_blocks( columns={stype: (name,)}, **{stype: tensor.narrow(-1, i, 1)}, ) @@ -1926,7 +2040,7 @@ def _split( stypes = tuple(tensors_dict.keys()) return tuple( - inp.__class__( + inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **dict(zip(stypes, blocks)), ) @@ -1955,7 +2069,7 @@ def _split_with_sizes( stypes = tuple(blocks_dict.keys()) return tuple( - inp.__class__( + inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **dict(zip(stypes, blocks)), ) @@ -1978,7 +2092,7 @@ def _index_select( stype: tensor.index_select(dim, index) for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -2010,7 +2124,7 @@ def _index( for stype, tensor in inp.items() } - return inp.__class__( + return inp._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], inp._columns), **blocks, ) @@ -2042,9 +2156,19 @@ def _cat(tensors: Sequence[Tensor], dim: int = 0) -> TableTensor: f"{tuple(tensor.size())})" ) + datetime_units = ( + tuple(chain.from_iterable(t.datetime_units for t in tensors)) + if is_column_dim + else _common_datetime_units(tensors) + ) blocks: dict[Stype, Tensor] = {} for stype, _ in TableTensor.items(ref): block_list = [getattr(tensor, stype.value) for tensor in tensors] + if stype == Stype.datetime and not is_column_dim: + block_list = [ + _datetime_in_units(tensor, datetime_units) + for tensor in tensors + ] if is_column_dim: # Skip empty blocks so their dtypes cannot promote the others, # but keep one so that an all-empty block retains its dtype. @@ -2066,8 +2190,9 @@ def _cat(tensors: Sequence[Tensor], dim: int = 0) -> TableTensor: for stype, _ in TableTensor.items(ref) } - return ref.__class__( + return ref._new_with_blocks( columns=cast("dict[StypeLike, tuple[str, ...]]", columns), + datetime_units=datetime_units, **blocks, ) @@ -2095,13 +2220,19 @@ def _stack(tensors: Sequence[Tensor], dim: int = 0) -> TableTensor: ref = tensors[0] tensors = (ref, *(_align_like(tensor, ref) for tensor in tensors[1:])) + datetime_units = _common_datetime_units(tensors) blocks = { stype: torch.stack([tensor.blocks[stype] for tensor in tensors], dim) for stype in ref._columns + if stype != Stype.datetime } + blocks[Stype.datetime] = torch.stack( + [_datetime_in_units(tensor, datetime_units) for tensor in tensors], dim + ) - return ref.__class__( + return ref._new_with_blocks( columns=cast(dict[StypeLike, tuple[str, ...]], ref._columns), + datetime_units=datetime_units, **blocks, ) @@ -2109,6 +2240,49 @@ def _stack(tensors: Sequence[Tensor], dim: int = 0) -> TableTensor: # Helpers ##################################################################### +def _common_datetime_units( + tables: Sequence[TableTensor], +) -> tuple[DatetimeUnit, ...]: + return tuple( + "ns" + if any(table.datetime_units[i] == "ns" for table in tables) + else "us" + for i in range(len(tables[0].datetime_units)) + ) + + +def _datetime_in_units( + table: TableTensor, units: Sequence[DatetimeUnit] +) -> Tensor: + if tuple(units) == table.datetime_units: + return table.datetime + return torch.cat( + [ + convert_datetime(column, source, target) + for column, source, target in zip( + table.datetime.split(1, dim=-1), table.datetime_units, units + ) + ], + dim=-1, + ) + + +def _datetime_equal(left: TableTensor, right: TableTensor) -> bool: + if left.datetime_units == right.datetime_units: + return left.datetime.equal(right.datetime) + for a, b, a_unit, b_unit in zip( + left.datetime.unbind(-1), + right.datetime.unbind(-1), + left.datetime_units, + right.datetime_units, + ): + a_parts = datetime_parts(a, a_unit) + b_parts = datetime_parts(b, b_unit) + if not all(a.equal(b) for a, b in zip(a_parts, b_parts)): + return False + return True + + def _block_size_repr(size: Sequence[int]) -> str: if len(size) == 0: return "(*,)" @@ -2152,7 +2326,7 @@ def _align_like(inp: TableTensor, ref: TableTensor) -> TableTensor: else: blocks[stype] = inp.blocks[stype] - return inp.__class__( + return inp._new_with_blocks( columns=cast(Mapping[StypeLike, Sequence[str]], ref._columns), **blocks, ) diff --git a/test/models/base/test_datetime_precision.py b/test/models/base/test_datetime_precision.py new file mode 100644 index 0000000000..fc417a4cdc --- /dev/null +++ b/test/models/base/test_datetime_precision.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +import pyarrow as pa +import pytest +import torch + +from sdm import RelatedTables, Stype, TableTensor +from sdm.cache import Cache +from sdm.models import ICLModel +from sdm.processing import AddCalendarFields, Recipe + + +class CalendarModel(ICLModel): + supported_feature_stypes = frozenset({Stype.datetime}) + supported_target_stypes = frozenset({Stype.numerical}) + supports_multi_target = False + supports_context_only_columns = False + supports_related_tables = False + + @classmethod + def default_recipe(cls) -> Recipe: + return Recipe() + + def _forward( + self, + x_context: TableTensor | None, + y_context: TableTensor | None, + x_query: TableTensor | None, + x_context_only: TableTensor | None, + related_context_tables: RelatedTables | None, + related_query_tables: RelatedTables | None, + cache: Cache | None, + generator: torch.Generator | None, + **kwargs: Any, + ) -> TableTensor: + table = x_context if x_query is None else x_query + assert table is not None + return ( + AddCalendarFields(["month"]) + .transform(table) + .select_columns("precise__month") + ) + + +@pytest.mark.parametrize("estimator_batch_size", [None, 2]) +def test_eight_member_model_keeps_datetime_units_during_execution( + estimator_batch_size: int | None, +) -> None: + midnight = 1_704_067_200_000_000_000 + table = TableTensor.from_arrow( + pa.table( + { + "legacy": pa.array( + [midnight // 1_000] * 2, type=pa.timestamp("us") + ), + "precise": pa.array( + [midnight - 1, midnight + 1], type=pa.timestamp("ns") + ), + } + ), + {"legacy": "datetime", "precise": "datetime"}, + ) + model = CalendarModel() + output = model( + x_context=table, + y_context=torch.zeros(2, 1), + x_query=table, + recipe=Recipe(), + num_estimators=8, + estimator_batch_size=estimator_batch_size, + ) + torch.testing.assert_close( + output.numerical, torch.tensor([[[11.0], [0.0]]]).expand(8, -1, -1) + ) + + +@pytest.mark.parametrize("estimator_batch_size", [None, 2]) +def test_fitted_eight_member_model_preserves_query_precision( + estimator_batch_size: int | None, +) -> None: + table = TableTensor.from_arrow( + pa.table({"precise": pa.array([-1, 1], type=pa.timestamp("ns"))}), + {"precise": "datetime"}, + ) + model = CalendarModel().eval() + model.fit( + table, + torch.zeros(2, 1), + recipe=Recipe(), + num_estimators=8, + estimator_batch_size=estimator_batch_size, + ) + output = model.predict(table) + torch.testing.assert_close( + output.numerical, torch.tensor([[[11.0], [0.0]]]).expand(8, -1, -1) + ) diff --git a/test/processing/common/test_select.py b/test/processing/common/test_select.py index 3a701f7278..968937f382 100644 --- a/test/processing/common/test_select.py +++ b/test/processing/common/test_select.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from typing import Literal + import pytest import torch @@ -46,3 +48,25 @@ def test_select_columns_round_robin_routes_members() -> None: columns=columns, ) assert out[member_id].equal(expected) + + +@pytest.mark.parametrize("method", ["first", "round_robin"]) +def test_select_columns_preserves_mixed_datetime_units( + method: Literal["first", "round_robin"], +) -> None: + table = TableTensor( + columns={"datetime": ("us", "ns", "other_us")}, + datetime=torch.tensor([[1, 1_001, 2], [2, 2_001, 3]]), + datetime_units=("us", "ns", "us"), + ) + output = SelectColumns( + max_columns=2, method=method + ).fit_transform_ensemble(EnsembleTable.from_table(table, num_members=3)) + columns = ( + [("us", "ns")] * 3 + if method == "first" + else [("us", "ns"), ("other_us", "us"), ("ns", "other_us")] + ) + original = table.to_arrow() + for i, names in enumerate(columns): + assert output[i].to_arrow().equals(original.select(names)) diff --git a/test/relational/test_datetime_precision.py b/test/relational/test_datetime_precision.py new file mode 100644 index 0000000000..2696d70dd2 --- /dev/null +++ b/test/relational/test_datetime_precision.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from datetime import datetime + +import pyarrow as pa +import pytest +import torch + +from sdm import RelationalData, Relationship, TableTensor, TaskLink +from sdm.testing import withCUDA + + +def _data(event_times: pa.Array, parent_time: pa.Array) -> RelationalData: + parents = TableTensor.from_arrow( + pa.table({"parent_id": [0], "time": parent_time}), + {"parent_id": "id", "time": "datetime"}, + ) + events = TableTensor.from_arrow( + pa.table( + { + "parent_id": [0] * len(event_times), + "event_id": range(len(event_times)), + "time": event_times, + } + ), + {"parent_id": "id", "event_id": "id", "time": "datetime"}, + ) + return RelationalData( + tables={"parents": parents, "events": events}, + relationships=[ + Relationship( + left_table="events", + left_columns=("parent_id",), + right_table="parents", + right_columns=("parent_id",), + ) + ], + ) + + +def _sample( + data: RelationalData, times: pa.Array +) -> tuple[list[list[int]], TableTensor]: + queries = TableTensor.from_arrow( + pa.table({"parent_id": [0] * len(times), "time": times}), + {"parent_id": "id", "time": "datetime"}, + device=data.device, + ) + output = data.sampler(time_columns={"parents": "time", "events": "time"})( + task_table=queries, + task_link=TaskLink( + task_columns=("parent_id",), + table="parents", + table_columns=("parent_id",), + ), + task_time_column="time", + num_neighbors=[10], + temporal_strategy="last", + ) + frame = output.related_tables.tables["events"].to_pandas() + groups = { + int(i): sorted(values.event_id.tolist()) + for i, values in frame.groupby("__example__") + } + return [groups.get(i, []) for i in range(len(times))], output.task_table + + +@withCUDA +@pytest.mark.parametrize("base", [-1_000_000, 1_704_067_200_000_000_000]) +def test_ns_cutoffs_preserve_anchor_observations_and_exclude_next_microsecond( + base: int, + device: torch.device, +) -> None: + pytest.importorskip("pyg_lib") + if device.type == "cuda": + pytest.importorskip("pylibcugraph") + data = _data( + pa.array( + [base // 1_000 - 1, base // 1_000, base // 1_000 + 1], + type=pa.timestamp("us"), + ), + pa.array([base // 1_000 - 10], type=pa.timestamp("us")), + ) + data = data.to(device) + times = pa.array( + [base - 1, base, base + 1, base + 999, base + 1_000, base + 1_001], + type=pa.timestamp("ns"), + ) + actual, returned_queries = _sample(data, times) + assert actual == [[0], [0, 1], [0, 1], [0, 1], [0, 1, 2], [0, 1, 2]] + assert returned_queries.datetime_units == ("ns",) + assert ( + returned_queries["time"] + .to_arrow()["time"] + .combine_chunks() + .equals(times) + ) + assert data.tables["events"].datetime_units == ("us",) + + +@withCUDA +def test_ns_event_equality_and_ordering_with_mixed_event_clocks( + device: torch.device, +) -> None: + pytest.importorskip("pyg_lib") + if device.type == "cuda": + pytest.importorskip("pylibcugraph") + base = 1_704_067_200_000_000_000 + data = _data( + pa.array( + [base - 1, base, base + 1, base + 1_000], type=pa.timestamp("ns") + ), + pa.array([base // 1_000 - 10], type=pa.timestamp("us")), + ) + data = data.to(device) + actual, _ = _sample( + data, + pa.array( + [base - 1, base, base + 1, base + 2], type=pa.timestamp("ns") + ), + ) + assert actual == [[0], [0, 1], [0, 1, 2], [0, 1, 2]] + # A microsecond query must not admit +1ns events in a nanosecond graph. + actual, returned = _sample( + data, pa.array([base // 1_000], type=pa.timestamp("us")) + ) + assert actual == [[0, 1]] + assert returned.datetime_units == ("us",) + + +def test_legacy_wide_us_graph_and_explicit_new_mixed_range_error() -> None: + pytest.importorskip("pyg_lib") + wide = pa.array([datetime(3000, 1, 1)], type=pa.timestamp("us")) + data = _data(wide, wide) + actual, _ = _sample(data, wide) + assert actual == [[0]] + mixed = _data(pa.array([1], type=pa.timestamp("ns")), wide) + with pytest.raises(ValueError, match="nanosecond timestamp range"): + mixed.sampler(time_columns={"parents": "time", "events": "time"}) + assert ( + data.tables["parents"].to_arrow()["time"].combine_chunks().equals(wide) + ) diff --git a/test/tensor/test_datetime_precision.py b/test/tensor/test_datetime_precision.py new file mode 100644 index 0000000000..981ae2989a --- /dev/null +++ b/test/tensor/test_datetime_precision.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pickle +from datetime import datetime +from typing import cast + +import pandas as pd +import pyarrow as pa +import pytest +import torch + +from sdm import EnsembleTable, NaT, TableTensor +from sdm.processing import AddCalendarFields +from sdm.tensor.datetime import relative_days +from sdm.testing import onlyCUDA, withCUDA + + +def test_lossless_datetime_ingress_and_roundtrip() -> None: + arrow = pa.table( + { + "wide": pa.array( + [datetime(3000, 1, 1), None], type=pa.timestamp("us") + ), + "precise": pa.array([1, -1], type=pa.timestamp("ns")), + "aligned": pa.array([1_000, -1_000], type=pa.timestamp("ns")), + } + ) + table = TableTensor.from_arrow( + arrow, dict.fromkeys(arrow.column_names, "datetime") + ) + assert table.datetime_units == ("us", "ns", "us") + result = table.to_arrow() + assert result["wide"].equals(arrow["wide"]) + assert result["precise"].equals(arrow["precise"]) + assert result["aligned"].cast(pa.timestamp("ns")).equals(arrow["aligned"]) + frame = pd.DataFrame({"time": pd.to_datetime([1, -1, None], unit="ns")}) + pandas_table = TableTensor.from_pandas(frame, {"time": "datetime"}) + assert pandas_table.datetime_units == ("ns",) + pd.testing.assert_frame_equal(pandas_table.to_pandas(), frame) + + +@withCUDA +def test_precision_metadata_survives_table_operations_and_pickle( + device: torch.device, +) -> None: + table = TableTensor.from_arrow( + pa.table({"time": pa.array([1, -1, None], type=pa.timestamp("ns"))}), + {"time": "datetime"}, + device=device, + ) + copies = [ + table.clone(), + table.to(dtype=torch.float64), + table[1:], + table.unsqueeze(0).expand(2, -1, -1)[0], + table["time"], + table.select_stypes("datetime"), + table.replace_blocks(datetime=table.datetime.clone()), + pickle.loads(pickle.dumps(table)), + torch.cat( + [table, TableTensor(numerical=torch.ones(3, 1, device=device))], + dim=-1, + ), + ] + for copy in copies: + copy = cast(TableTensor, copy) + assert copy.datetime_units == ("ns",) + expected = table[1:] if len(copy) == 2 else table + assert copy["time"].to_arrow().equals(expected.to_arrow()) + assert table.drop_stypes("datetime").datetime_units == () + + +def test_mixed_precision_rows_promote_losslessly() -> None: + micros = TableTensor(datetime=torch.tensor([[1], [-1], [NaT]])) + nanos = TableTensor( + datetime=torch.tensor([[1_000], [-1_000], [NaT]]), + datetime_units=("ns",), + ) + assert micros.equal(nanos) + assert micros.allclose(nanos) + for result in [torch.cat([micros, nanos]), torch.stack([micros, nanos])]: + result = cast(TableTensor, result) + assert result.datetime_units == ("ns",) + assert result.datetime.flatten().tolist() == [1_000, -1_000, NaT] * 2 + named_nanos = TableTensor( + columns={"datetime": ("ns",)}, + datetime=nanos.datetime, + datetime_units=nanos.datetime_units, + ) + mixed_columns = cast(TableTensor, torch.cat([micros, named_nanos], dim=-1)) + assert mixed_columns.datetime_units == ("us", "ns") + assert mixed_columns["ns"].datetime_units == ("ns",) + + +def test_heterogeneous_ensemble_preserves_precision_and_wide_dates() -> None: + wide = TableTensor(datetime=torch.tensor([[32_503_680_000_000_000]])) + precise = TableTensor(datetime=torch.tensor([[1]]), datetime_units=("ns",)) + tables = [wide, precise] + ensemble = EnsembleTable.from_tables(tables, member_table_ids=[0, 1]) + packed = ensemble.replace_tables(tables=tables, member_table_ids=[0, 1]) + for output in [ensemble, packed]: + for i, expected in enumerate(tables): + assert output[i].datetime_units == expected.datetime_units + assert output[i].to_arrow().equals(expected.to_arrow()) + + +def test_new_mixed_precision_range_error_does_not_change_legacy_data() -> None: + wide = TableTensor(datetime=torch.tensor([[32_503_680_000_000_000]])) + assert wide.datetime_units == ("us",) + assert wide.to_arrow()["dt_0"][0].as_py() == datetime(3000, 1, 1) + precise = TableTensor(datetime=torch.tensor([[1]]), datetime_units=("ns",)) + with pytest.raises(ValueError, match="nanosecond timestamp range"): + torch.cat([wide, precise]) + assert wide.to_arrow()["dt_0"][0].as_py() == datetime(3000, 1, 1) + + +def test_mixed_precision_inplace_view_assignment_is_rejected() -> None: + parent = TableTensor(datetime=torch.tensor([[1], [2]])) + before = parent.to_arrow() + view = parent[:1] + source = TableTensor( + datetime=torch.tensor([[1_001]]), datetime_units=("ns",) + ) + with pytest.raises(ValueError, match="matching datetime units"): + view.copy_(source) + with pytest.raises(ValueError, match="matching datetime units"): + parent[:1] = source + assert parent.to_arrow().equals(before) + + +def test_legacy_serialized_constructor_defaults_to_microseconds() -> None: + original = TableTensor(datetime=torch.tensor([[1_700_000_000_000_000]])) + + class LegacyPayload: + def __reduce__(self): + # Constructor arguments emitted by the old pickle format. + return TableTensor, ( + original.size()[:-1], + original.columns, + original.numerical, + original.categorical, + original.datetime, + original.text, + original.id, + ) + + restored = pickle.loads(pickle.dumps(LegacyPayload())) + assert restored.datetime_units == ("us",) + assert restored.to_arrow().equals(original.to_arrow()) + + +@withCUDA +def test_nanosecond_calendar_rollover_with_wide_microsecond_column( + device: torch.device, +) -> None: + midnight = 1_704_067_200_000_000_000 # 2024-01-01, a Monday. + arrow = pa.table( + { + "legacy": pa.array( + [datetime(3000, 1, 1)] * 4, type=pa.timestamp("us") + ), + "precise": pa.array( + [midnight - 1, midnight, midnight + 1, None], + type=pa.timestamp("ns"), + ), + } + ) + table = TableTensor.from_arrow( + arrow, dict.fromkeys(arrow.column_names, "datetime"), device=device + ) + output = AddCalendarFields( + ["minute", "hour", "weekday", "day_of_month", "month"] + ).transform(table) + torch.testing.assert_close( + output.numerical[:, 5:], + torch.tensor( + [ + [59, 23, 6, 30, 11], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [float("nan")] * 5, + ], + device=device, + ), + equal_nan=True, + ) + assert output.datetime_units == ("us", "ns") + assert output.select_stypes("datetime").to_arrow().equals(arrow) + assert output.numerical[0, :5].tolist() == [ + 0, + 0, + datetime(3000, 1, 1).weekday(), + 0, + 0, + ] + + +@pytest.mark.parametrize("base", [0, 1_704_067_200_000_000_000]) +def test_relative_day_features_preserve_nanosecond_sign_and_legacy_arithmetic( + base: int, +) -> None: + value = torch.tensor([[base - 1, base + 1]]) + seed = torch.tensor([[base // 1_000]]) + actual = relative_days(value, ("ns", "ns"), seed, ("us",)) + torch.testing.assert_close( + actual, + torch.tensor([[[1 / 86_400_000_000_000], [-1 / 86_400_000_000_000]]]), + rtol=1e-6, + atol=0, + ) + legacy = torch.tensor([[1_700_000_000_000_000]]) + later = legacy + 86_400_000_000 + assert relative_days(legacy, ("us",), later, ("us",)).equal( + (later.unsqueeze(-2) - legacy.unsqueeze(-1)) / 86_400_000_000 + ) + + +@onlyCUDA +def test_cudf_datetime_precision_roundtrip() -> None: + cudf = pytest.importorskip("cudf") + frame = cudf.DataFrame( + { + "precise": cudf.Series([1, -1, None], dtype="datetime64[ns]"), + "aligned": cudf.Series( + [1_000, -1_000, None], dtype="datetime64[ns]" + ), + } + ) + table = TableTensor.from_cudf( + frame, {"precise": "datetime", "aligned": "datetime"} + ) + assert table.datetime_units == ("ns", "us") + result = table.to_cudf() + assert result["precise"].equals(frame["precise"]) + assert result["aligned"].astype("datetime64[ns]").equals(frame["aligned"]) + # Existing string-to-datetime ingress remains supported. + dates = cudf.DataFrame({"date": ["2024-01-01", "2024-01-02"]}) + converted = TableTensor.from_cudf(dates, {"date": "datetime"}) + assert converted.datetime_units == ("us",) + + +@onlyCUDA +def test_rejected_cuda_conversion_keeps_device_usable() -> None: + wide = TableTensor( + datetime=torch.tensor([[32_503_680_000_000_000]], device="cuda") + ) + precise = TableTensor( + datetime=torch.tensor([[1]], device="cuda"), datetime_units=("ns",) + ) + with pytest.raises(ValueError, match="nanosecond timestamp range"): + torch.cat([wide, precise]) + legacy = TableTensor(datetime=torch.tensor([[2]], device="cuda")) + combined = cast(TableTensor, torch.cat([legacy, precise])) + assert combined.datetime_units == ("ns",) + assert combined.datetime.flatten().tolist() == [2_000, 1] + assert (torch.ones(1, device="cuda") + 1).item() == 2 + + +def test_compiled_same_unit_table_operation_preserves_precision( + compilation: None, +) -> None: + table = TableTensor( + datetime=torch.tensor([[-1], [1]]), datetime_units=("ns",) + ) + operation = torch.compile( + lambda table: table.unsqueeze(0), + backend="eager", + fullgraph=True, + ) + output = operation(table) + assert output.datetime_units == ("ns",) + assert output[0].to_arrow().equals(table.to_arrow())