Preserve nanosecond datetime precision with legacy microsecond defaults

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,

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:]
         ),
     ]

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:

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,
                 )

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
 

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)

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)

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)

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__(

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())