diff --git a/kumo/connector/connector.py b/kumo/connector/connector.py index 7d90762f6..ac10da613 100644 --- a/kumo/connector/connector.py +++ b/kumo/connector/connector.py @@ -107,6 +107,7 @@ def _compute_col_stats_from_sample( *, approx: bool, data_version: Any, + timestamp_precision: str = 'seconds', ) -> Dict[BaseColumnConfig, Dict[ColStatType, Stats]]: r"""Reduce an already-bounded sample without another warehouse read.""" out: Dict[BaseColumnConfig, Dict[ColStatType, Stats]] = {} @@ -121,7 +122,8 @@ def _compute_col_stats_from_sample( # A ``timestamp``-stype column read raw (cast_response=False) can # arrive as a bare epoch integer/float; coerce it to datetime so # the MIN/MAX col-stats are typed ``pd.Timestamp`` values. - ser = coerce_timestamp_stype_series_for_stats(ser, col_cfg) + ser = coerce_timestamp_stype_series_for_stats(ser, col_cfg, + timestamp_precision) # A warehouse decimal arrives as an object column of `decimal.Decimal`, # which the gate below reads as `string` and strips every numerical # stat from -- leaving a `numerical` column without the stats diff --git a/kumo/connector/timestamp_convert.py b/kumo/connector/timestamp_convert.py index cd0a56822..2e68162ef 100644 --- a/kumo/connector/timestamp_convert.py +++ b/kumo/connector/timestamp_convert.py @@ -13,10 +13,11 @@ interpreted. What remains on this side is the standardized output model: - * a timestamp collapses to an ``i64`` epoch-seconds value (UTC); + * a timestamp collapses to an ``i64`` epoch value (UTC) in ticks of an index + clock, seconds unless a snapshot's clock is given (:func:`to_epoch_ticks`); * a date collapses to an ``i32`` days-since-epoch value. -Sub-second precision and timezone are never retained, and everything is computed -in the integer epoch-seconds / epoch-days domain; only +Timezone and precision finer than the clock are never retained, and everything +is computed in the integer epoch-ticks / epoch-days domain; only :func:`epoch_seconds_to_timestamp` lifts the final scalar back into a ``pd.Timestamp`` for the stats value. @@ -44,6 +45,15 @@ _NATIVE_UNIT_DIVISOR = {'s': 1, 'ms': 1_000, 'us': 1_000_000, 'ns': 1_000_000_000} +PANDAS_UNIT_OF_INDEX_CLOCK = { + 'nanoseconds': 'ns', + 'microseconds': 'us', + 'milliseconds': 'ms', + 'seconds': 's', + 'hours': 'h', + 'days': 'D', +} + # --------------------------------------------------------------------------- # Core integer converters. @@ -94,15 +104,17 @@ def days_from_secs(secs: int) -> int: # --------------------------------------------------------------------------- # The DiskGraph crossing: one call per column, no per-value Python work. # --------------------------------------------------------------------------- -def _resolve(series: pd.Series, hint: TemporalHint, *, - days: bool) -> np.ndarray: - r"""Apply ``hint`` to a whole column, returning epoch seconds or epoch days. +def _resolve(series: pd.Series, hint: TemporalHint, *, days: bool, + timestamp_precision: str = 'seconds') -> np.ndarray: + r"""Apply ``hint`` to a whole column, returning epoch days or epoch ticks + of the ``timestamp_precision`` clock. An ``UNRESOLVED`` column declares seconds rather than declaring nothing: a declared hint is what stops DiskGraph re-detecting the column at apply time. """ unit = hint.unit or ('s' if hint.kind is HintKind.UNRESOLVED else None) - epochs, _, _ = resolve_temporal(_as_arrow(series), hint.strftime, unit, days) + epochs, _, _ = resolve_temporal(_as_arrow(series), hint.strftime, unit, + days, timestamp_unit=timestamp_precision) out = np.empty(len(epochs), dtype=object) out[:] = epochs return out @@ -146,13 +158,15 @@ def _native_divisor(dtype) -> int: return _NATIVE_UNIT_DIVISOR.get(unit, 1_000_000_000) -def timestamp_secs_from_array(series: pd.Series) -> np.ndarray: - r"""Convert a native temporal series to epoch seconds. +def timestamp_ticks_from_array( + series: pd.Series, timestamp_precision: str = 'seconds') -> np.ndarray: + r"""Convert a native temporal series to ticks of ``timestamp_precision``. - Reads the series' own native unit and ceiling-divides each value to seconds - (second resolution passes through, ms/us/ns go through the ceiling), so even a - native timestamp is ceiling-converted. Null (``NaT``) values map to ``None``. - Returns a numpy object array aligned with the input positions. + Reads the series' own native unit and ceiling-divides each value into a + coarser clock (the same unit passes through), so even a native timestamp is + ceiling-converted; a finer clock multiplies, saturating at the int64 range. + Null (``NaT``) values map to ``None``. Returns a numpy object array aligned + with the input positions. Fully vectorized: the ceiling is ``-((-a) // divisor)`` on the int64 view (floor-division floors toward negative infinity, so negating the floor of the @@ -162,7 +176,9 @@ def timestamp_secs_from_array(series: pd.Series) -> np.ndarray: becomes a bogus epoch. """ s = _as_naive_utc_datetime_series(series) - divisor = _native_divisor(s.dtype) + native_ns = 1_000_000_000 // _native_divisor(s.dtype) + tick_ns = pd.Timedelta( + 1, unit=PANDAS_UNIT_OF_INDEX_CLOCK[timestamp_precision]).value ints = s.to_numpy().view('int64') isna = pd.isna(s).to_numpy() # An object dtype array initializes to None, which is the NaT/null result. @@ -171,11 +187,17 @@ def timestamp_secs_from_array(series: pd.Series) -> np.ndarray: if not computable.any(): return out sub = ints[computable] # excludes NaT (the minimum-int64 sentinel) - if divisor == 1: - secs = sub # Second resolution: already whole epoch seconds + if tick_ns == native_ns: + ticks = sub + elif tick_ns > native_ns: + ticks = -((-sub) // (tick_ns // native_ns)) else: - secs = -((-sub) // divisor) - out[computable] = secs.tolist() # store Python ints, leave nulls as None + factor = native_ns // tick_ns + limit = np.iinfo(np.int64).max // factor + ticks = np.where( + sub > limit, np.iinfo(np.int64).max, + np.where(sub < -limit, np.iinfo(np.int64).min, sub * factor)) + out[computable] = ticks.tolist() # store Python ints, leave nulls as None return out @@ -190,9 +212,18 @@ def to_epoch_seconds(series: pd.Series, hint: TemporalHint) -> np.ndarray: unconvertible or null value becomes ``None``, never raising. Returns a numpy object array of ints and ``None`` aligned with the input. """ + return to_epoch_ticks(series, hint, 'seconds') + + +def to_epoch_ticks(series: pd.Series, hint: TemporalHint, + timestamp_precision: str) -> np.ndarray: + r"""Like :func:`to_epoch_seconds`, in ticks of the ``timestamp_precision`` + index clock. + """ if hint.kind is HintKind.NATIVE or is_datetime64_any_dtype(series): - return timestamp_secs_from_array(series) - return _resolve(series, hint, days=False) + return timestamp_ticks_from_array(series, timestamp_precision) + return _resolve(series, hint, days=False, + timestamp_precision=timestamp_precision) def to_epoch_days(series: pd.Series, hint: TemporalHint) -> np.ndarray: @@ -201,7 +232,7 @@ def to_epoch_days(series: pd.Series, hint: TemporalHint) -> np.ndarray: Best-effort: unconvertible or null values become ``None``. """ if hint.kind is HintKind.NATIVE or is_datetime64_any_dtype(series): - secs = timestamp_secs_from_array(series) + secs = timestamp_ticks_from_array(series) return np.array( [None if x is None else days_from_secs(int(x)) for x in secs], dtype=object) diff --git a/kumo/connector/utils.py b/kumo/connector/utils.py index 99a5261a9..0e9eedacb 100644 --- a/kumo/connector/utils.py +++ b/kumo/connector/utils.py @@ -26,7 +26,12 @@ from io_utils import fsutil, kumo_dataclass, logging from kumo.config import BaseColumnConfig -from kumo.connector.timestamp_convert import to_epoch_days, to_epoch_seconds +from kumo.connector.timestamp_convert import ( + PANDAS_UNIT_OF_INDEX_CLOCK, + to_epoch_days, + to_epoch_seconds, + to_epoch_ticks, +) from kumo.connector.timestamp_infer import ( HintKind, TemporalHint, @@ -372,6 +377,7 @@ def name_inference_temporal_supported( def coerce_timestamp_stype_series_for_stats( ser: pd.Series, col_config: BaseColumnConfig, + timestamp_precision: str = 'seconds', ) -> pd.Series: r"""Coerce a ``timestamp``-stype series to a datetime series before the Python column-stat kernel reduces it. @@ -385,11 +391,11 @@ def coerce_timestamp_stype_series_for_stats( The coercion resolves the column's :class:`TemporalHint` (:func:`_resolve_stats_temporal_hint`) and applies - :func:`~kumo.connector.timestamp_convert.to_epoch_seconds`, then wraps the - epoch-seconds back into a second-precision ``datetime64`` series. Because every - temporal column collapses to epoch-seconds with ceiling division, the resulting - MIN/MAX (and the other timestamp stats) are computed in the same epoch-seconds - domain the on-disk index uses and match it rather than drifting. + :func:`~kumo.connector.timestamp_convert.to_epoch_ticks` in the + ``timestamp_precision`` index clock, then wraps the ticks back into a + ``datetime64`` series of that precision. The resulting MIN/MAX (and the + other timestamp stats) are computed in the same domain the on-disk index + uses and match it rather than drifting. Best-effort and narrowly scoped: * non-``timestamp``-stype columns pass through untouched; @@ -406,8 +412,9 @@ def coerce_timestamp_stype_series_for_stats( return ser hint = _resolve_stats_temporal_hint(ser, col_config) - epoch_seconds = to_epoch_seconds(ser, hint) - return pd.to_datetime(pd.Series(epoch_seconds, index=ser.index), unit='s', + ticks = to_epoch_ticks(ser, hint, timestamp_precision) + return pd.to_datetime(pd.Series(ticks, index=ser.index), + unit=PANDAS_UNIT_OF_INDEX_CLOCK[timestamp_precision], errors='coerce') diff --git a/kumo/data/generic/diskgraph_batch_pred.py b/kumo/data/generic/diskgraph_batch_pred.py index 56e71d445..174237bf3 100644 --- a/kumo/data/generic/diskgraph_batch_pred.py +++ b/kumo/data/generic/diskgraph_batch_pred.py @@ -110,7 +110,7 @@ def generate_batch_pred_table_via_diskgraph( raise RuntimeError('diskgraph index path is required') with kumo_timing("batch_pred_table.create_plan_compiler"): - compiler = DiskGraphPlanCompiler() + compiler = DiskGraphPlanCompiler(index_path=index_path) # Determine execution mode if task_type == core_types.TaskType.FORECASTING: @@ -756,7 +756,7 @@ def generate_batch_pred_table_via_diskgraph_ttng( if not index_path: raise RuntimeError('diskgraph index path is required') - compiler = DiskGraphPlanCompiler() + compiler = DiskGraphPlanCompiler(index_path=index_path) if task_type == core_types.TaskType.FORECASTING: return _execute_forecasting_ttng( @@ -804,6 +804,7 @@ def _execute_node_prediction_ttng( # Node-pred: no lp_resolution; tt-ingest synthesizes ENTITY1 -> entity_node_id. plan_json = build_ttng_predict_plan( query, + index_path=index_path, task_type=task_type, graph=graph, timestamps=timestamps, @@ -927,6 +928,7 @@ def _execute_link_prediction_ttng( lhs_plan = build_ttng_predict_plan( original_query, + index_path=index_path, task_type=task_type, graph=graph, timestamps=timestamps, @@ -1141,6 +1143,7 @@ def _rename_dg_to_kumo(batch: pa.RecordBatch) -> pa.RecordBatch: def build_ttng_predict_plan( query: ValidatedPredictiveQuery, *, + index_path: str, task_type: core_types.TaskType, graph: Graph, timestamps: List[int], @@ -1157,7 +1160,7 @@ def build_ttng_predict_plan( from diskgraph import DiskGraphPlanCompiler from kumo_core.training.table import apply_reference_resolution - compiler = DiskGraphPlanCompiler() + compiler = DiskGraphPlanCompiler(index_path=index_path) tt_str = task_type.value is_link_pred = bool(getattr(task_type, 'is_link_pred', False)) diff --git a/kumo/data/generic/diskgraph_train_table.py b/kumo/data/generic/diskgraph_train_table.py index c8e7ddfec..e052d6427 100644 --- a/kumo/data/generic/diskgraph_train_table.py +++ b/kumo/data/generic/diskgraph_train_table.py @@ -757,7 +757,7 @@ def generate_train_table_via_diskgraph( if not fsutil.exists(index_path): raise RuntimeError(f'diskgraph index path does not exist: {index_path}') - compiler = DiskGraphPlanCompiler() + compiler = DiskGraphPlanCompiler(index_path=index_path) split_spec = _build_split_spec(split_config, split_family, SplitSpec) if split_spec is None: raise RuntimeError( diff --git a/kumo/diskgraph/ingest.py b/kumo/diskgraph/ingest.py index 543958cb3..1962bae34 100644 --- a/kumo/diskgraph/ingest.py +++ b/kumo/diskgraph/ingest.py @@ -18,6 +18,7 @@ from kumo_core.diskgraph.converter import convert_diskgraph_config from kumo_core.diskgraph.ingest import ( normalize_legacy_column_aliases, + set_timestamp_precision, validate_diskgraph_ingest_config, ) from kumo_core.diskgraph.runner import ( @@ -159,6 +160,7 @@ def run_diskgraph_ingest( source_specs: Optional[dict] = None, subprocess_environment: Optional[dict[str, Optional[str]]] = None, heartbeat: Optional[Callable[[], None]] = None, + timestamp_precision: str = 'seconds', ) -> DiskgraphIngestResult: """Convert the kumo graph config and launch the DiskGraph supervisor. @@ -190,6 +192,7 @@ def run_diskgraph_ingest( source_specs=source_specs, subprocess_environment=subprocess_environment, ) + set_timestamp_precision(ingest_config, timestamp_precision) validate_diskgraph_ingest_config(ingest_config) # DiskGraph owns phase detection, concurrency reduction after OOM, and diff --git a/kumo/diskgraph/service.py b/kumo/diskgraph/service.py index 0adc700a7..f230bbd49 100644 --- a/kumo/diskgraph/service.py +++ b/kumo/diskgraph/service.py @@ -660,6 +660,7 @@ def ingest( data_persist_dir: str, table_data_versions: dict, dataplane_root: str, + timestamp_precision: str = 'seconds', source_paths: Optional[dict] = None, source_specs: Optional[dict] = None, subprocess_environment: Optional[dict[str, Optional[str]]] = None, @@ -689,6 +690,7 @@ def ingest( graph_config=graph_config, data_persist_dir=data_persist_dir, table_data_versions=table_data_versions, + timestamp_precision=timestamp_precision, output_dir=output_dir, source_paths=source_paths, source_specs=source_specs, @@ -1721,6 +1723,7 @@ def compute_col_stats_for_parquet( parquet_paths: List[str], kumo_column_specs: List[dict], data_version: str, + timestamp_precision: str = 'seconds', ) -> Tuple[str, str]: """Compute Kumo-format column statistics over BP-output parquet. @@ -1733,6 +1736,7 @@ def compute_col_stats_for_parquet( ``name`` / ``stype`` / ``dtype`` (raw Kumo enum value strings); pybind maps them via ``map_kumo_stype`` / ``map_kumo_dtype``. data_version: Stamped into every ``kumo_col_stats_v1`` entry. + timestamp_precision: Index clock that timestamp statistics are in. Returns: ``(full_json, min_json)`` — the full map and its min variant. @@ -1742,6 +1746,7 @@ def compute_col_stats_for_parquet( parquet_paths=parquet_paths, kumo_column_specs_json=json.dumps(kumo_column_specs), data_version=data_version, + timestamp_unit=timestamp_precision, ) diff --git a/packages/diskgraph/Cargo.lock b/packages/diskgraph/Cargo.lock index be75fbfdd..216a17da4 100644 --- a/packages/diskgraph/Cargo.lock +++ b/packages/diskgraph/Cargo.lock @@ -1853,6 +1853,7 @@ dependencies = [ "base64 0.22.1", "bincode", "chrono", + "diskgraph-core", "fast-float2", "gxhash", "half", diff --git a/packages/diskgraph/diskgraph-core/src/metadata.rs b/packages/diskgraph/diskgraph-core/src/metadata.rs index c10793774..313b0128f 100644 --- a/packages/diskgraph/diskgraph-core/src/metadata.rs +++ b/packages/diskgraph/diskgraph-core/src/metadata.rs @@ -130,6 +130,7 @@ pub enum EdgeSortKeyTimestampKind { Int8, Int16, Int32, + Int64, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/packages/diskgraph/diskgraph-core/src/sortkey.rs b/packages/diskgraph/diskgraph-core/src/sortkey.rs index f6e652f2e..214a94a8c 100644 --- a/packages/diskgraph/diskgraph-core/src/sortkey.rs +++ b/packages/diskgraph/diskgraph-core/src/sortkey.rs @@ -37,6 +37,7 @@ impl SortKeyDecoder { EdgeSortKeyTimestampKind::Int8 => 8, EdgeSortKeyTimestampKind::Int16 => 16, EdgeSortKeyTimestampKind::Int32 => 32, + EdgeSortKeyTimestampKind::Int64 => 64, }) .unwrap_or(0); @@ -231,6 +232,7 @@ impl SortKeyEncoder { EdgeSortKeyTimestampKind::Int8 => 8, EdgeSortKeyTimestampKind::Int16 => 16, EdgeSortKeyTimestampKind::Int32 => 32, + EdgeSortKeyTimestampKind::Int64 => 64, }) .unwrap_or(0); @@ -853,4 +855,36 @@ mod tests { "encode_upper(30) incorrectly includes ts=30" ); } + + #[test] + fn int64_timestamp_keys_round_trip_in_order() { + let meta = EdgeSortKeyMetadata { + width_bits: 64, + timestamp: Some(EdgeSortKeyTimestampMetadata { + table: "t".into(), + column: "c".into(), + encoding: EdgeSortKeyTimestampEncoding { + kind: EdgeSortKeyTimestampKind::Int64, + anchor: 0, + step: 1, + }, + }), + weight: None, + }; + let encoder = SortKeyEncoder::new(&meta); + let decoder = SortKeyDecoder::new(&meta); + let times = [ + -1_000_000_001, + -1, + 0, + 1, + 1_700_000_000_123_456_789, + i64::MAX, + ]; + let keys: Vec = times.iter().map(|&t| encoder.encode_lower(t)).collect(); + assert!(keys.windows(2).all(|pair| pair[0] < pair[1])); + for (time, key) in times.iter().zip(keys) { + assert_eq!(decoder.decode(key).0, Some(*time)); + } + } } diff --git a/packages/diskgraph/diskgraph-core/src/time.rs b/packages/diskgraph/diskgraph-core/src/time.rs index 8fe8c1e4b..bf237799f 100644 --- a/packages/diskgraph/diskgraph-core/src/time.rs +++ b/packages/diskgraph/diskgraph-core/src/time.rs @@ -26,7 +26,46 @@ pub enum TimeUnit { Years, } +impl From for TimeUnit { + fn from(unit: arrow_schema::TimeUnit) -> Self { + match unit { + arrow_schema::TimeUnit::Second => Self::Seconds, + arrow_schema::TimeUnit::Millisecond => Self::Milliseconds, + arrow_schema::TimeUnit::Microsecond => Self::Microseconds, + arrow_schema::TimeUnit::Nanosecond => Self::Nanoseconds, + } + } +} + impl TimeUnit { + /// The units an index can store its timestamps in. + pub const INDEX_CLOCKS: [TimeUnit; 6] = [ + Self::Nanoseconds, + Self::Microseconds, + Self::Milliseconds, + Self::Seconds, + Self::Hours, + Self::Days, + ]; + + pub fn is_index_clock(self) -> bool { + Self::INDEX_CLOCKS.contains(&self) + } + + pub const fn tick_nanos(self) -> i64 { + match self { + Self::Nanoseconds => 1, + Self::Microseconds => 1_000, + Self::Milliseconds => 1_000_000, + Self::Seconds => NS_PER_SEC, + Self::Minutes => 60 * NS_PER_SEC, + Self::Hours => 3_600 * NS_PER_SEC, + Self::Days => 86_400 * NS_PER_SEC, + Self::Weeks => 604_800 * NS_PER_SEC, + Self::Months | Self::Years => panic!("calendar units have no fixed tick"), + } + } + /// Parse a unit name string (used by builder methods that still accept strings). pub fn from_name(s: &str) -> Option { match s { @@ -45,6 +84,109 @@ impl TimeUnit { } } +// --------------------------------------------------------------------------- +// Converting raw values between fixed units. Each rounding mode is a rule that +// `main` applies to some kind of value; coarse-to-fine conversions saturate. +// --------------------------------------------------------------------------- + +pub fn convert_trunc(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { + let (from, to) = (from.tick_nanos(), to.tick_nanos()); + if from >= to { + value.saturating_mul(from / to) + } else { + value / (to / from) + } +} + +pub fn convert_floor(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { + let (from, to) = (from.tick_nanos(), to.tick_nanos()); + if from >= to { + value.saturating_mul(from / to) + } else { + value.div_euclid(to / from) + } +} + +pub fn convert_ceil(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { + let (from, to) = (from.tick_nanos(), to.tick_nanos()); + if from >= to { + value.saturating_mul(from / to) + } else { + div_ceil(value, to / from) + } +} + +#[inline] +pub fn div_ceil(value: i64, divisor: i64) -> i64 { + value.div_euclid(divisor) + i64::from(value.rem_euclid(divisor) != 0) +} + +// The conversions on the ingest hot path are generic over the clock and match +// the source unit arm by arm, so every factor is a compile-time constant (for +// the seconds clock, exactly main's hard-coded divisions). + +/// Converts `value` in `unit` to ticks of the clock `U`, rounding up; `None` +/// when the clock cannot represent it. +#[inline(always)] +pub fn ceil_to_clock(value: i64, unit: TimeUnit) -> Option { + match unit { + TimeUnit::Nanoseconds => ceil_from_tick::(value, 1), + TimeUnit::Microseconds => ceil_from_tick::(value, 1_000), + TimeUnit::Milliseconds => ceil_from_tick::(value, 1_000_000), + TimeUnit::Seconds => ceil_from_tick::(value, NS_PER_SEC), + other => ceil_from_tick::(value, other.tick_nanos()), + } +} + +#[inline(always)] +fn ceil_from_tick(value: i64, tick: i64) -> Option { + let clock_tick = U::TICK_NANOS; + if tick >= clock_tick { + value.checked_mul(tick / clock_tick) + } else { + Some(div_ceil(value, clock_tick / tick)) + } +} + +/// Floors `secs + nanos()` (with `0 <= nanos() < 1e9`) to a tick of the clock +/// `U`, or `None` when the clock cannot represent it; `nanos` is only +/// evaluated when a tick is shorter than a second. +#[inline(always)] +pub fn floor_parts(secs: i64, nanos: impl FnOnce() -> u32) -> Option { + let clock_tick = U::TICK_NANOS; + if clock_tick < NS_PER_SEC { + let fraction = i64::from(nanos()) / clock_tick; + secs.checked_mul(NS_PER_SEC / clock_tick)? + .checked_add(fraction) + } else { + Some(secs.div_euclid(clock_tick / NS_PER_SEC)) + } +} + +/// Rounds `secs + nanos` (with `0 <= nanos < 1e9`) up to a tick of `to`. +pub fn ceil_parts(secs: i64, nanos: u32, to: TimeUnit) -> i64 { + if to.tick_nanos() < NS_PER_SEC { + let fraction = convert_ceil(i64::from(nanos), TimeUnit::Nanoseconds, to); + convert_floor(secs, TimeUnit::Seconds, to).saturating_add(fraction) + } else { + convert_ceil( + secs.saturating_add(i64::from(nanos > 0)), + TimeUnit::Seconds, + to, + ) + } +} + +/// Splits a value in `unit` into whole epoch seconds and a nanosecond fraction. +pub fn split_parts(value: i64, unit: TimeUnit) -> (i64, u32) { + let ticks_per_sec = NS_PER_SEC / unit.tick_nanos(); + if ticks_per_sec == 0 { + return (convert_floor(value, unit, TimeUnit::Seconds), 0); + } + let nanos = value.rem_euclid(ticks_per_sec) * unit.tick_nanos(); + (value.div_euclid(ticks_per_sec), nanos as u32) +} + // --------------------------------------------------------------------------- // UnitSystem trait + Abs / Dur phantom-typed wrappers // --------------------------------------------------------------------------- @@ -53,13 +195,16 @@ impl TimeUnit { /// /// Implementations provide construction from various source units and /// extraction back to raw i64. Current implementations: -/// - [`Seconds`] — epoch seconds, used by rowpack storage and the core engine -/// - [`Nanoseconds`] — epoch nanoseconds, used at the PyO3/Arrow boundary +/// - [`Seconds`] — epoch seconds, the default index clock +/// - [`Nanoseconds`] — epoch nanoseconds, also used at the PyO3/Arrow boundary /// where timestamps arrive as `TimestampNanosecondArray` +/// - [`Microseconds`], [`Milliseconds`], [`Hours`], [`Days`] — the other index clocks /// /// Adding a new unit requires a new marker struct + impl block + one match arm /// at the top-level dispatch site. pub trait UnitSystem: 'static + Copy + Send + Sync + Eq + Ord + std::hash::Hash { + const UNIT: TimeUnit; + const TICK_NANOS: i64 = Self::UNIT.tick_nanos(); fn abs_from_secs(secs: i64) -> Abs; fn abs_from_millis(ms: i64) -> Abs; fn abs_from_micros(us: i64) -> Abs; @@ -89,6 +234,10 @@ pub struct Dur(i64, PhantomData); // -- Abs convenience methods (delegate to UnitSystem) -------------------- impl Abs { + /// `self + dur`, or `None` when the result leaves the clock. + pub fn checked_add(self, dur: Dur) -> Option { + self.0.checked_add(dur.0).map(Abs::from_raw) + } pub fn from_secs(secs: i64) -> Self { U::abs_from_secs(secs) } @@ -243,13 +392,12 @@ impl<'de, U: UnitSystem> Deserialize<'de> for Dur { // -- Seconds unit system ---------------------------------------------------- /// The seconds unit system. `Abs` stores epoch seconds, `Dur` -/// stores duration in seconds. This is the only implementation today — it -/// matches the rowpack storage contract where `access.get_timestamp()` always -/// returns epoch seconds. +/// stores duration in seconds. This is the default index clock. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Seconds; impl UnitSystem for Seconds { + const UNIT: TimeUnit = TimeUnit::Seconds; fn abs_from_secs(secs: i64) -> Abs { Abs(secs, PhantomData) } @@ -299,6 +447,7 @@ pub struct Nanoseconds; const NS_PER_SEC: i64 = 1_000_000_000; impl UnitSystem for Nanoseconds { + const UNIT: TimeUnit = TimeUnit::Nanoseconds; fn abs_from_secs(secs: i64) -> Abs { Abs(secs.saturating_mul(NS_PER_SEC), PhantomData) } @@ -337,6 +486,95 @@ impl UnitSystem for Nanoseconds { } } +// -- Microseconds, milliseconds, hours and days ------------------------------ + +macro_rules! fixed_clock { + ($name:ident) => { + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] + pub struct $name; + + impl UnitSystem for $name { + const UNIT: TimeUnit = TimeUnit::$name; + fn abs_from_secs(secs: i64) -> Abs { + Abs::from_raw(convert_trunc(secs, TimeUnit::Seconds, Self::UNIT)) + } + fn abs_from_millis(ms: i64) -> Abs { + Abs::from_raw(convert_trunc(ms, TimeUnit::Milliseconds, Self::UNIT)) + } + fn abs_from_micros(us: i64) -> Abs { + Abs::from_raw(convert_trunc(us, TimeUnit::Microseconds, Self::UNIT)) + } + fn abs_from_nanos(ns: i64) -> Abs { + Abs::from_raw(convert_trunc(ns, TimeUnit::Nanoseconds, Self::UNIT)) + } + fn abs_from_days(days: i32) -> Abs { + Abs::from_raw(convert_trunc(i64::from(days), TimeUnit::Days, Self::UNIT)) + } + fn abs_raw(ts: Abs) -> i64 { + ts.0 + } + fn abs_to_secs(ts: Abs) -> i64 { + convert_trunc(ts.0, Self::UNIT, TimeUnit::Seconds) + } + fn dur_from_secs(secs: i64) -> Dur { + Dur::from_raw(convert_trunc(secs, TimeUnit::Seconds, Self::UNIT)) + } + fn dur_from_nanos(ns: i64) -> Dur { + Dur::from_raw(convert_trunc(ns, TimeUnit::Nanoseconds, Self::UNIT)) + } + fn dur_from_days(days: i64) -> Dur { + Dur::from_raw(convert_trunc(days, TimeUnit::Days, Self::UNIT)) + } + fn dur_raw(dur: Dur) -> i64 { + dur.0 + } + fn abs_to_nanos(ts: Abs) -> i64 { + convert_trunc(ts.0, Self::UNIT, TimeUnit::Nanoseconds) + } + } + }; +} + +fixed_clock!(Microseconds); +fixed_clock!(Milliseconds); +fixed_clock!(Hours); +fixed_clock!(Days); + +/// Runs `$body` with `$U` bound to the `UnitSystem` of an index clock. +/// Callers validate the clock with [`TimeUnit::is_index_clock`] first. +#[macro_export] +macro_rules! with_index_clock { + ($clock:expr, $U:ident => $body:expr) => { + match $clock { + $crate::time::TimeUnit::Nanoseconds => { + type $U = $crate::time::Nanoseconds; + $body + } + $crate::time::TimeUnit::Microseconds => { + type $U = $crate::time::Microseconds; + $body + } + $crate::time::TimeUnit::Milliseconds => { + type $U = $crate::time::Milliseconds; + $body + } + $crate::time::TimeUnit::Seconds => { + type $U = $crate::time::Seconds; + $body + } + $crate::time::TimeUnit::Hours => { + type $U = $crate::time::Hours; + $body + } + $crate::time::TimeUnit::Days => { + type $U = $crate::time::Days; + $body + } + other => unreachable!("{:?} is not an index clock", other), + } + }; +} + // -- Helpers for converting from a dynamic TimeUnit ------------------------- /// Create an `Abs` from a raw value whose unit is given dynamically. @@ -371,13 +609,12 @@ pub fn dur_from_unit(value: i64, unit: TimeUnit) -> Dur { /// Add calendar months to a timestamp. Uses chrono for correct month /// boundaries (e.g. Jan 31 + 1 month = Feb 28). Matches /// `pd.DateOffset(months=N)` behavior. -pub fn add_calendar_months(ts: Abs, months: i32) -> Abs { +pub fn add_calendar_months(ts: Abs, months: i32) -> Option> { use chrono::{DateTime, Datelike, NaiveDate, Timelike}; let secs = ts.to_secs(); - let dt = DateTime::from_timestamp(secs, 0) - .expect("timestamp out of range for calendar arithmetic") - .naive_utc(); + let subsecond = ts - Abs::::from_secs(secs); + let dt = DateTime::from_timestamp(secs, 0)?.naive_utc(); let total_months = dt.month0() as i32 + months; // Euclidean division to handle negative months correctly @@ -388,12 +625,10 @@ pub fn add_calendar_months(ts: Abs, months: i32) -> Abs { let max_day = days_in_month(target_year, target_month); let day = dt.day().min(max_day); - let result = NaiveDate::from_ymd_opt(target_year, target_month, day) - .expect("invalid date in calendar arithmetic") - .and_hms_opt(dt.hour(), dt.minute(), dt.second()) - .expect("invalid time in calendar arithmetic"); + let date = NaiveDate::from_ymd_opt(target_year, target_month, day)?; + let result = date.and_hms_opt(dt.hour(), dt.minute(), dt.second())?; - Abs::::from_secs(result.and_utc().timestamp()) + Abs::from_raw(floor_parts::(result.and_utc().timestamp(), || 0)?).checked_add(subsecond) } /// Add an offset to a timestamp. For fixed-size units (seconds through weeks) @@ -402,11 +637,11 @@ pub fn add_calendar_months(ts: Abs, months: i32) -> Abs { /// /// This is the typed replacement for the ad-hoc `add_unit` that operated on /// raw nanosecond i64 values. -pub fn add_offset(ts: Abs, unit: TimeUnit, value: i64) -> Abs { +pub fn add_offset(ts: Abs, unit: TimeUnit, value: i64) -> Option> { match unit { TimeUnit::Months => add_calendar_months(ts, value as i32), TimeUnit::Years => add_calendar_months(ts, (value * 12) as i32), - _ => ts + dur_from_unit(value, unit), + _ => ts.checked_add(dur_from_unit(value, unit)), } } @@ -709,7 +944,7 @@ mod tests { fn add_one_month_from_jan() { // 2024-01-15 00:00:00 UTC let ts = Abs::::from_secs(1_705_276_800); - let result = add_calendar_months(ts, 1); + let result = add_calendar_months(ts, 1).unwrap(); // 2024-02-15 00:00:00 UTC assert_eq!(result.raw(), 1_707_955_200); } @@ -718,7 +953,7 @@ mod tests { fn add_month_clamps_day_jan31_to_feb28() { // 2023-01-31 00:00:00 UTC = 1675123200 let ts = Abs::::from_secs(1_675_123_200); - let result = add_calendar_months(ts, 1); + let result = add_calendar_months(ts, 1).unwrap(); // 2023-02-28 00:00:00 UTC = 1677542400 assert_eq!(result.raw(), 1_677_542_400); } @@ -727,7 +962,7 @@ mod tests { fn subtract_one_month() { // 2024-03-15 00:00:00 UTC = 1710460800 let ts = Abs::::from_secs(1_710_460_800); - let result = add_calendar_months(ts, -1); + let result = add_calendar_months(ts, -1).unwrap(); // 2024-02-15 00:00:00 UTC = 1707955200 assert_eq!(result.raw(), 1_707_955_200); } @@ -736,7 +971,7 @@ mod tests { fn add_twelve_months_equals_one_year() { // 2024-06-01 00:00:00 UTC = 1717200000 let ts = Abs::::from_secs(1_717_200_000); - let result = add_calendar_months(ts, 12); + let result = add_calendar_months(ts, 12).unwrap(); // 2025-06-01 00:00:00 UTC = 1748736000 assert_eq!(result.raw(), 1_748_736_000); } @@ -858,7 +1093,7 @@ mod tests { #[test] fn add_offset_days_nanoseconds() { let ts = Abs::::from_secs(1000); - let result = add_offset(ts, TimeUnit::Days, 2); + let result = add_offset(ts, TimeUnit::Days, 2).unwrap(); assert_eq!(result.raw(), ts.raw() + 2 * 86_400_000_000_000); } @@ -866,7 +1101,7 @@ mod tests { fn add_offset_months_nanoseconds() { // 2024-01-15 00:00:00 UTC in nanos let ts = Abs::::from_secs(1_705_276_800); - let result = add_offset(ts, TimeUnit::Months, 1); + let result = add_offset(ts, TimeUnit::Months, 1).unwrap(); // 2024-02-15 00:00:00 UTC in nanos let expected = Abs::::from_secs(1_707_955_200); assert_eq!(result.raw(), expected.raw()); @@ -876,7 +1111,7 @@ mod tests { fn add_offset_years_nanoseconds() { // 2024-01-15 00:00:00 UTC in nanos let ts = Abs::::from_secs(1_705_276_800); - let result = add_offset(ts, TimeUnit::Years, 1); + let result = add_offset(ts, TimeUnit::Years, 1).unwrap(); // 2025-01-15 00:00:00 UTC in nanos let expected = Abs::::from_secs(1_736_899_200); assert_eq!(result.raw(), expected.raw()); @@ -898,6 +1133,128 @@ mod tests { assert_eq!(TimeUnit::from_name("microseconds"), Some(TimeUnit::Microseconds)); } + // -- Conversions between fixed units -------------------------------------- + + #[test] + fn finer_to_coarser_rounds_by_mode() { + let ms = TimeUnit::Milliseconds; + let s = TimeUnit::Seconds; + assert_eq!(convert_ceil(1_500, ms, s), 2); + assert_eq!(convert_ceil(-1_500, ms, s), -1); + assert_eq!(convert_floor(1_500, ms, s), 1); + assert_eq!(convert_floor(-1_500, ms, s), -2); + assert_eq!(convert_trunc(1_500, ms, s), 1); + assert_eq!(convert_trunc(-1_500, ms, s), -1); + } + + #[test] + fn coarser_to_finer_multiplies_and_saturates() { + for convert in [convert_ceil, convert_floor, convert_trunc] { + assert_eq!(convert(2, TimeUnit::Hours, TimeUnit::Seconds), 7_200); + assert_eq!( + convert(i64::MAX / 2, TimeUnit::Seconds, TimeUnit::Nanoseconds), + i64::MAX + ); + assert_eq!( + convert(i64::MIN / 2, TimeUnit::Days, TimeUnit::Seconds), + i64::MIN + ); + assert_eq!(convert(-7, TimeUnit::Days, TimeUnit::Days), -7); + } + } + + #[test] + fn floor_parts_drops_the_fraction_below_one_tick() { + assert_eq!(floor_parts::(10, || 999_999_999), Some(10)); + assert_eq!( + floor_parts::(10, || 123_456_789), + Some(10_123) + ); + assert_eq!( + floor_parts::(10, || 123_456_789), + Some(10_123_456_789) + ); + assert_eq!(floor_parts::(-1, || 500_000_000), Some(-1)); + assert_eq!(floor_parts::(253_402_300_799, || 0), None); + } + + #[test] + fn ceil_to_clock_is_none_beyond_the_clock() { + assert_eq!( + ceil_to_clock::(253_402_300_799, TimeUnit::Seconds), + None + ); + assert_eq!( + ceil_to_clock::(i64::MAX, TimeUnit::Nanoseconds), + Some(9_223_372_037) + ); + } + + #[test] + fn ceil_to_clock_matches_the_dynamic_conversion() { + for value in [-1_500_000_001, -1, 0, 1, 1_705_312_800_750] { + for unit in [ + TimeUnit::Nanoseconds, + TimeUnit::Milliseconds, + TimeUnit::Seconds, + TimeUnit::Days, + ] { + for clock in TimeUnit::INDEX_CLOCKS { + let ticks = + crate::with_index_clock!(clock, U => ceil_to_clock::(value, unit)); + let saturated = convert_ceil(value, unit, clock); + let expected = Some(saturated).filter(|t| *t != i64::MAX && *t != i64::MIN); + assert_eq!(ticks, expected, "{unit:?} -> {clock:?}"); + } + } + } + } + + #[test] + fn ceil_parts_rounds_any_fraction_up_to_one_tick() { + assert_eq!(ceil_parts(10, 0, TimeUnit::Seconds), 10); + assert_eq!(ceil_parts(10, 1, TimeUnit::Seconds), 11); + assert_eq!(ceil_parts(10, 1, TimeUnit::Milliseconds), 10_001); + assert_eq!(ceil_parts(3_600, 0, TimeUnit::Hours), 1); + assert_eq!(ceil_parts(3_600, 1, TimeUnit::Hours), 2); + assert_eq!(ceil_parts(-1, 500_000_000, TimeUnit::Hours), 0); + } + + #[test] + fn split_parts_returns_seconds_and_a_positive_fraction() { + assert_eq!(split_parts(-1, TimeUnit::Nanoseconds), (-1, 999_999_999)); + assert_eq!( + split_parts(10_123, TimeUnit::Milliseconds), + (10, 123_000_000) + ); + assert_eq!(split_parts(2, TimeUnit::Hours), (7_200, 0)); + assert_eq!(split_parts(i64::MAX, TimeUnit::Days), (i64::MAX, 0)); + } + + #[test] + fn every_index_clock_has_a_unit_system() { + for clock in TimeUnit::INDEX_CLOCKS { + let day = crate::with_index_clock!(clock, U => Abs::::from_days(1).to_secs()); + assert_eq!(day, 86_400, "{clock:?}"); + } + } + + #[test] + fn calendar_months_keep_the_time_of_day_below_one_second() { + let ts = Abs::::from_raw(1_705_276_800_250); + assert_eq!(add_calendar_months(ts, 1).unwrap().raw(), 1_707_955_200_250); + } + + #[test] + fn calendar_months_beyond_the_clock_are_none() { + let last_ns = Abs::::from_raw(i64::MAX - 1); + assert_eq!(add_calendar_months(last_ns, 1), None); + assert_eq!( + add_calendar_months(Abs::::from_raw(i64::MAX), 1), + None + ); + } + #[test] fn time_unit_from_name_invalid() { assert_eq!(TimeUnit::from_name("invalid"), None); diff --git a/packages/diskgraph/diskgraph-core/src/tt_builder.rs b/packages/diskgraph/diskgraph-core/src/tt_builder.rs index ad317ca83..35153c750 100644 --- a/packages/diskgraph/diskgraph-core/src/tt_builder.rs +++ b/packages/diskgraph/diskgraph-core/src/tt_builder.rs @@ -537,6 +537,127 @@ pub enum StageFilter { Not(Box), } +// ─── Index clocks ──────────────────────────────────────────────────────────── +// +// A serialized plan stores raw ticks of one clock, named by a top-level +// `timestamp_unit` (omitted for seconds). + +impl DiskGraphTTPlan { + pub fn to_json(&self) -> serde_json::Result { + if U::UNIT == TimeUnit::Seconds { + return serde_json::to_string(self); + } + let mut json = serde_json::to_value(self)?; + json["timestamp_unit"] = serde_json::to_value(U::UNIT)?; + serde_json::to_string(&json) + } +} + +impl Op { + /// The same op in clock `V`, with every offset converted by `offset`. + pub fn map_clock(self, offset: &impl Fn(Dur) -> Dur) -> Op { + match self { + Op::Aggregate { + alias, + path, + func, + target_col, + time_col, + window, + window_unit, + row_filter, + limit, + target_entity_table, + } => Op::Aggregate { + alias, + path, + func, + target_col, + time_col, + window, + window_unit, + row_filter, + limit, + target_entity_table, + }, + Op::Column { + alias, + table, + column, + path, + } => Op::Column { + alias, + table, + column, + path, + }, + Op::Compare { + alias, + lhs, + op, + rhs, + } => Op::Compare { + alias, + lhs, + op, + rhs, + }, + Op::LogicalOp { + alias, + op, + operands, + } => Op::LogicalOp { + alias, + op, + operands, + }, + Op::Autoregression { + alias, + target_op, + lag_offsets, + } => Op::Autoregression { + alias, + target_op: Box::new(target_op.map_clock(offset)), + lag_offsets: lag_offsets.into_iter().map(offset).collect(), + }, + Op::FutureLabels { + alias, + target_op, + future_offsets, + } => Op::FutureLabels { + alias, + target_op: Box::new(target_op.map_clock(offset)), + future_offsets: future_offsets.into_iter().map(offset).collect(), + }, + Op::HasEdges { + alias, + path, + time_col, + window, + window_unit, + row_filter, + } => Op::HasEdges { + alias, + path, + time_col, + window, + window_unit, + row_filter, + }, + } + } +} + +/// The clock a serialized plan is written in. +pub fn plan_timestamp_unit(plan_json: &str) -> serde_json::Result { + #[derive(Deserialize)] + struct ClockTag { + #[serde(default)] + timestamp_unit: TimeUnit, + } + Ok(serde_json::from_str::(plan_json)?.timestamp_unit) +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index fa4deee73..d7326819b 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -10,8 +10,8 @@ use crate::state::{ }; use diskgraph_core::sampling::*; use diskgraph_core::sortkey::SortKeyDecoder; +use diskgraph_core::time::{convert_trunc, TimeUnit}; use diskgraph_core::tt_builder::{Value, WEIGHT_COL}; -use diskgraph_core::time::{abs_from_unit, Seconds, TimeUnit}; use diskgraph_index::layout::{EdgeListView, RelationSegment}; use diskgraph_index::reader::ShardReader; use gxhash::HashMap as GxHashMap; @@ -685,9 +685,9 @@ impl DiskGraphEngine { diskgraph_rowpack::RowPackContext::from_bytes(&ctx_bytes).map_err(|e| { EngineError::Internal(format!("failed to parse rowpack ctx: {}", e)) })?; - let decoder = RowPackDecoder::new(ctx).map_err(|e| { - EngineError::Internal(format!("failed to create decoder: {}", e)) - })?; + let decoder = RowPackDecoder::new(ctx) + .map_err(|e| EngineError::Internal(format!("failed to create decoder: {}", e)))? + .with_clock(metadata.timestamp_unit); node_decoders_map.insert(node_type.name.clone(), decoder); } } @@ -1498,11 +1498,7 @@ impl DiskGraphEngine { col.as_any() .downcast_ref::() .map(|arr| (0..arr.len()).map(|i| arr.value(i)).collect()) - .or_else(|| { - col.as_any() - .downcast_ref::() - .map(|arr| (0..arr.len()).map(|i| arr.value(i)).collect()) - }) + .or_else(|| seed_timestamps_in_clock(col, self.timestamp_unit)) }) } else { None @@ -2109,26 +2105,18 @@ impl DiskGraphEngine { ) -> Result<()> { let sample_start = Instant::now(); - // Convert seed timestamps from the caller's unit to the index unit (Seconds). - // The index always stores timestamps in seconds today; assert to catch if - // this assumption changes (at which point the engine should be made generic - // over UnitSystem). - assert_eq!( - self.timestamp_unit, - TimeUnit::Seconds, - "engine assumes seconds index; make engine generic over UnitSystem if this changes" - ); + // Convert seed timestamps from the caller's unit to the index clock. let mut request_cow = std::borrow::Cow::Borrowed(request); - if request.seed_timestamp_unit != TimeUnit::Seconds { + if request.seed_timestamp_unit != self.timestamp_unit { if let Some(mut ts_map) = request.seed_timestamps.clone() { for ts_vec in ts_map.values_mut() { for ts in ts_vec.iter_mut() { - *ts = abs_from_unit::(*ts, request.seed_timestamp_unit).raw(); + *ts = convert_trunc(*ts, request.seed_timestamp_unit, self.timestamp_unit); } } let mut owned = request.clone(); owned.seed_timestamps = Some(ts_map); - owned.seed_timestamp_unit = TimeUnit::Seconds; + owned.seed_timestamp_unit = self.timestamp_unit; request_cow = std::borrow::Cow::Owned(owned); } } @@ -2341,9 +2329,13 @@ impl DiskGraphEngine { if is_reserved_feature_column(&col.name) { continue; } - let (_, builder) = - diskgraph_rowpack::make_arrow_builder(&col.name, col.logical_type, 1024) - .map_err(|e| EngineError::Internal(e.to_string()))?; + let (_, builder) = diskgraph_rowpack::make_arrow_builder( + &col.name, + col.logical_type, + 1024, + decoder.clock(), + ) + .map_err(|e| EngineError::Internal(e.to_string()))?; builders.push(builder); } type_state.builders = Some(builders); @@ -3384,6 +3376,29 @@ impl DiskGraphEngine { } } +/// Read a decoded `TIMESTAMP` column as ticks of the index clock. +fn seed_timestamps_in_clock(col: &arrow_array::ArrayRef, clock: TimeUnit) -> Option> { + use arrow_array::cast::AsArray; + use arrow_array::types::*; + use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit}; + let DataType::Timestamp(unit, _) = col.data_type() else { + return None; + }; + let values = match unit { + ArrowTimeUnit::Second => col.as_primitive::().values(), + ArrowTimeUnit::Millisecond => col.as_primitive::().values(), + ArrowTimeUnit::Microsecond => col.as_primitive::().values(), + ArrowTimeUnit::Nanosecond => col.as_primitive::().values(), + }; + let unit = TimeUnit::from(*unit); + Some( + values + .iter() + .map(|&value| convert_trunc(value, unit, clock)) + .collect(), + ) +} + /// Project a seed-rowpack `RecordBatch` into the two consumer buckets /// `sample_runtime_batch` needs: /// diff --git a/packages/diskgraph/diskgraph-index/src/bin/explore.rs b/packages/diskgraph/diskgraph-index/src/bin/explore.rs index 815f986c8..6fed6edeb 100644 --- a/packages/diskgraph/diskgraph-index/src/bin/explore.rs +++ b/packages/diskgraph/diskgraph-index/src/bin/explore.rs @@ -1788,6 +1788,7 @@ fn format_sort_key_summary(rel: &EdgeRelationMetadata) -> String { EdgeSortKeyTimestampKind::Int8 => "int8", EdgeSortKeyTimestampKind::Int16 => "int16", EdgeSortKeyTimestampKind::Int32 => "int32", + EdgeSortKeyTimestampKind::Int64 => "int64", }; parts.push(format!("ts({kind})")); } @@ -1869,6 +1870,7 @@ fn timestamp_bits_for_meta(meta: &diskgraph_core::metadata::EdgeSortKeyTimestamp EdgeSortKeyTimestampKind::Int8 => 8, EdgeSortKeyTimestampKind::Int16 => 16, EdgeSortKeyTimestampKind::Int32 => 32, + EdgeSortKeyTimestampKind::Int64 => 64, } } diff --git a/packages/diskgraph/diskgraph-index/src/graph.rs b/packages/diskgraph/diskgraph-index/src/graph.rs index 0fcb6c2cc..c2063c0cb 100644 --- a/packages/diskgraph/diskgraph-index/src/graph.rs +++ b/packages/diskgraph/diskgraph-index/src/graph.rs @@ -17,6 +17,7 @@ use diskgraph_core::metadata::{ DatasetMetadata, EdgeRelationMetadata, NodeDataCompression, NodeShardMetadata, NodeTypeMetadata, PrimaryKeyMap, }; +use diskgraph_core::time::{convert_trunc, TimeUnit}; use diskgraph_core::tt_builder::Value; use serde::{Deserialize, Serialize}; use serde_json; @@ -1055,7 +1056,8 @@ impl GraphIndexReader { let bytes = fs::read(&ctx_path) .with_context(|| format!("read node data ctx {}", ctx_path.display()))?; let ctx = diskgraph_rowpack::RowPackContext::from_bytes(&bytes)?; - let decoder = diskgraph_rowpack::RowPackDecoder::new(ctx)?; + let decoder = diskgraph_rowpack::RowPackDecoder::new(ctx)? + .with_clock(metadata.timestamp_unit); (Some(cfg.clone()), Some(decoder)) } None => (None, None), @@ -1746,7 +1748,7 @@ impl TableReader { } /// Check whether the entity identified by `node_id` is temporally valid at - /// `observation_ts` (epoch seconds). Returns `true` if the entity's + /// `observation_ts` (index-clock ticks). Returns `true` if the entity's /// `create_time_col <= observation_ts` and `(end_time_col IS NULL OR /// end_time_col > observation_ts)`. /// @@ -1793,12 +1795,13 @@ impl TableReader { None => return Ok(true), }; + let clock = decoder.clock(); // Check create_time_col <= create_ref_ts if let Some(col_idx) = create_col_idx { let val = decoder.decode_column_value(raw, col_idx, scratch)?; - let ts_secs = temporal_value_to_epoch_secs(&val); - if let Some(secs) = ts_secs { - if secs > create_ref_ts { + let ticks = temporal_value_in_clock(&val, clock); + if let Some(ticks) = ticks { + if ticks > create_ref_ts { return Ok(false); } } @@ -1808,9 +1811,9 @@ impl TableReader { // Check end_time_col IS NULL OR end_time_col > end_ref_ts if let Some(col_idx) = end_col_idx { let val = decoder.decode_column_value(raw, col_idx, scratch)?; - let ts_secs = temporal_value_to_epoch_secs(&val); - if let Some(secs) = ts_secs { - if secs <= end_ref_ts { + let ticks = temporal_value_in_clock(&val, clock); + if let Some(ticks) = ticks { + if ticks <= end_ref_ts { return Ok(false); } } @@ -1878,12 +1881,16 @@ impl TableReader { } } -/// Convert a rowpack `ValueRef` representing a temporal column to epoch seconds. -fn temporal_value_to_epoch_secs(val: &diskgraph_rowpack::ValueRef<'_>) -> Option { +/// Convert a rowpack `ValueRef` representing a temporal column to the index clock. +fn temporal_value_in_clock(val: &diskgraph_rowpack::ValueRef<'_>, clock: TimeUnit) -> Option { match val { - diskgraph_rowpack::ValueRef::Timestamp(secs) => Some(*secs), - diskgraph_rowpack::ValueRef::Int64(ns) => Some(*ns / 1_000_000_000), - diskgraph_rowpack::ValueRef::Date(days) => Some(*days as i64 * 86400), + diskgraph_rowpack::ValueRef::Timestamp(ticks) => Some(*ticks), + diskgraph_rowpack::ValueRef::Int64(ns) => { + Some(convert_trunc(*ns, TimeUnit::Nanoseconds, clock)) + } + diskgraph_rowpack::ValueRef::Date(days) => { + Some(convert_trunc(i64::from(*days), TimeUnit::Days, clock)) + } diskgraph_rowpack::ValueRef::Null => None, _ => None, } diff --git a/packages/diskgraph/diskgraph-infer/src/lib.rs b/packages/diskgraph/diskgraph-infer/src/lib.rs index d6970ea61..f390fefc8 100644 --- a/packages/diskgraph/diskgraph-infer/src/lib.rs +++ b/packages/diskgraph/diskgraph-infer/src/lib.rs @@ -394,6 +394,7 @@ pub fn infer_source_metadata(root: &RootSource, options: InferOptions) -> Result let num_shards = diskgraph_ingest::config_converter::inference::estimate_num_shards(total_bytes); let metadata = GraphConfig { + timestamp_unit: Default::default(), dataset_name: root_dataset_name(root), tables: table_specs, rowpack: None, diff --git a/packages/diskgraph/diskgraph-ingest/README.md b/packages/diskgraph/diskgraph-ingest/README.md index 1689f19ad..59dd44cab 100644 --- a/packages/diskgraph/diskgraph-ingest/README.md +++ b/packages/diskgraph/diskgraph-ingest/README.md @@ -179,6 +179,18 @@ Key matching behavior: - parse failures or missing targets skip the edge and are counted. - duplicate primary keys are skipped and counted. +### Index clock + +The top-level `"timestamp_unit"` sets the clock timestamps are stored in: +`nanoseconds`, `microseconds`, `milliseconds`, `seconds` (the default), `hours` +or `days`. +- A value the clock cannot represent is stored as null; a null `create_time` + drops the row. +- `seconds`, `hours` and `days` keep the 8/16/32-bit edge time key; sub-second + clocks use the full 64-bit tick, so a `time` + `weight` sort key is a config + error on them. +- A training-table plan must be written in the index's clock. + ## Supported Parquet types Ingestion expects schemas to match `diskgraph_config.json` and supports: diff --git a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs index b8bbd94b9..9212c54fc 100644 --- a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs +++ b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs @@ -15,6 +15,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use crate::graph_config::{ColumnSpec, IngestedType, SemanticType}; +use diskgraph_core::time::{split_parts, TimeUnit}; use diskgraph_rowpack::{ColumnValueObserver, FloatListView, IntListView, StringListView}; use serde_json::{json, Value as JsonValue}; @@ -562,6 +563,8 @@ pub struct FileColStats { rng: u64, /// Per-column cached graduated set snapshots (Graduation strategy only). grad_snapshots: Vec>>>, + /// The index clock that timestamp statistics are in. + clock: TimeUnit, } /// Per-column stats collector, chosen by SemanticType. @@ -596,8 +599,8 @@ pub enum ColumnStatsCollector { count: u64, null_count: u64, parse_fail_count: u64, - min_secs: i64, - max_secs: i64, + min: i64, + max: i64, hll: HyperLogLog, }, Sequence { @@ -677,6 +680,7 @@ impl FileColStats { row_subsample: false, rng: 1, grad_snapshots: (0..n).map(|_| None).collect(), + clock: TimeUnit::Seconds, } } @@ -730,9 +734,15 @@ impl FileColStats { row_subsample, rng, grad_snapshots, + clock: TimeUnit::Seconds, } } + pub fn with_clock(mut self, clock: TimeUnit) -> Self { + self.clock = clock; + self + } + /// Merge another `FileColStats` into this one (per-column dispatch). /// Used on the channel receive side to accumulate per-table stats. pub fn merge(&mut self, other: &mut FileColStats) { @@ -847,7 +857,7 @@ impl FileColStats { ) -> serde_json::Map { let mut map = serde_json::Map::new(); for (col_idx, (col, name)) in self.cols.iter_mut().zip(column_names.iter()).enumerate() { - if let Some(mut val) = col.to_json() { + if let Some(mut val) = col.to_json(self.clock) { if has_counts[col_idx] { if let JsonValue::Object(ref mut obj) = val { obj.insert("category_counts".to_string(), json!([])); @@ -1102,8 +1112,8 @@ impl ColumnStatsCollector { count: 0, null_count: 0, parse_fail_count: 0, - min_secs: i64::MAX, - max_secs: i64::MIN, + min: i64::MAX, + max: i64::MIN, hll: HyperLogLog::new(), }, Some(SemanticType::Sequence) => ColumnStatsCollector::Sequence { @@ -1259,16 +1269,16 @@ impl ColumnStatsCollector { count, null_count, parse_fail_count, - min_secs, - max_secs, + min, + max, hll, }, ColumnStatsCollector::Timestamp { count: c2, null_count: nc2, parse_fail_count: pf2, - min_secs: ms2, - max_secs: xs2, + min: ms2, + max: xs2, hll: hll2, }, ) => { @@ -1276,11 +1286,11 @@ impl ColumnStatsCollector { *null_count += *nc2; *parse_fail_count += *pf2; hll.merge(hll2); - if *ms2 < *min_secs { - *min_secs = *ms2; + if *ms2 < *min { + *min = *ms2; } - if *xs2 > *max_secs { - *max_secs = *xs2; + if *xs2 > *max { + *max = *xs2; } } ( @@ -1405,7 +1415,7 @@ impl ColumnStatsCollector { /// Serialize this collector's stats to JSON. Returns None for Skip. #[allow(clippy::wrong_self_convention)] // needs &mut self for tdigest flush - fn to_json(&mut self) -> Option { + fn to_json(&mut self, clock: TimeUnit) -> Option { match self { ColumnStatsCollector::Numerical { null_count, @@ -1475,21 +1485,26 @@ impl ColumnStatsCollector { count, null_count, parse_fail_count, - min_secs, - max_secs, + min, + max, hll, .. } => { let total = (*count).saturating_add(*null_count); - Some(json!({ + let mut stats = json!({ "semantic_type": "timestamp", "count": total, "null_count": *null_count, "parse_fail_count": *parse_fail_count, "num_unique": hll.estimate() as u64, - "min_secs": *min_secs, - "max_secs": *max_secs, - })) + }); + let suffix = timestamp_stat_suffix(clock); + stats[format!("min_{suffix}")] = json!(*min); + stats[format!("max_{suffix}")] = json!(*max); + if clock != TimeUnit::Seconds { + stats["timestamp_unit"] = json!(clock); + } + Some(stats) } ColumnStatsCollector::Sequence { count, @@ -1568,14 +1583,28 @@ impl ColumnStatsCollector { // Kumo-format stats output (replaces Python bridge) // --------------------------------------------------------------------------- -/// Format an epoch-second timestamp to kumo timestamp string. +/// The `col_stats.json` key suffix of timestamp bounds (`min_secs` for seconds). +fn timestamp_stat_suffix(clock: TimeUnit) -> &'static str { + match clock { + TimeUnit::Nanoseconds => "ns", + TimeUnit::Microseconds => "us", + TimeUnit::Milliseconds => "ms", + TimeUnit::Hours => "hours", + TimeUnit::Days => "days", + TimeUnit::Seconds => "secs", + other => unreachable!("{other:?} is not an index clock"), + } +} + +/// Format a timestamp in `clock` ticks to kumo timestamp string. /// Sentinels (i64::MAX/MIN) and unparseable values → "kumo_timestamp_NaT". -fn format_kumo_timestamp(epoch_secs: i64) -> String { - if epoch_secs == i64::MAX || epoch_secs == i64::MIN { +fn format_kumo_timestamp(ticks: i64, clock: TimeUnit) -> String { + if ticks == i64::MAX || ticks == i64::MIN { return "kumo_timestamp_NaT".to_string(); } - match chrono::DateTime::from_timestamp(epoch_secs, 0) { - Some(dt) => format!("kumo_timestamp_{}", dt.format("%Y-%m-%d %H:%M:%S")), + let (epoch_secs, nanos) = split_parts(ticks, clock); + match chrono::DateTime::from_timestamp(epoch_secs, nanos) { + Some(dt) => format!("kumo_timestamp_{}", dt.format("%Y-%m-%d %H:%M:%S%.f")), None => "kumo_timestamp_NaT".to_string(), } } @@ -2050,8 +2079,8 @@ impl FileColStats { count: c, null_count, parse_fail_count, - min_secs, - max_secs, + min, + max, hll, .. } => { @@ -2082,11 +2111,21 @@ impl FileColStats { ); map.insert( kumo_key(name, sk, "MIN", false), - kumo_entry(json!(format_kumo_timestamp(*min_secs)), false, dv, total), + kumo_entry( + json!(format_kumo_timestamp(*min, self.clock)), + false, + dv, + total, + ), ); map.insert( kumo_key(name, sk, "MAX", false), - kumo_entry(json!(format_kumo_timestamp(*max_secs)), false, dv, total), + kumo_entry( + json!(format_kumo_timestamp(*max, self.clock)), + false, + dv, + total, + ), ); let num_unique = (hll.estimate() as u64).min(total); map.insert( @@ -2516,17 +2555,17 @@ impl ColumnValueObserver for FileColStats { } ColumnStatsCollector::Timestamp { count, - min_secs, - max_secs, + min, + max, hll, .. } => { *count += 1; - if value < *min_secs { - *min_secs = value; + if value < *min { + *min = value; } - if value > *max_secs { - *max_secs = value; + if value > *max { + *max = value; } hll.observe_i64(value); } @@ -2643,18 +2682,18 @@ impl ColumnValueObserver for FileColStats { fn observe_timestamp(&mut self, col_idx: usize, secs: i64) { if let ColumnStatsCollector::Timestamp { count, - min_secs, - max_secs, + min, + max, hll, .. } = &mut self.cols[col_idx] { *count += 1; - if secs < *min_secs { - *min_secs = secs; + if secs < *min { + *min = secs; } - if secs > *max_secs { - *max_secs = secs; + if secs > *max { + *max = secs; } hll.observe_i64(secs); } @@ -3857,15 +3896,21 @@ mod tests { #[test] fn test_format_kumo_timestamp_normal() { assert_eq!( - format_kumo_timestamp(1537401600), + format_kumo_timestamp(1537401600, TimeUnit::Seconds), "kumo_timestamp_2018-09-20 00:00:00" ); } #[test] fn test_format_kumo_timestamp_sentinels() { - assert_eq!(format_kumo_timestamp(i64::MAX), "kumo_timestamp_NaT"); - assert_eq!(format_kumo_timestamp(i64::MIN), "kumo_timestamp_NaT"); + assert_eq!( + format_kumo_timestamp(i64::MAX, TimeUnit::Seconds), + "kumo_timestamp_NaT" + ); + assert_eq!( + format_kumo_timestamp(i64::MIN, TimeUnit::Seconds), + "kumo_timestamp_NaT" + ); } #[test] @@ -3925,6 +3970,27 @@ mod tests { assert!(map.contains_key("ts__kumo.timestamp__NUM_UNIQUE__False")); } + #[test] + fn kumo_timestamps_carry_a_fraction_only_when_one_is_stored() { + let at = |ticks, clock| format_kumo_timestamp(ticks, clock); + assert_eq!( + at(1_700_000_000_123_456_789, TimeUnit::Nanoseconds), + "kumo_timestamp_2023-11-14 22:13:20.123456789" + ); + assert_eq!( + at(1_700_000_000_250, TimeUnit::Milliseconds), + "kumo_timestamp_2023-11-14 22:13:20.250" + ); + assert_eq!( + at(1_700_000_000_000, TimeUnit::Milliseconds), + "kumo_timestamp_2023-11-14 22:13:20" + ); + assert_eq!( + at(1_700_000_000, TimeUnit::Seconds), + "kumo_timestamp_2023-11-14 22:13:20" + ); + } + #[test] fn test_to_kumo_json_timestamp_sentinel() { let stypes = vec![Some(SemanticType::Timestamp)]; diff --git a/packages/diskgraph/diskgraph-ingest/src/config_converter/convert.rs b/packages/diskgraph/diskgraph-ingest/src/config_converter/convert.rs index 2971f0f02..e2ea961ed 100644 --- a/packages/diskgraph/diskgraph-ingest/src/config_converter/convert.rs +++ b/packages/diskgraph/diskgraph-ingest/src/config_converter/convert.rs @@ -89,6 +89,7 @@ pub fn convert_kumo_to_diskgraph( } let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: options.dataset_name.clone(), tables, rowpack: None, diff --git a/packages/diskgraph/diskgraph-ingest/src/config_converter/inference.rs b/packages/diskgraph/diskgraph-ingest/src/config_converter/inference.rs index d81977ddb..b49f1c2fe 100644 --- a/packages/diskgraph/diskgraph-ingest/src/config_converter/inference.rs +++ b/packages/diskgraph/diskgraph-ingest/src/config_converter/inference.rs @@ -533,13 +533,14 @@ pub struct ResolvedTemporal { } /// Detect (unless declared) and apply the temporal hint for one column, -/// yielding epoch seconds, or epoch days when `as_days`. Best-effort: a null or -/// unparseable value yields `None`, never an epoch. +/// yielding epochs in `clock` ticks, or epoch days when `as_days`. Best-effort: +/// a null or unparseable value yields `None`, never an epoch. pub fn resolve_temporal( array: &arrow_array::ArrayRef, declared_format: Option<&str>, declared_unit: Option, as_days: bool, + clock: diskgraph_core::time::TimeUnit, ) -> ResolvedTemporal { let normalized = normalize_epoch_column(array.as_ref()); let array = normalized.as_ref().unwrap_or(array); @@ -559,7 +560,8 @@ pub fn resolve_temporal( let mut access = diskgraph_rowpack::RowAccess::new( vec![hints], diskgraph_rowpack::RowAccessStats::default(), - ); + ) + .with_clock(clock); access.set_cols(std::slice::from_ref(array)); let epochs = (0..array.len()) .map(|row| { @@ -1690,6 +1692,7 @@ mod tests { mod tests_resolve { use super::*; use arrow_array::ArrayRef; + use diskgraph_core::time::TimeUnit; use std::sync::Arc; const JUL_4_2019: i64 = 1_562_198_400; @@ -1744,7 +1747,7 @@ mod tests_resolve { fmt: Option<&str>, unit: Option, ) -> ResolvedTemporal { - resolve_temporal(array, fmt, unit, false) + resolve_temporal(array, fmt, unit, false, TimeUnit::Seconds) } #[test] @@ -1913,7 +1916,7 @@ mod tests_resolve { #[test] fn resolve_string_epoch_column_is_a_unit_not_a_date_format() { let array = strings(&[Some("1562198400"), Some("1562284800")]); - let r = resolve_temporal(&array, None, None, false); + let r = resolve_temporal(&array, None, None, false, TimeUnit::Seconds); assert_eq!(r.format, None); assert_eq!(r.unit, Some(TimestampUnit::Seconds)); assert_eq!(r.epochs, vec![Some(JUL_4_2019), Some(JUL_4_2019 + 86_400)]); @@ -1927,7 +1930,7 @@ mod tests_resolve { Some("2019-07-04 12:00:00"), None, ]); - let r = resolve_temporal(&array, None, None, true); + let r = resolve_temporal(&array, None, None, true, TimeUnit::Seconds); let day = JUL_4_2019 / 86_400; assert_eq!(r.epochs, vec![Some(day), Some(day + 1), None]); assert_eq!(r.format.as_deref(), Some("%Y-%m-%d %H:%M:%S")); diff --git a/packages/diskgraph/diskgraph-ingest/src/config_converter/tests.rs b/packages/diskgraph/diskgraph-ingest/src/config_converter/tests.rs index 63cec5480..1afc3ed64 100644 --- a/packages/diskgraph/diskgraph-ingest/src/config_converter/tests.rs +++ b/packages/diskgraph/diskgraph-ingest/src/config_converter/tests.rs @@ -1427,6 +1427,7 @@ fn test_detect_timestamp_formats_from_parquet() { // Build a GraphConfig with a Timestamp column pointing to the parquet file. let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![TableSpec { name: "events".to_string(), @@ -1558,6 +1559,7 @@ fn test_detect_timestamp_formats_allows_unrecognized_create_time_strings() { writer.close().unwrap(); let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![TableSpec { name: "accounts".to_string(), @@ -1647,6 +1649,7 @@ fn test_detect_timestamp_formats_skips_native_timestamp_columns() { writer.close().unwrap(); let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![TableSpec { name: "events".to_string(), @@ -1718,6 +1721,7 @@ fn test_detect_timestamp_formats_integer_yyyymmdd() { writer.close().unwrap(); let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![TableSpec { name: "events".to_string(), @@ -1803,6 +1807,7 @@ fn test_detect_timestamp_formats_multi_value_sampling() { writer.close().unwrap(); let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![TableSpec { name: "events".to_string(), @@ -1896,6 +1901,7 @@ fn test_detect_timestamp_formats_skips_empty_parquet_files() { // Point the config at the directory (not individual files). let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![TableSpec { name: "transactions".to_string(), @@ -2182,6 +2188,7 @@ fn test_inference_skips_columns_with_explicit_timestamp_format() { writer.close().unwrap(); let mut config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "ds".to_string(), tables: vec![TableSpec { name: "events".to_string(), diff --git a/packages/diskgraph/diskgraph-ingest/src/graph_config.rs b/packages/diskgraph/diskgraph-ingest/src/graph_config.rs index f484431cd..837b9d5fd 100644 --- a/packages/diskgraph/diskgraph-ingest/src/graph_config.rs +++ b/packages/diskgraph/diskgraph-ingest/src/graph_config.rs @@ -4,6 +4,7 @@ use anyhow::{anyhow, ensure, Context, Result}; use diskgraph_core::metadata::EdgeStorageFormat; +use diskgraph_core::time::TimeUnit; use diskgraph_rowpack::RowPackConfig; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; @@ -53,6 +54,8 @@ fn source_dtype_is_decimal(dtype: Option<&str>) -> bool { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GraphConfig { + #[serde(default)] + pub timestamp_unit: TimeUnit, pub dataset_name: String, pub tables: Vec, #[serde(default)] @@ -67,6 +70,11 @@ pub struct GraphConfig { pub coercion_mode: CoercionMode, } +/// Sub-second ticks overflow a 32-bit key within years, so they key on the full value. +pub(crate) fn keys_full_timestamp(clock: TimeUnit) -> bool { + clock.tick_nanos() < TimeUnit::Seconds.tick_nanos() +} + fn is_default_coercion_mode(mode: &CoercionMode) -> bool { *mode == CoercionMode::default() } @@ -783,6 +791,28 @@ impl GraphConfig { if self.tables.is_empty() { errors.push("at least one table is required".to_string()); } + if !self.timestamp_unit.is_index_clock() { + errors.push(format!( + "unsupported timestamp_unit {:?}", + self.timestamp_unit + )); + } + let mut directions = self.tables.iter().flat_map(|table| { + let fkeys = table.fkeys.iter().flat_map(|fk| [&fk.forward, &fk.reverse]); + let edges = table + .edge_table + .iter() + .flat_map(|e| [&e.forward, &e.reverse]); + fkeys.chain(edges).flatten() + }); + if keys_full_timestamp(self.timestamp_unit) + && directions.any(|d| d.time.is_some() && d.weight.is_some()) + { + errors.push(format!( + "a {:?} time key leaves no room for a sort key weight", + self.timestamp_unit + )); + } if let Some(shards) = self.num_shards { if shards == 0 { errors.push("num_shards must be >= 1".to_string()); @@ -1763,6 +1793,7 @@ mod tests { }); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![customers, orders], rowpack: None, @@ -1819,6 +1850,7 @@ mod tests { }); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![users, products, reviews], rowpack: None, @@ -1877,6 +1909,7 @@ mod tests { }); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![customers, orders], rowpack: None, @@ -1923,6 +1956,7 @@ mod tests { }); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".to_string(), tables: vec![customers, orders], rowpack: None, @@ -1993,6 +2027,7 @@ mod tests { fn normalize_flips_string_multicat_with_delim_to_stringarray() { let col = make_multicat_column("tags", IngestedType::String, Some(",")); let mut cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2019,6 +2054,7 @@ mod tests { fn normalize_leaves_string_multicat_without_delim_alone() { let col = make_multicat_column("tags", IngestedType::String, None); let mut cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2038,6 +2074,7 @@ mod tests { fn normalize_leaves_native_stringarray_alone() { let col = make_multicat_column("tags", IngestedType::StringArray, None); let mut cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2059,6 +2096,7 @@ mod tests { col.semantic_type = Some(SemanticType::Text); col.multicategorical_delimiter = Some(",".into()); let mut cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2077,6 +2115,7 @@ mod tests { fn normalize_is_idempotent() { let col = make_multicat_column("tags", IngestedType::String, Some(",")); let mut cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2107,6 +2146,7 @@ mod tests { fn normalize_latlong_forces_fixed_len_two() { let col = make_latlong_column("location", IngestedType::Float32Array, None); let mut cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2125,6 +2165,7 @@ mod tests { fn validate_latlong_requires_numeric_array() { let col = make_latlong_column("location", IngestedType::String, Some(2)); let cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, @@ -2145,6 +2186,7 @@ mod tests { fn validate_latlong_requires_fixed_len_two() { let col = make_latlong_column("location", IngestedType::Float32Array, Some(3)); let cfg = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "t".into(), tables: vec![make_table("items", vec![col], None)], rowpack: None, diff --git a/packages/diskgraph/diskgraph-ingest/src/lib.rs b/packages/diskgraph/diskgraph-ingest/src/lib.rs index a79192a08..bbc6e467d 100644 --- a/packages/diskgraph/diskgraph-ingest/src/lib.rs +++ b/packages/diskgraph/diskgraph-ingest/src/lib.rs @@ -83,6 +83,7 @@ use diskgraph_core::metadata::{ EdgeSortKeyWeightEncoding, EdgeSortKeyWeightKind, EdgeSortKeyWeightMetadata, EdgeStorageFormat, NodeDataCompression, NodeShardMetadata, NodeTypeMetadata, PrimaryKeyMap, }; +use diskgraph_core::time::{convert_floor, convert_trunc, Seconds, TimeUnit}; use diskgraph_fs::{CloudObjectCondition, CloudTransfer, PooledBytes}; use diskgraph_index::builder::{ EdgeBufferPool as IndexEdgeBufferPool, MemoryBudget, MemoryBudgetGuard, @@ -658,6 +659,15 @@ fn encoded_num_shards(root: &Path) -> Result> { .map(|n| n as usize)) } +/// An Encode marker written before the clock was always recorded omits it, which +/// reads as seconds. +fn encode_marker_matches(marker: &[u8], payload: &[u8]) -> bool { + marker == payload + || serde_json::from_slice::(marker) + .and_then(|config| serde_json::to_vec(&config)) + .is_ok_and(|normalized| normalized == payload) +} + fn table_sized_bytes(declared: Option, compute_local_bytes: impl FnOnce() -> u64) -> u64 { match declared { Some(bytes) if bytes > 0 => bytes, @@ -890,7 +900,7 @@ pub(crate) fn ingest_with_stats_from_source_and_sidecars( } if pred == IngestPhase::Encode { let marker = fs::read(phase_marker_path(root, pred.name()))?; - if marker != encode_marker_payload { + if !encode_marker_matches(&marker, &encode_marker_payload) { return Err(anyhow!( "resume source identity does not match the completed Encode marker" )); @@ -1386,10 +1396,25 @@ mod tests_resume_phase { use super::*; use std::collections::HashSet; + #[test] + fn a_config_without_a_clock_is_seconds_and_resumes_its_encode_marker() { + let source: GraphConfig = + serde_json::from_str(r#"{"dataset_name": "old", "tables": []}"#).unwrap(); + assert_eq!(source.timestamp_unit, TimeUnit::Seconds); + let payload = serde_json::to_vec(&source).unwrap(); + let mut marker = serde_json::to_value(&source).unwrap(); + assert_eq!(marker["timestamp_unit"], "seconds"); + marker.as_object_mut().unwrap().remove("timestamp_unit"); + let marker = serde_json::to_vec(&marker).unwrap(); + assert!(encode_marker_matches(&marker, &payload)); + assert!(!encode_marker_matches(br#"{"dataset_name": "other", "tables": []}"#, &payload)); + } + #[test] fn public_ingest_rejects_graduation_for_exact_graph_feature_stats() { let output = tempfile::tempdir().unwrap(); let source = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "graduation_rejected".to_owned(), tables: Vec::new(), rowpack: None, @@ -1830,6 +1855,7 @@ mod tests_repopulate_helpers { fn graph_config(tables: Vec) -> GraphConfig { GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_dataset".to_string(), tables, rowpack: None, @@ -6391,6 +6417,7 @@ pub(crate) enum TimestampEncodingKind { Int8, Int16, Int32, + Int64, } #[derive(Clone, Debug)] @@ -6415,7 +6442,14 @@ struct WeightEncodingSpec { } /// Convert TimeSpec from config to internal encoding spec. -fn time_spec_to_encoding(spec: &TimeSpec) -> TimestampEncodingSpec { +fn time_spec_to_encoding(spec: &TimeSpec, clock: TimeUnit) -> TimestampEncodingSpec { + if crate::graph_config::keys_full_timestamp(clock) { + return TimestampEncodingSpec { + kind: TimestampEncodingKind::Int64, + anchor: 0, + step: 1, + }; + } let kind = match spec.bitwidth { Some(8) => TimestampEncodingKind::Int8, Some(16) => TimestampEncodingKind::Int16, @@ -6441,8 +6475,8 @@ fn time_spec_to_encoding(spec: &TimeSpec) -> TimestampEncodingSpec { }; TimestampEncodingSpec { kind, - anchor: spec.anchor_time.unwrap_or(0), - step, + anchor: convert_floor(spec.anchor_time.unwrap_or(0), TimeUnit::Seconds, clock), + step: convert_trunc(step, TimeUnit::Seconds, clock).max(1), } } @@ -6541,6 +6575,7 @@ fn timestamp_width_bits(kind: &TimestampEncodingKind) -> u8 { TimestampEncodingKind::Int8 => 8, TimestampEncodingKind::Int16 => 16, TimestampEncodingKind::Int32 => 32, + TimestampEncodingKind::Int64 => 64, } } @@ -6554,10 +6589,13 @@ fn weight_width_bits(kind: &WeightEncodingKind) -> u8 { } } -fn build_relation_sort_key_config(spec: &DirectionSpec) -> Option { +fn build_relation_sort_key_config( + spec: &DirectionSpec, + clock: TimeUnit, +) -> Option { let mut width_bits: u16 = 0; let timestamp = spec.time.as_ref().map(|ts| { - let encoding = time_spec_to_encoding(ts); + let encoding = time_spec_to_encoding(ts, clock); width_bits += timestamp_width_bits(&encoding.kind) as u16; SortKeyTimestampConfig { table: ts.table.clone(), @@ -6599,6 +6637,7 @@ fn build_sort_key_metadata(cfg: &RelationSortKeyConfig) -> EdgeSortKeyMetadata { TimestampEncodingKind::Int8 => EdgeSortKeyTimestampKind::Int8, TimestampEncodingKind::Int16 => EdgeSortKeyTimestampKind::Int16, TimestampEncodingKind::Int32 => EdgeSortKeyTimestampKind::Int32, + TimestampEncodingKind::Int64 => EdgeSortKeyTimestampKind::Int64, }, anchor: ts.encoding.anchor, step: ts.encoding.step, @@ -6771,6 +6810,7 @@ pub(crate) fn encode_timestamp_value(value: Option, encoding: &TimestampEnc let clamped = delta.clamp(i32::MIN as i64, i32::MAX as i64) as i32; ordered_signed_to_u64(clamped as i64, 32) } + TimestampEncodingKind::Int64 => ordered_signed_to_u64(delta, 64), } } @@ -7170,11 +7210,11 @@ impl IngestContext { let dst_sort = fkey .forward .as_ref() - .and_then(build_relation_sort_key_config); + .and_then(|spec| build_relation_sort_key_config(spec, source.timestamp_unit)); let src_sort = fkey .reverse .as_ref() - .and_then(build_relation_sort_key_config); + .and_then(|spec| build_relation_sort_key_config(spec, source.timestamp_unit)); let rev_name = format!("{base}_reverse"); // Forward relation: paired with reverse only when bidirectional. @@ -7219,20 +7259,18 @@ impl IngestContext { name: fwd_name.clone(), src_table: edge.src_fk.ref_table.clone(), dst_table: edge.dst_fk.ref_table.clone(), - sort_key: edge - .forward - .as_ref() - .and_then(build_relation_sort_key_config), + sort_key: edge.forward.as_ref().and_then(|spec| { + build_relation_sort_key_config(spec, source.timestamp_unit) + }), reverse_of: Some(rev_name.clone()), }; let dst_rel = RelationDef { name: rev_name, src_table: edge.dst_fk.ref_table.clone(), dst_table: edge.src_fk.ref_table.clone(), - sort_key: edge - .reverse - .as_ref() - .and_then(build_relation_sort_key_config), + sort_key: edge.reverse.as_ref().and_then(|spec| { + build_relation_sort_key_config(spec, source.timestamp_unit) + }), reverse_of: Some(fwd_name), }; relation_defs.push(src_rel); @@ -7372,6 +7410,7 @@ impl IngestContext { let col_stats_enabled = self.col_stats; let cat_count_strategy = self.cat_count_strategy; let row_subsample = self.row_subsample; + let clock = self.source.timestamp_unit; let mut table_cat_counts: HashMap< String, @@ -7494,6 +7533,7 @@ impl IngestContext { &dest_key_errors, &relation_defs, &fk_edge_targets, + clock, )?; Ok((dropped, key_failures)) }, @@ -7523,7 +7563,7 @@ impl IngestContext { let mut pk_max_key_len = 0usize; let row_access_hints = build_row_access_hints(table); let stats = RowAccessStats::new(Some(trimmed_strings.as_ref())); - let mut access = RowAccess::new(row_access_hints, stats); + let mut access = RowAccess::new(row_access_hints, stats).with_clock(clock); let mut prefix = Vec::new(); let mut prefix_rows = 0usize; let mut prefix_bytes = 0usize; @@ -7738,6 +7778,7 @@ impl IngestContext { ref_resolver.as_deref(), &sidecars, &mut observer, + clock, )?; observer.flush_categories(); Ok((node, key_failures, Some(observer))) @@ -7760,6 +7801,7 @@ impl IngestContext { ref_resolver.as_deref(), &sidecars, &mut diskgraph_rowpack::NoOpObserver, + clock, )?; Ok((node, key_failures, None)) } @@ -7857,7 +7899,7 @@ impl IngestContext { if let Some(stats) = table_col_stats.remove(&table.name) { outputs.push(crate::col_stats::ColStatsTable { name: table.name.clone(), - stats, + stats: stats.with_clock(clock), column_names: table.columns.iter().map(|c| c.name.clone()).collect(), ingested_types: table .columns @@ -8383,6 +8425,7 @@ impl IngestContext { ref_resolver: Option<&crate::mapping::RefResolver>, sidecars: &StageSidecarColumns, observer: &mut O, + clock: TimeUnit, ) -> Result where O: diskgraph_rowpack::ColumnValueObserver, @@ -8436,7 +8479,7 @@ impl IngestContext { let mut total_rows = 0u64; let row_access_hints = build_row_access_hints(table); let stats = RowAccessStats::new(Some(trimmed_strings.as_ref())); - let mut access = RowAccess::new(row_access_hints, stats); + let mut access = RowAccess::new(row_access_hints, stats).with_clock(clock); let mut key_bufs: Vec> = vec![Vec::new(); table.columns.len()]; let mut f16_buf: Vec = Vec::new(); let mut node_tags = Vec::new(); @@ -8852,8 +8895,9 @@ impl IngestContext { // (`group_ts_idx.is_some()`), so the non-grouping path is unchanged. let group_ts = match group_ts_idx { Some(idx) => Some( + // Buckets stay one second wide whatever the index clock. access - .get_timestamp(idx) + .timestamp_in::(idx) .with_context(|| { format!( "TT-NG ts-grouping: failed to decode TIMESTAMP for table `{}` row {}", @@ -8973,6 +9017,7 @@ impl IngestContext { dest_key_errors: &std::sync::Arc>>, relation_defs: &[RelationDef], fk_edge_targets: &HashMap, + clock: TimeUnit, ) -> Result where I: IntoIterator>, @@ -9030,7 +9075,7 @@ impl IngestContext { let row_access_hints = build_row_access_hints(table); let stats = RowAccessStats::new(Some(trimmed_strings.as_ref())); - let mut access = RowAccess::new(row_access_hints, stats); + let mut access = RowAccess::new(row_access_hints, stats).with_clock(clock); macro_rules! process_edge_batch { ($batch:expr, $src_key_buf:expr, $dst_key_buf:expr, $forward_row:expr) => { let timestamp_idx = create_time_info.as_ref().map(|(idx, _)| *idx); @@ -11969,7 +12014,7 @@ impl IngestContext { } let index_metadata = DatasetMetadata { - timestamp_unit: diskgraph_core::time::TimeUnit::Seconds, + timestamp_unit: self.source.timestamp_unit, dataset_name: self.source.dataset_name.clone(), node_types, edge_relations, @@ -14856,6 +14901,7 @@ impl SortKeyComponentBatcher { column_idx, encoding, } => { + let clock = access.clock(); let ts_value = if let Some(secs) = access.get_timestamp(*column_idx).ok().flatten() { Some(secs) @@ -14864,7 +14910,7 @@ impl SortKeyComponentBatcher { .get_date(*column_idx) .ok() .flatten() - .map(|days| i64::from(days) * 86_400) + .map(|days| convert_trunc(i64::from(days), TimeUnit::Days, clock)) }; encode_timestamp_value(ts_value, encoding) } diff --git a/packages/diskgraph/diskgraph-ingest/src/stage_extension.rs b/packages/diskgraph/diskgraph-ingest/src/stage_extension.rs index f68a8eb9e..cc7d367a4 100644 --- a/packages/diskgraph/diskgraph-ingest/src/stage_extension.rs +++ b/packages/diskgraph/diskgraph-ingest/src/stage_extension.rs @@ -24,9 +24,11 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; +use diskgraph_core::metadata::DatasetMetadata; use crate::config::{ResolveColumnSpec, ResolveKind}; use crate::graph_config::{ColumnSpec, DataSource, FKeySpec, GraphConfig, IngestedType, TableSpec}; +use crate::index_metadata::read_index_metadata; use crate::{ingest_with_stats_from_source_and_sidecars, Dest, IngestOptions}; /// One stage table to ingest into a stage extension. @@ -228,7 +230,10 @@ pub fn ingest_stage_extension_with_resolution_and_sidecars( } }; + let parent_metadata: DatasetMetadata = + serde_json::from_str(&read_index_metadata(parent_index_path)?)?; let graph_config = GraphConfig { + timestamp_unit: parent_metadata.timestamp_unit, dataset_name: dataset_name.to_string(), tables, rowpack: None, diff --git a/packages/diskgraph/diskgraph-ingest/src/tests_unit.rs b/packages/diskgraph/diskgraph-ingest/src/tests_unit.rs index 0201a4aa4..9c35d71b5 100644 --- a/packages/diskgraph/diskgraph-ingest/src/tests_unit.rs +++ b/packages/diskgraph/diskgraph-ingest/src/tests_unit.rs @@ -466,6 +466,46 @@ fn encode_timestamp_int8() { assert_eq!(encoded_clamp, ordered_signed_to_u64(i8::MAX as i64, 8)); } +#[test] +fn far_future_date32_edge_keys_sort_after_2037_in_every_clock() { + use crate::{ + encode_timestamp_value, time_spec_to_encoding, timestamp_width_bits, + SortKeyComponentBatcher, SortKeyComponentEncoder, SortKeyComponentKind, + }; + use diskgraph_core::time::{convert_trunc, TimeUnit}; + let year_9999: ArrayRef = Arc::new(arrow_array::Date32Array::from(vec![2_932_896])); + let spec = crate::graph_config::TimeSpec { + table: "events".into(), + column: "at".into(), + bitwidth: None, + anchor_time: None, + unit: "seconds".into(), + }; + for clock in TimeUnit::INDEX_CLOCKS { + let mut access = RowAccess::new(vec![RowAccessHints::default()], RowAccessStats::new(None)) + .with_clock(clock); + access.set_cols(std::slice::from_ref(&year_9999)); + let encoding = time_spec_to_encoding(&spec, clock); + let width = timestamp_width_bits(&encoding.kind) / 8; + let year_2037 = convert_trunc(24_472, TimeUnit::Days, clock); + let key_2037 = encode_timestamp_value(Some(year_2037), &encoding); + let mut batcher = SortKeyComponentBatcher::new( + Arc::from("events"), + SortKeyComponentKind::Timestamp, + width, + SortKeyComponentEncoder::Timestamp { + column_idx: 0, + encoding, + }, + 1, + ); + batcher.push_row(0, &mut access).unwrap(); + let mut key = [0u8; 8]; + key[..usize::from(width)].copy_from_slice(&batcher.buffers[0]); + assert!(u64::from_le_bytes(key) > key_2037, "{clock:?}"); + } +} + #[test] fn convert_string_date_to_rowpack_date() { let table = TableSpec { diff --git a/packages/diskgraph/diskgraph-ingest/tests/graph_config.rs b/packages/diskgraph/diskgraph-ingest/tests/graph_config.rs index ee50bd286..e346bc880 100644 --- a/packages/diskgraph/diskgraph-ingest/tests/graph_config.rs +++ b/packages/diskgraph/diskgraph-ingest/tests/graph_config.rs @@ -57,6 +57,7 @@ fn base_table(name: &str) -> TableSpec { fn base_metadata() -> GraphConfig { GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".into(), tables: vec![base_table("User"), base_table("Device")], rowpack: None, @@ -80,6 +81,7 @@ fn assert_err_contains(meta: GraphConfig, needle: &str) { #[test] fn validate_minimal_metadata() { let meta = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".into(), tables: vec![base_table("User")], rowpack: None, @@ -106,6 +108,7 @@ fn validate_rejects_unsupported_direct_multicategorical_dtype() { multicategorical_delimiter: None, }); let meta = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".into(), tables: vec![table], rowpack: None, @@ -122,6 +125,7 @@ fn validate_rejects_unsupported_direct_multicategorical_dtype() { #[test] fn validate_missing_pk_column() { let meta = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test".into(), tables: vec![TableSpec { name: "User".into(), diff --git a/packages/diskgraph/diskgraph-ingest/tests/resolution_ingest.rs b/packages/diskgraph/diskgraph-ingest/tests/resolution_ingest.rs index 297fde0df..8b5e5d7da 100644 --- a/packages/diskgraph/diskgraph-ingest/tests/resolution_ingest.rs +++ b/packages/diskgraph/diskgraph-ingest/tests/resolution_ingest.rs @@ -98,6 +98,7 @@ fn ingest_tiny_parent_index(parent_root: &Path, num_rows: i64) -> Result<()> { group_by_timestamp: false, }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "tiny_parent".to_string(), tables: vec![table], rowpack: None, @@ -276,6 +277,7 @@ fn ingest_resolution_with_stats( let mut table = stage_table_spec(stage_parquet.to_str().unwrap()); table.fkeys = vec![]; // explicit: no FK edges. let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "resolve_run".to_string(), tables: vec![table], rowpack: None, @@ -677,6 +679,7 @@ fn ingest_exclude_resolution( }); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "resolve_exclude_run".to_string(), tables: vec![table], rowpack: None, diff --git a/packages/diskgraph/diskgraph-ingest/tests/resolution_runtime_batch.rs b/packages/diskgraph/diskgraph-ingest/tests/resolution_runtime_batch.rs index c3836bb2b..5e6d0ed4f 100644 --- a/packages/diskgraph/diskgraph-ingest/tests/resolution_runtime_batch.rs +++ b/packages/diskgraph/diskgraph-ingest/tests/resolution_runtime_batch.rs @@ -179,6 +179,7 @@ fn ingest_parent_users_orders(parent_root: &Path, n_users: i64) -> Result<()> { }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "users_orders_parent".to_string(), tables: vec![users, orders], rowpack: None, @@ -318,6 +319,7 @@ fn build_fixture_to( let table = stage_table_spec(stage_parquet.to_str().unwrap()); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "resolve_run".to_string(), tables: vec![table], rowpack: None, @@ -794,6 +796,7 @@ fn build_fixture_with_exclude( table.columns.push(typed_col("exclude_target_keys", IngestedType::Int64Array)); table.columns.push(typed_col("exclude_node_ids", IngestedType::Int64Array)); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "resolve_exclude_run".to_string(), tables: vec![table], rowpack: None, diff --git a/packages/diskgraph/diskgraph-ingest/tests/seed_table_ingest.rs b/packages/diskgraph/diskgraph-ingest/tests/seed_table_ingest.rs index 4c3030499..fbf9100db 100644 --- a/packages/diskgraph/diskgraph-ingest/tests/seed_table_ingest.rs +++ b/packages/diskgraph/diskgraph-ingest/tests/seed_table_ingest.rs @@ -293,6 +293,7 @@ fn test_ingest_skips_reverse_when_bidirectional_false_in_config_dst() { bidirectional: false, }]; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "fwd_only".to_string(), tables: vec![users, orders], rowpack: None, @@ -459,6 +460,7 @@ fn ingest_tiny_parent_index_impl( }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "tiny_parent".to_string(), tables: vec![table], rowpack: None, @@ -826,6 +828,7 @@ fn ingest_users_orders_parent_index(parent_root: &Path, n_users: i64) -> Result< }]; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "users_orders_parent".to_string(), tables: vec![users, orders], rowpack: None, @@ -877,6 +880,7 @@ fn ingest_one_extension( let seed_table = seed_table_spec(seed_parquet.to_str().unwrap()); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: format!("seed_{tt_id}"), tables: vec![seed_table], rowpack: None, @@ -1637,6 +1641,7 @@ fn test_sample_runtime_batch_seed_topology_excluded_features_preserved() { }, ); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: format!("seed_{tt_id}"), tables: vec![seed_table], rowpack: None, @@ -2647,6 +2652,7 @@ fn ingest_one_extension_with_weight( }, ); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: format!("seed_{tt_id}"), tables: vec![seed_table], rowpack: None, diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/compiler.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/compiler.py index 020a35b55..15e5c75c8 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/compiler.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/compiler.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json from typing import Any, Dict, List, Optional, Tuple from kumoapi.pquery.AST import (Aggregation, ASTNode, Column, Condition, @@ -11,7 +12,7 @@ from ._c import (DiskgraphConfigError, DiskgraphDataError, DiskgraphSystemError, DiskGraphTTBuilder, EdgeRef, RowFilter, - StageFilter, TimeUnit) + StageFilter, TimeUnit, read_index_metadata) from .split_contract import SplitSpec VALID_TASK_TYPES = { @@ -74,14 +75,22 @@ def _validate_lp_layout(name: str, value: Optional[str], allowed: set[str]) -> N def _requires_non_empty_target_list(task_type: str) -> bool: return task_type in TASK_TYPES_REQUIRING_NON_EMPTY_TARGET_LIST +def index_timestamp_unit(index_path: Optional[str]) -> str: + """The clock the index at ``index_path`` stores timestamps in (seconds when omitted).""" + if index_path is None: + return "seconds" + return json.loads(read_index_metadata(index_path)).get("timestamp_unit", "seconds") + class DiskGraphPlanCompiler: """ Translates a ValidatedPredictiveQuery (AST) into a DiskGraphTTPlan (JSON). + Plans are written in the clock of the index at ``index_path`` (seconds when omitted). """ - def __init__(self, builder: Optional[DiskGraphTTBuilder] = None): - self.builder = builder or DiskGraphTTBuilder() + def __init__(self, index_path: Optional[str] = None): + self._timestamp_unit = index_timestamp_unit(index_path) + self.builder = DiskGraphTTBuilder(self._timestamp_unit) self._alias_counter = 0 self._is_link_pred: bool = False self._lp_rhs_entity_table: Optional[str] = None @@ -401,7 +410,7 @@ def _compile_row_filter(self, node: ASTNode, query: ValidatedPredictiveQuery, co return RowFilter.compare_to_entity(col_name, op, ent_col) else: val = node.get_typed_value() - return RowFilter.compare(col_name, op, val) + return RowFilter.compare(col_name, op, val, self._timestamp_unit) path, leaf = self._extract_path(node.target, query, root_table=root_table) col_name = leaf.fqn.split('.')[1] @@ -420,9 +429,10 @@ def _compile_row_filter(self, node: ASTNode, query: ValidatedPredictiveQuery, co op = "is_null" if op in ("==", "eq") else "is_not_null" val = 0 if not path: - return RowFilter.compare(col_name, op, val) + return RowFilter.compare(col_name, op, val, self._timestamp_unit) else: - return RowFilter.compare_remote_lhs(path, col_name, op, val) + return RowFilter.compare_remote_lhs(path, col_name, op, val, + self._timestamp_unit) elif isinstance(node, LogicalOperation): left = self._compile_row_filter(node.left, query, context_table=context_table) @@ -453,7 +463,7 @@ def _compile_stage_filter(self, node: ASTNode, query: ValidatedPredictiveQuery) if val is None: op = "is_null" if op in ("==", "eq") else "is_not_null" val = 0 - return StageFilter.compare(alias, op, val) + return StageFilter.compare(alias, op, val, self._timestamp_unit) # LHS is a Join wrapping an aggregation (edge traversal in stage filter) elif isinstance(node.target, Join): @@ -468,7 +478,7 @@ def _compile_stage_filter(self, node: ASTNode, query: ValidatedPredictiveQuery) if val is None: op = "is_null" if op in ("==", "eq") else "is_not_null" val = 0 - return StageFilter.compare(alias, op, val) + return StageFilter.compare(alias, op, val, self._timestamp_unit) # Simple Column LHS in StageFilter — must emit a column-fetch # op so the Rust engine can resolve the alias at execution time. @@ -481,7 +491,7 @@ def _compile_stage_filter(self, node: ASTNode, query: ValidatedPredictiveQuery) if val is None: op = "is_null" if op in ("==", "eq") else "is_not_null" val = 0 - return StageFilter.compare(alias, op, val) + return StageFilter.compare(alias, op, val, self._timestamp_unit) elif isinstance(node, LogicalOperation): left = self._compile_stage_filter(node.left, query) diff --git a/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py b/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py index 9fc993ee8..0fe10551c 100644 --- a/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py +++ b/packages/diskgraph/diskgraph-pybind/python/diskgraph/tt_artifacts.py @@ -25,6 +25,7 @@ build_split_spec_json as _native_build_split_spec_json, build_train_table_artifacts_for_existing_table as _native_build, ) +from .compiler import index_timestamp_unit from .split_contract import SplitSpec @@ -45,9 +46,11 @@ def _split_spec_to_plan_json( split_spec: Optional[SplitSpec], *, split_family_hint: Optional[str], + index_path: Optional[str] = None, ) -> Optional[str]: """Serialize a Python `SplitSpec` to the JSON the Rust executor consumes - (a `SplitSpec` from `diskgraph-core`), for + (a `SplitSpec` from `diskgraph-core` in the clock of the index at + `index_path`, seconds when omitted), for `build_train_table_artifacts_for_existing_table`'s `split_spec_json`. Returns None when split_spec is None (no split path). @@ -80,7 +83,8 @@ def _split_spec_to_plan_json( resolved_family = split_spec.kind or split_family_hint if resolved_family != split_spec.kind: split_spec = dataclasses.replace(split_spec, kind=resolved_family) - return _native_build_split_spec_json(split_spec) + return _native_build_split_spec_json(split_spec, + index_timestamp_unit(index_path)) def _resolve_common_config( @@ -242,7 +246,7 @@ def build_train_table_artifacts_for_existing_table( ) plan_split_json = _split_spec_to_plan_json( - split_config, split_family_hint=split_family_hint) + split_config, split_family_hint=split_family_hint, index_path=index_path) # The producer_summary_json + train_table_edges_config kwargs are # currently consumed elsewhere (kumo's separate edge-bundle pass). diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py new file mode 100644 index 000000000..3123fe0ab --- /dev/null +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""An index clock changes only the stored tick: every clock gives main's results.""" + +import datetime +import decimal +import json +import pathlib + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +import diskgraph +from diskgraph import (DiskGraphTTBuilder, EdgeRef, RowFilter, SamplerEngine, + SplitSpec, TimeUnit, execute_tt_gen) + +CLOCKS = ("nanoseconds", "microseconds", "milliseconds", "seconds", "hours", "days") +TICK_NS = { + "nanoseconds": 1, + "microseconds": 1_000, + "milliseconds": 1_000_000, + "seconds": 1_000_000_000, + "hours": 3_600_000_000_000, + "days": 86_400_000_000_000, +} +DAY = 86_400 +EPOCH_DAY = 19_700 # 2023-12-09 +USERS = 6 + + +def _ingest(index_dir: pathlib.Path, clock: str, tables: list) -> str: + diskgraph.ingest( + config_json=json.dumps({ + "dataset_name": "clocks", + "num_shards": 1, + "timestamp_unit": clock, + "tables": tables, + }), + output_dir=str(index_dir), + ) + return str(index_dir) + + +def _as_nanos(array: pa.Array) -> list: + unit = array.type.unit + factor = {"s": 10**9, "ms": 10**6, "us": 10**3, "ns": 1}[unit] + return [None if v is None else v * factor for v in array.cast(pa.int64()).to_pylist()] + + +def _features(index: str, table: str, column: str, num_nodes: int) -> list: + engine = SamplerEngine(index) + events = engine.fetch_features([ + {"node_type": table, "node_ids": list(range(num_nodes)), "feature_names": [column]}, + ]) + batches = [e["batch"] for e in events if e["type"] == "feature"] + return _as_nanos(pa.concat_arrays([b.column(column) for b in batches])) + + +@pytest.fixture(scope="module") +def whole_day_data(tmp_path_factory) -> pathlib.Path: + data_dir = tmp_path_factory.mktemp("whole_days") + pq.write_table( + pa.table({ + "user_id": pa.array(range(USERS), pa.int64()), + "signup": pa.array([EPOCH_DAY + u for u in range(USERS)], pa.date32()), + }), + data_dir / "users.parquet", + ) + orders = [(u * 10 + k, u, (EPOCH_DAY + u + 3 * k) * DAY) for u in range(USERS) + for k in range(5)] + pq.write_table( + pa.table({ + "order_id": pa.array([o[0] for o in orders], pa.int64()), + "user_id": pa.array([o[1] for o in orders], pa.int64()), + "ordered_at": pa.array([o[2] for o in orders], pa.timestamp("s")), + }), + data_dir / "orders.parquet", + ) + return data_dir + + +def _whole_day_tables(data_dir: pathlib.Path) -> list: + return [ + { + "name": "users", + "columns": [ + {"name": "user_id", "ingested_dtype": "int64"}, + {"name": "signup", "ingested_dtype": "timestamp"}, + ], + "primary_key": "user_id", + "create_time": "signup", + "source": {"type": "parquet", "paths": [str(data_dir / "users.parquet")]}, + }, + { + "name": "orders", + "columns": [ + {"name": "order_id", "ingested_dtype": "int64"}, + {"name": "user_id", "ingested_dtype": "int64"}, + {"name": "ordered_at", "ingested_dtype": "timestamp", + "semantic_type": "timestamp"}, + ], + "primary_key": "order_id", + "create_time": "ordered_at", + "fkeys": [{ + "src_column": "user_id", + "dst_table": "users", + "reverse": {"time": {"table": "orders", "column": "ordered_at", + "unit": "s"}}, + }], + "source": {"type": "parquet", "paths": [str(data_dir / "orders.parquet")]}, + }, + ] + + +def _sampled_edges(index: str) -> dict: + plan = {"layers": [{"fanouts": {}, "default_fanout": { + "fanout": 100, "mode": "temporal_most_recent", "absent_key_policy": "reject_plan"}}]} + anchors = [(EPOCH_DAY + 7 + u) * DAY * 10**9 for u in range(USERS)] + events = SamplerEngine(index, random_seed=0).sample( + {"users": list(range(USERS))}, json.dumps(plan), fused=True, + seed_timestamps={"users": anchors}, seed_timestamp_unit=TimeUnit.Nanoseconds) + edges = [] + for event in events: + if event["type"] == "edge": + batch = event["batch"] + edges += zip(batch.column("id_a").to_pylist(), batch.column("id_b").to_pylist()) + return sorted(edges) + + +def _training_table(index: str, clock: str, window_days: int = 7) -> list: + builder = DiskGraphTTBuilder(clock) + builder.temporal([(EPOCH_DAY + d) * DAY for d in (4, 8, 12, 16)], forecast_length=1, + unit="seconds") + builder.node_prediction() + builder.for_each("users", "user_id") + builder.entity_time_col("signup") + builder.set_target_stype("numerical") + builder.aggregate( + "orders_next_week", EdgeRef("users", "user_id", "orders", "user_id"), "count", + time_col="ordered_at", window=(0, window_days), window_unit="days", + row_filter=RowFilter.compare( + "ordered_at", ">", datetime.datetime(2023, 12, 10, tzinfo=datetime.timezone.utc), + clock), + ) + builder.split_ratio(SplitSpec( + kind="TimeRangeSplit", + ranges=[((EPOCH_DAY + 4) * DAY, (EPOCH_DAY + 9) * DAY), + ((EPOCH_DAY + 9) * DAY, (EPOCH_DAY + 13) * DAY), + ((EPOCH_DAY + 13) * DAY, (EPOCH_DAY + 20) * DAY)], + is_temporal=True, + timestamp_unit="seconds", + offset_unit="seconds", + unit="seconds", + )) + generated = execute_tt_gen(builder.build(), index, "clock_tt") + staging = pathlib.Path(generated["staging_dir"]) + rows = [] + for part in sorted(staging.rglob("*.parquet")): + table = pq.read_table(part) + rows += [(part.relative_to(staging).as_posix(), row) for row in table.to_pylist()] + return sorted(rows, key=repr) + + +def _timestamp_stats(index: str, clock: str) -> dict: + stats = json.loads((pathlib.Path(index) / "col_stats.json").read_text())["orders"] + stats = stats["ordered_at"] + suffix = {"nanoseconds": "ns", "microseconds": "us", "milliseconds": "ms", + "seconds": "secs", "hours": "hours", "days": "days"}[clock] + return { + "count": stats["count"], + "null_count": stats["null_count"], + "min_ns": stats[f"min_{suffix}"] * TICK_NS[clock], + "max_ns": stats[f"max_{suffix}"] * TICK_NS[clock], + } + + +def _whole_day_results(data_dir: pathlib.Path, index_dir: pathlib.Path, clock: str) -> dict: + index = _ingest(index_dir, clock, _whole_day_tables(data_dir)) + return { + "orders": _features(index, "orders", "ordered_at", USERS * 5), + "signups": _features(index, "users", "signup", USERS), + "sampled": _sampled_edges(index), + "training_table": _training_table(index, clock), + "stats": _timestamp_stats(index, clock), + } + + +@pytest.mark.parametrize("clock", [c for c in CLOCKS if c != "seconds"]) +def test_every_clock_matches_the_seconds_index_on_whole_days(whole_day_data, tmp_path, + clock): + seconds = _whole_day_results(whole_day_data, tmp_path / "seconds", "seconds") + assert _whole_day_results(whole_day_data, tmp_path / clock, clock) == seconds + + +def test_index_metadata_and_config_record_every_clock(whole_day_data, tmp_path): + for clock in CLOCKS: + index = _ingest(tmp_path / clock, clock, _whole_day_tables(whole_day_data)) + metadata = json.loads(diskgraph.read_index_metadata(index)) + assert metadata["timestamp_unit"] == clock + config = json.loads((pathlib.Path(index) / "diskgraph_config.json").read_text()) + assert config["timestamp_unit"] == clock + + +# Instant 2024-01-15 10:00:00.750 UTC, written in each source representation. +INSTANT_NS = 1_705_312_800_750_000_000 + + +def _floor(ns: int, clock: str) -> int: + return ns // TICK_NS[clock] + + +def _ceil(ns: int, clock: str) -> int: + return -(-ns // TICK_NS[clock]) + + +def _decimal(ns: int, unit_ns: int, clock: str) -> int: + step = min(unit_ns, TICK_NS[clock]) + truncated = ns // step * step if ns >= 0 else -(-ns // step * step) + return _ceil(truncated, clock) + + +@pytest.fixture(scope="module") +def messy_data(tmp_path_factory) -> pathlib.Path: + data_dir = tmp_path_factory.mktemp("messy") + pq.write_table( + pa.table({ + "event_id": pa.array([0, 1, 2], pa.int64()), + "text_with_format": ["2024-01-15 10:00:00.750", None, "not a time"], + "text_without_format": ["2024-01-15T10:00:00.750Z", "2024-01-15T10:00:00", + "2024-01-15 10:00:00"], + "epoch_millis": pa.array([1_705_312_800_750, None, -1], pa.int64()), + "epoch_decimal": pa.array([decimal.Decimal("1705312800.75"), None, + decimal.Decimal("-0.75")], pa.decimal128(20, 2)), + "epoch_millis_decimal": pa.array([decimal.Decimal("1705312800750.5"), None, None], + pa.decimal128(20, 1)), + "typed_nanos": pa.array([INSTANT_NS, INSTANT_NS + 1, None], + pa.timestamp("ns")), + "date": pa.array([19_737, 29_999, 2_932_896], pa.date32()), + }), + data_dir / "events.parquet", + ) + return data_dir + + +def _messy_tables(data_dir: pathlib.Path) -> list: + timestamp = {"ingested_dtype": "timestamp"} + return [{ + "name": "events", + "columns": [ + {"name": "event_id", "ingested_dtype": "int64"}, + {"name": "text_with_format", **timestamp, + "timestamp_format": "%Y-%m-%d %H:%M:%S%.f"}, + {"name": "text_without_format", **timestamp}, + {"name": "epoch_millis", **timestamp, "timestamp_unit": "millis"}, + {"name": "epoch_decimal", **timestamp}, + {"name": "epoch_millis_decimal", **timestamp, "timestamp_unit": "millis"}, + {"name": "typed_nanos", **timestamp}, + {"name": "date", **timestamp}, + ], + "primary_key": "event_id", + "source": {"type": "parquet", "paths": [str(data_dir / "events.parquet")]}, + }] + + +@pytest.mark.parametrize("clock", CLOCKS) +def test_messy_values_follow_mains_rules_in_every_clock(messy_data, tmp_path, clock): + index = _ingest(tmp_path / clock, clock, _messy_tables(messy_data)) + jan_15 = 19_737 * DAY * 10**9 + expected = { + # Text truncates; epochs and typed values round up; decimals truncate to + # the finer of their unit and the clock, then round up. + "text_with_format": [_floor(INSTANT_NS, clock), None, None], + "text_without_format": [_floor(INSTANT_NS, clock), None, _floor(jan_15 + 36_000 * 10**9, clock)], + "epoch_millis": [_ceil(INSTANT_NS, clock), None, _ceil(-10**6, clock)], + "epoch_decimal": [_decimal(INSTANT_NS, 10**9, clock), None, + _decimal(-750_000_000, 10**9, clock)], + "epoch_millis_decimal": [_decimal(INSTANT_NS + 500_000, 10**6, clock), None, None], + "typed_nanos": [_ceil(INSTANT_NS, clock), _ceil(INSTANT_NS + 1, clock), None], + # Dates outside main's range (after 2038-01-19) are null. + "date": [_floor(jan_15, clock), None, None], + } + for column, ticks in expected.items(): + decoded = _features(index, "events", column, 3) + # Hours and days decode as Arrow seconds, which are exact multiples. + assert [None if v is None else _floor(v, clock) for v in decoded] == ticks, column + + +def test_values_beyond_the_clock_are_null_and_drop_required_rows(tmp_path): + year_9999 = 253_402_300_799 + pq.write_table( + pa.table({ + "event_id": pa.array([0, 1], pa.int64()), + "happened_at": pa.array([1_705_312_800, 1_705_312_800], pa.timestamp("s")), + "noted": ["9999-12-31 23:59:59", "2024-01-15 10:00:00"], + "expires_at": pa.array([year_9999, 1_705_312_800], pa.timestamp("s")), + }), + tmp_path / "events.parquet", + ) + timestamp = {"ingested_dtype": "timestamp"} + tables = [{ + "name": "events", + "columns": [ + {"name": "event_id", "ingested_dtype": "int64"}, + {"name": "happened_at", **timestamp}, + {"name": "noted", **timestamp}, + {"name": "expires_at", **timestamp}, + ], + "primary_key": "event_id", + "source": {"type": "parquet", "paths": [str(tmp_path / "events.parquet")]}, + }] + index = _ingest(tmp_path / "ns", "nanoseconds", tables) + assert _features(index, "events", "noted", 2) == [None, 1_705_312_800 * 10**9] + assert _features(index, "events", "expires_at", 2) == [None, 1_705_312_800 * 10**9] + + tables[0]["create_time"] = "expires_at" + index = _ingest(tmp_path / "ns_required", "nanoseconds", tables) + nodes = json.loads(diskgraph.read_index_metadata(index))["node_types"] + assert [node["max_id"] for node in nodes if node["name"] == "events"] == [0] + + +def test_rows_whose_window_leaves_the_clock_are_skipped(whole_day_data, tmp_path): + index = _ingest(tmp_path / "ns", "nanoseconds", _whole_day_tables(whole_day_data)) + last_day_in_clock = (2**63 - 1) // (DAY * 10**9) + window_days = last_day_in_clock - (EPOCH_DAY + 8) + rows = _training_table(index, "nanoseconds", window_days) + anchors = {row["TIMESTAMP"].date() for part, row in rows if part.startswith("tt.parquet")} + assert anchors == {datetime.date(2023, 12, 13), datetime.date(2023, 12, 17)} + + +@pytest.mark.parametrize("clock", ["seconds", "milliseconds"]) +def test_plans_keep_anchors_beyond_the_nanosecond_range(clock): + year_9999 = 253_402_300_799 + builder = DiskGraphTTBuilder(clock) + builder.temporal([year_9999], forecast_length=1) + builder.for_each("users", "user_id", default_time_value=year_9999) + plan = json.loads(builder.build()) + ticks = year_9999 * TICK_NS["seconds"] // TICK_NS[clock] + assert plan["temporal"]["timestamps"] == [ticks] + assert plan["entity"]["default_time_value"] == ticks + + +def test_an_unreadable_datetime_is_a_config_error(): + class NoTimestamp(datetime.datetime): + def timestamp(self): + raise OverflowError("out of range") + + with pytest.raises(diskgraph.DiskgraphConfigError): + RowFilter.compare("ordered_at", ">", NoTimestamp(2024, 1, 1), "milliseconds") diff --git a/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs b/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs index 923e1798c..8eb65f329 100644 --- a/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs +++ b/packages/diskgraph/diskgraph-pybind/src/bp_col_stats.rs @@ -23,6 +23,7 @@ use serde_json::{json, Map as JsonMap, Value as JsonValue}; use diskgraph_core::accumulators::{ kumo_entry, kumo_key, CategoryCountStrategy, WelfordAccumulator, MAX_CATEGORIES, }; +use diskgraph_core::time::{convert_ceil, TimeUnit}; use diskgraph_ingest::col_stats::{FileColStats, TableCategoryCounts}; use diskgraph_ingest::graph_config::{ColumnSpec, IngestedType, SemanticType, TimestampUnit}; use diskgraph_ingest::PublicIngestError; @@ -249,9 +250,13 @@ fn integer_timestamp( })) } -fn timestamp_seconds(access: &RowAccess<'_>, col: &ColumnSpec) -> Result> { +fn timestamp_ticks(access: &RowAccess<'_>, col: &ColumnSpec) -> Result> { if let Some((unit, raw)) = integer_timestamp(access, col)? { - return Ok(Some(unit.scale_to_secs(raw))); + return Ok(Some(convert_ceil( + raw, + unit.to_arrow().into(), + access.clock(), + ))); } access.get_timestamp(0) } @@ -261,7 +266,7 @@ fn histogram_timestamp_seconds( access: &RowAccess<'_>, raw_timestamp: Option<(TimestampUnit, i64)>, ) -> Result> { - let Some(seconds) = timestamp_seconds(access, col)? else { + let Some(seconds) = timestamp_ticks(access, col)? else { return Ok(None); }; let raw_unit = raw_timestamp.or(integer_timestamp(access, col)?); @@ -288,7 +293,8 @@ fn for_each_timestamp_second( _ => None, }; let arrays = [array.clone()]; - let mut access = RowAccess::new(vec![row_access_hints(col)], RowAccessStats::new(None)); + let mut access = RowAccess::new(vec![row_access_hints(col)], RowAccessStats::new(None)) + .with_clock(TimeUnit::Seconds); access.set_cols(&arrays); for row in 0..array.len() { access.set_row(row); @@ -636,14 +642,16 @@ fn drive_from_arrow_column( col: &ColumnSpec, observer: &mut O, mut integer_range: Option<&mut IntegerRange>, + clock: TimeUnit, ) -> Result<()> { let arrays = [array.clone()]; - let mut access = RowAccess::new(vec![row_access_hints(col)], RowAccessStats::new(None)); + let mut access = + RowAccess::new(vec![row_access_hints(col)], RowAccessStats::new(None)).with_clock(clock); access.set_cols(&arrays); for row in 0..array.len() { access.set_row(row); let present = if col.semantic_type == Some(SemanticType::Timestamp) { - access_or_parse_failure(timestamp_seconds(&access, col), col_idx, observer) + access_or_parse_failure(timestamp_ticks(&access, col), col_idx, observer) .map(|value| observer.observe_timestamp(col_idx, value)) .is_some() } else { @@ -694,7 +702,7 @@ fn drive_from_arrow_column( access_or_parse_failure(value, col_idx, observer).is_some() } IngestedType::Timestamp => { - access_or_parse_failure(timestamp_seconds(&access, col), col_idx, observer) + access_or_parse_failure(timestamp_ticks(&access, col), col_idx, observer) .map(|value| observer.observe_timestamp(col_idx, value)) .is_some() } @@ -795,6 +803,7 @@ fn stats_for_row_group( integer_spec_indices: &[usize], table_cat: &Arc, score_class: Option<(usize, usize)>, + clock: TimeUnit, ) -> Result<(FileColStats, ByClassStats, Vec)> { let mut fcs = FileColStats::new_with_categories(columns, Arc::clone(table_cat)); let mut by_class = ByClassStats::new(); @@ -820,6 +829,7 @@ fn stats_for_row_group( col, &mut fcs, range, + clock, ); if let Err(error) = result { if error.is::() { @@ -846,17 +856,20 @@ fn stats_for_row_group( Ok((fcs, by_class, integer_ranges)) } -/// Compute full and min Kumo-format statistics for prediction parquet files. +/// Compute full and min Kumo-format statistics for prediction parquet files, +/// with timestamp statistics in `clock` ticks. pub fn compute_parquet_col_stats( parquet_paths: &[&Path], columns: &[ColumnSpec], data_version: &str, + clock: TimeUnit, ) -> Result<(String, String)> { compute_parquet_col_stats_with_max_categories( parquet_paths, columns, data_version, MAX_CATEGORIES, + clock, ) } @@ -867,6 +880,7 @@ fn accumulate_parquet_col_stats<'a>( parquet_paths: &'a [&'a Path], columns: &'a [ColumnSpec], max_categories: u64, + clock: TimeUnit, ) -> Result<(FileColStats, Arc, Accumulated<'a>)> { if max_categories == 0 { bail!("category count limit must be greater than zero"); @@ -927,6 +941,7 @@ fn accumulate_parquet_col_stats<'a>( &integer_spec_indices, &table_cat, score_class, + clock, ) }) .try_reduce( @@ -951,7 +966,7 @@ fn accumulate_parquet_col_stats<'a>( )? }; Ok(( - acc, + acc.with_clock(clock), table_cat, Accumulated { parquet_paths, @@ -986,9 +1001,10 @@ pub(crate) fn compute_parquet_col_stats_with_max_categories( columns: &[ColumnSpec], data_version: &str, max_categories: u64, + clock: TimeUnit, ) -> Result<(String, String)> { let (acc, table_cat, run) = - accumulate_parquet_col_stats(parquet_paths, columns, max_categories)?; + accumulate_parquet_col_stats(parquet_paths, columns, max_categories, clock)?; let names: Vec = columns.iter().map(|c| c.name.clone()).collect(); let types: Vec = columns.iter().map(|c| c.ingested_dtype.clone()).collect(); let mut kumo = acc.into_kumo_stats( @@ -1337,9 +1353,14 @@ mod tests { SemanticType::Categorical, )]; - let (full, _) = - compute_parquet_col_stats_with_max_categories(&[path.as_path()], &columns, "v1", 2) - .unwrap(); + let (full, _) = compute_parquet_col_stats_with_max_categories( + &[path.as_path()], + &columns, + "v1", + 2, + TimeUnit::Seconds, + ) + .unwrap(); assert_eq!( val(&full, "status__kumo.categorical__CATEGORY_COUNTS__False"), json!([["red", 3], ["blue", 2]]) @@ -1368,7 +1389,8 @@ mod tests { ) -> (String, String) { use diskgraph_ingest::col_stats::minify_kumo_stats; let (mut acc, table_cat, run) = - accumulate_parquet_col_stats(paths, columns, max_categories).unwrap(); + accumulate_parquet_col_stats(paths, columns, max_categories, TimeUnit::Seconds) + .unwrap(); let names: Vec = columns.iter().map(|c| c.name.clone()).collect(); let types: Vec = columns.iter().map(|c| c.ingested_dtype.clone()).collect(); let mut full = acc.to_kumo_json(&names, &types, Some(&table_cat), None, "v1"); @@ -1440,8 +1462,14 @@ mod tests { col("SCORE", IngestedType::Float64, SemanticType::Numerical), col("CLASS", IngestedType::Int64, SemanticType::Categorical), ]; - let streamed = - compute_parquet_col_stats_with_max_categories(&paths, &columns, "v1", cap).unwrap(); + let streamed = compute_parquet_col_stats_with_max_categories( + &paths, + &columns, + "v1", + cap, + TimeUnit::Seconds, + ) + .unwrap(); assert_eq!( streamed, legacy_col_stats(&paths, &columns, cap), @@ -1453,7 +1481,9 @@ mod tests { #[test] fn configurable_category_cap_rejects_zero() { - let error = compute_parquet_col_stats_with_max_categories(&[], &[], "v1", 0).unwrap_err(); + let error = + compute_parquet_col_stats_with_max_categories(&[], &[], "v1", 0, TimeUnit::Seconds) + .unwrap_err(); assert!(error .to_string() .contains("category count limit must be greater than zero")); @@ -1490,7 +1520,9 @@ mod tests { col("SCORE", IngestedType::Float64, SemanticType::Numerical), col("CLASS", IngestedType::String, SemanticType::Categorical), ]; - let (full, min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); // COUNT == n across columns (merge across row groups correct). assert_eq!(val(&full, "entity__kumo.ID__COUNT__False"), n as u64); @@ -1601,7 +1633,9 @@ mod tests { ts_col("month", None), ts_col("day", None), ]; - let (full, _) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); for (name, stype) in [ ("rank", "kumo.numerical"), ("revenue", "kumo.numerical"), @@ -1649,7 +1683,9 @@ mod tests { IngestedType::Float32Array, SemanticType::Latlong, )]; - let (full, _) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!( val(&full, "location__kumo.latlong__SEQUENCE_MAX_LENGTH__False"), 2 @@ -1683,7 +1719,9 @@ mod tests { col("SCORE", IngestedType::Float32, SemanticType::Numerical), col("CLASS", IngestedType::Bool, SemanticType::Categorical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); // Numerical scalar stats populated (not empty-Welford defaults). let mx = val(&full, "SCORE__kumo.numerical__MAX__False") .as_f64() @@ -1732,7 +1770,9 @@ mod tests { col("SCORE", IngestedType::Float64, SemanticType::Numerical), col("CLASS", IngestedType::String, SemanticType::Categorical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); let bc = val( &full, "SCORE__kumo.numerical__SUMMARY_STATS_BY_CLASS__False", @@ -1782,8 +1822,13 @@ mod tests { col("entity", IngestedType::String, SemanticType::Id), col("score", IngestedType::Float64, SemanticType::Numerical), ]; - let (full, _) = - compute_parquet_col_stats(&[p1.as_path(), p2.as_path()], &columns, "v1").unwrap(); + let (full, _) = compute_parquet_col_stats( + &[p1.as_path(), p2.as_path()], + &columns, + "v1", + TimeUnit::Seconds, + ) + .unwrap(); assert_eq!(val(&full, "entity__kumo.ID__COUNT__False"), 4); assert_eq!(val(&full, "score__kumo.numerical__MIN__False"), 1.0); assert_eq!(val(&full, "score__kumo.numerical__MAX__False"), 4.0); @@ -1818,6 +1863,7 @@ mod tests { &[p1.as_path(), p2.as_path()], &[col("score", IngestedType::Float64, SemanticType::Numerical)], "v1", + TimeUnit::Seconds, ) .unwrap_err(); assert!(format!("{err:#}").contains("column 'score'")); @@ -1882,7 +1928,9 @@ mod tests { ); let columns = vec![ts_col("created_at", None), ts_col("updated_at", None)]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!( val( &full, @@ -1936,8 +1984,13 @@ mod tests { 1, ); - let (full, _min) = - compute_parquet_col_stats(&[path.as_path()], &[ts_col("ts", None)], "v1").unwrap(); + let (full, _min) = compute_parquet_col_stats( + &[path.as_path()], + &[ts_col("ts", None)], + "v1", + TimeUnit::Seconds, + ) + .unwrap(); let histogram = val(&full, "ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False"); let histogram = histogram.as_array().unwrap(); assert_eq!(histogram.len(), TIMESTAMP_HISTOGRAM_BINS); @@ -1980,8 +2033,13 @@ mod tests { 1, ); - let (full, _) = - compute_parquet_col_stats(&[path.as_path()], &[ts_col("ts", None)], "v1").unwrap(); + let (full, _) = compute_parquet_col_stats( + &[path.as_path()], + &[ts_col("ts", None)], + "v1", + TimeUnit::Seconds, + ) + .unwrap(); let histogram = val(&full, "ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False"); let histogram = histogram.as_array().unwrap(); assert_eq!(histogram.len(), TIMESTAMP_HISTOGRAM_BINS); @@ -2033,7 +2091,9 @@ mod tests { ts_col("null_ts", None), col("score", IngestedType::Float64, SemanticType::Numerical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); let stats: JsonMap = serde_json::from_str(&full).unwrap(); assert!(!stats.contains_key("ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False")); assert!(!stats.contains_key("null_ts__kumo.timestamp__HISTOGRAM_HALF_MONTH__False")); @@ -2077,7 +2137,9 @@ mod tests { let mut timestamp = ts_col("ts", Some(TimestampUnit::Millis)); timestamp.ingested_dtype = IngestedType::Int64; let columns = vec![timestamp]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!(val(&full, "ts__kumo.timestamp__MIN__False"), expected); assert_eq!(val(&full, "ts__kumo.timestamp__MAX__False"), expected); assert_eq!( @@ -2088,7 +2150,8 @@ mod tests { let mut timestamp = ts_col("ts", None); timestamp.ingested_dtype = IngestedType::Int64; let (full, _min) = - compute_parquet_col_stats(&[path.as_path()], &[timestamp], "v1").unwrap(); + compute_parquet_col_stats(&[path.as_path()], &[timestamp], "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!(val(&full, "ts__kumo.timestamp__MIN__False"), expected); } @@ -2106,7 +2169,9 @@ mod tests { 64, ); let columns = vec![ts_col("ts", Some(TimestampUnit::Seconds))]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!( val(&full, "ts__kumo.timestamp__MIN__False"), "kumo_timestamp_2022-07-31 22:00:01" @@ -2129,7 +2194,8 @@ mod tests { IngestedType::Float64, SemanticType::Numerical, )]; - let err = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap_err(); + let err = compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap_err(); let msg = format!("{err:#}"); assert!(msg.contains("NOT_THERE") && msg.contains("schema error")); @@ -2142,7 +2208,9 @@ mod tests { 1024, ); let columns = vec![col("SCORE", IngestedType::Float64, SemanticType::Numerical)]; - let (full, _) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!( val(&full, "SCORE__kumo.numerical__MIN__False"), JsonValue::Null @@ -2175,7 +2243,9 @@ mod tests { col("cnt", IngestedType::Int32, SemanticType::Numerical), col("b", IngestedType::UInt8, SemanticType::Numerical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!(val(&full, "cnt__kumo.numerical__COUNT__False"), 3u64); assert_eq!(val(&full, "b__kumo.numerical__COUNT__False"), 3u64); let mfull: JsonMap = serde_json::from_str(&full).unwrap(); @@ -2203,6 +2273,7 @@ mod tests { &[path.as_path()], &[col("v", IngestedType::Int64, SemanticType::Numerical)], "v1", + TimeUnit::Seconds, ) .unwrap(); assert_eq!(val(&full, "v__kumo.numerical__MIN__False"), i64::MAX - 2); @@ -2231,6 +2302,7 @@ mod tests { &[path.as_path()], &[col("v", IngestedType::Int64, SemanticType::Numerical)], "v1", + TimeUnit::Seconds, ) .unwrap(); let histogram = val(&full, "v__kumo.numerical__HISTOGRAM__False"); @@ -2267,6 +2339,7 @@ mod tests { SemanticType::Numerical, )], "v1", + TimeUnit::Seconds, ) .unwrap(); @@ -2307,6 +2380,7 @@ mod tests { SemanticType::Numerical, )], "v1", + TimeUnit::Seconds, ) .unwrap(); @@ -2377,6 +2451,7 @@ mod tests { SemanticType::Numerical, )], "v1", + TimeUnit::Seconds, ) .unwrap_err(); @@ -2418,6 +2493,7 @@ mod tests { col("float", IngestedType::Float64, SemanticType::Numerical), ], "v1", + TimeUnit::Seconds, ) .unwrap(); let stats: JsonMap = serde_json::from_str(&full).unwrap(); @@ -2452,7 +2528,9 @@ mod tests { col("subnormal", IngestedType::Float64, SemanticType::Numerical), ]; - let (full, _min) = compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + let (full, _min) = + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); let histogram = val(&full, "extreme__kumo.numerical__HISTOGRAM__False"); let histogram = histogram.as_array().unwrap(); assert_eq!(histogram.len(), NUMERICAL_HISTOGRAM_BINS); @@ -2509,7 +2587,8 @@ mod tests { ); let columns = vec![ts_col("ts", None)]; let (full, _min) = - compute_parquet_col_stats(&[path.as_path()], &columns, "v1").unwrap(); + compute_parquet_col_stats(&[path.as_path()], &columns, "v1", TimeUnit::Seconds) + .unwrap(); assert_eq!( val(&full, "ts__kumo.timestamp__MIN__False"), "kumo_timestamp_1969-12-31 23:59:59", diff --git a/packages/diskgraph/diskgraph-pybind/src/custom_tt_entry.rs b/packages/diskgraph/diskgraph-pybind/src/custom_tt_entry.rs index 423473e92..b79bfb6d1 100644 --- a/packages/diskgraph/diskgraph-pybind/src/custom_tt_entry.rs +++ b/packages/diskgraph/diskgraph-pybind/src/custom_tt_entry.rs @@ -43,11 +43,12 @@ use std::path::{Path, PathBuf}; -use diskgraph_core::time::{Seconds, TimeUnit}; +use diskgraph_core::time::{TimeUnit, UnitSystem}; use diskgraph_core::tt_builder::{ DiskGraphTTPlan, EntitySpec, FactPseudoTimeSpec, PredictionType, RefResolveKind, RefResolveSpec, SplitSpec as PlanSplitSpec, }; +use diskgraph_core::with_index_clock; use diskgraph_train_table::custom_tt_row_errors::{ IngestOnError, RowFailuresReport, TimestampReference, }; @@ -57,7 +58,7 @@ use diskgraph_train_table::ttgen_output::{PRIMARY_ENTITY_COL, PRIMARY_TARGET_COL use pyo3::prelude::*; use pyo3::types::PyDict; -use crate::tt_executor::{run_gen_bundle, validate_tt_id}; +use crate::tt_executor::{index_clock, run_gen_bundle, validate_tt_id}; fn custom_tt_artifact_base(parquet_path: &Path) -> PathBuf { // Multi-part contract (kumo Phase 2): kumo passes the train-table @@ -98,7 +99,7 @@ fn path_with_suffix(base: &Path, suffix: &str) -> PathBuf { /// /// Args mirror what kumo's /// `ensure_diskgraph_native_artifacts_for_train_table` calls into. The -/// `split_spec_json` is a JSON-serialized `SplitSpec` shape (the +/// `split_spec_json` is a JSON-serialized `SplitSpec` in the index clock (the /// pure-Rust plan SplitSpec, NOT the Python SplitSpec dataclass — the /// Python wrapper does the conversion). /// @@ -160,7 +161,7 @@ pub fn build_train_table_artifacts_for_existing_table( .map_err(|e| { crate::DiskgraphConfigError::new_err(format!("invalid fact_pseudo_time_json: {e}")) })?; - build_bundle_for_existing_table( + with_index_clock!(index_clock(index_path)?, U => build_bundle_for_existing_table::( py, output_path, index_path, @@ -180,7 +181,7 @@ pub fn build_train_table_artifacts_for_existing_table( allow_negative_weight, timestamp_reference_json, /* is_prediction */ false, - ) + )) } /// Build TT artifacts (manifest + sidecars) for an already-materialized @@ -226,7 +227,7 @@ pub fn build_prediction_table_artifacts_for_existing_table( task_type: Option<&str>, ingest_on_error: &str, ) -> PyResult> { - build_bundle_for_existing_table( + with_index_clock!(index_clock(index_path)?, U => build_bundle_for_existing_table::( py, output_path, index_path, @@ -251,7 +252,7 @@ pub fn build_prediction_table_artifacts_for_existing_table( /* allow_negative_weight */ false, /* timestamp_reference_json */ None, /* is_prediction */ true, - ) + )) } /// Build the plan's `fkey_mapping`: always the LHS `ENTITY1 → entity_table` @@ -375,7 +376,7 @@ fn row_failures_to_py<'py>( /// entries. `is_prediction` toggles the plan flag (predict skips the customer /// TARGET); the PREDICT wrapper additionally passes `split_spec_json=None` so /// the executor takes the `SplitMode::None` single-`predict`-stage path. -fn build_bundle_for_existing_table( +fn build_bundle_for_existing_table( py: Python<'_>, output_path: &str, index_path: &str, @@ -427,7 +428,7 @@ fn build_bundle_for_existing_table( // the executor takes the SplitMode::None path. A custom split (e.g. // `ManualSplit` on a customer-provided split column) flows through here as a // SplitSpec and routes via the column deciders — there is no split.bin. - let split: Option> = match split_spec_json { + let split: Option> = match split_spec_json { Some(s) if !s.is_empty() => Some(serde_json::from_str(s).map_err(|e| { crate::DiskgraphConfigError::new_err(format!("invalid split_spec_json: {e}")) })?), @@ -445,7 +446,7 @@ fn build_bundle_for_existing_table( // (ENTITYk, TIMESTAMP, TARGETk) is read directly from the customer // parquet by the chunk processor. LP CTT uses LinkArray output plus // the same FK-free resolution contract as PQL TT-NG LP. - let plan: DiskGraphTTPlan = DiskGraphTTPlan { + let plan: DiskGraphTTPlan = DiskGraphTTPlan { entity: EntitySpec { table: entity_table_name.to_string(), column: entity_pk_column.to_string(), diff --git a/packages/diskgraph/diskgraph-pybind/src/lib.rs b/packages/diskgraph/diskgraph-pybind/src/lib.rs index c907055d1..df8117373 100644 --- a/packages/diskgraph/diskgraph-pybind/src/lib.rs +++ b/packages/diskgraph/diskgraph-pybind/src/lib.rs @@ -1414,6 +1414,8 @@ fn estimate_ingest_size(py: Python<'_>, config_json: &str) -> PyResult { /// invalid. Semantic and logical types are mapped through the same /// ``map_kumo_stype`` / ``map_kumo_dtype`` helpers ingest uses. /// data_version: version string stamped into every emitted entry. +/// timestamp_unit: keyword-only index clock that timestamp statistics are +/// computed in (default ``"seconds"``). /// _category_count_limit_for_testing: keyword-only test hook that lowers /// the retained category limit. ``None`` preserves the production /// 1,000,000-category limit; a supplied value must be greater than 0. @@ -1428,6 +1430,7 @@ fn estimate_ingest_size(py: Python<'_>, config_json: &str) -> PyResult { kumo_column_specs_json, data_version, *, + timestamp_unit="seconds", _category_count_limit_for_testing=None, ))] #[tracing::instrument(level = "debug", skip(py, kumo_column_specs_json))] @@ -1436,6 +1439,7 @@ fn compute_col_stats_for_parquet( parquet_paths: Vec, kumo_column_specs_json: &str, data_version: &str, + timestamp_unit: &str, _category_count_limit_for_testing: Option, ) -> PyResult<(String, String)> { use diskgraph_ingest::config_converter::{ @@ -1505,6 +1509,7 @@ fn compute_col_stats_for_parquet( .collect::>>()?; let data_version = data_version.to_string(); let paths_owned = parquet_paths; + let clock = tt_builder::parse_index_clock(timestamp_unit)?; py.detach(|| -> anyhow::Result<(String, String)> { let paths: Vec<&std::path::Path> = @@ -1515,8 +1520,9 @@ fn compute_col_stats_for_parquet( &columns, &data_version, limit, + clock, ), - None => compute_parquet_col_stats(&paths, &columns, &data_version), + None => compute_parquet_col_stats(&paths, &columns, &data_version, clock), } }) // Route schema failures through the shared exception classifier rather diff --git a/packages/diskgraph/diskgraph-pybind/src/temporal.rs b/packages/diskgraph/diskgraph-pybind/src/temporal.rs index 2b15184e3..7c2441768 100644 --- a/packages/diskgraph/diskgraph-pybind/src/temporal.rs +++ b/packages/diskgraph/diskgraph-pybind/src/temporal.rs @@ -19,21 +19,31 @@ type ResolvedColumn = (Vec>, Option, Option); /// Detect (unless declared) and apply one pyarrow column's temporal hint in a /// single call, so the resolved format never round-trips through Python. -/// Returns epoch seconds — epoch days when `as_days` — aligned with the input, -/// `None` for a null or unparseable value, plus the hint that was resolved. +/// Returns epochs in `timestamp_unit` ticks — epoch days when `as_days` — +/// aligned with the input, `None` for a null or unparseable value, plus the +/// hint that was resolved. #[pyfunction] -#[pyo3(signature = (values, declared_format=None, declared_unit=None, as_days=false))] +#[pyo3(signature = ( + values, + declared_format=None, + declared_unit=None, + as_days=false, + timestamp_unit="seconds", +))] pub fn resolve_temporal( py: Python<'_>, values: &Bound<'_, PyAny>, declared_format: Option, declared_unit: Option, as_days: bool, + timestamp_unit: &str, ) -> PyResult { let array = make_array(ArrayData::from_pyarrow_bound(values)?); let unit = declared_unit.as_deref().and_then(map_timestamp_unit); - let resolved = py - .detach(|| inference::resolve_temporal(&array, declared_format.as_deref(), unit, as_days)); + let clock = crate::tt_builder::parse_index_clock(timestamp_unit)?; + let resolved = py.detach(|| { + inference::resolve_temporal(&array, declared_format.as_deref(), unit, as_days, clock) + }); Ok(( resolved.epochs, resolved.format, diff --git a/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs b/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs index 2b0f97e1f..81255e69e 100644 --- a/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs +++ b/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs @@ -17,8 +17,12 @@ use pyo3::prelude::*; use pyo3::types::PyList; -use diskgraph_core::time::{abs_from_unit, dur_from_unit, Abs, Seconds}; +use diskgraph_core::time::{ + abs_from_unit, ceil_parts, convert_trunc, dur_from_unit, split_parts, Abs, Dur, Seconds, + TimeUnit, UnitSystem, +}; use diskgraph_core::tt_builder as tt; +use diskgraph_core::with_index_clock; // ─── EdgeRef ───────────────────────────────────────────────────────────────── @@ -66,12 +70,19 @@ impl RowFilter { /// RowFilter.compare("status", "==", "paid") /// ``` #[staticmethod] - fn compare(column: String, op: &str, value: &Bound<'_, PyAny>) -> PyResult { + #[pyo3(signature = (column, op, value, timestamp_unit="seconds"))] + fn compare( + column: String, + op: &str, + value: &Bound<'_, PyAny>, + timestamp_unit: &str, + ) -> PyResult { + let clock = parse_index_clock(timestamp_unit)?; Ok(Self { inner: tt::RowFilter::Compare { lhs: tt::RowFilterLhs::Column(column), op: parse_compare_op(op)?, - rhs: tt::RowFilterRhs::Constant(extract_value(value)?), + rhs: tt::RowFilterRhs::Constant(extract_value(value, clock)?), }, }) } @@ -98,12 +109,15 @@ impl RowFilter { /// RowFilter.compare_remote_lhs(order_to_merchant, "country", "==", "US") /// ``` #[staticmethod] + #[pyo3(signature = (path, column, op, value, timestamp_unit="seconds"))] fn compare_remote_lhs( path: &Bound<'_, PyAny>, column: String, op: &str, value: &Bound<'_, PyAny>, + timestamp_unit: &str, ) -> PyResult { + let clock = parse_index_clock(timestamp_unit)?; Ok(Self { inner: tt::RowFilter::Compare { lhs: tt::RowFilterLhs::Remote { @@ -111,7 +125,7 @@ impl RowFilter { column, }, op: parse_compare_op(op)?, - rhs: tt::RowFilterRhs::Constant(extract_value(value)?), + rhs: tt::RowFilterRhs::Constant(extract_value(value, clock)?), }, }) } @@ -206,12 +220,19 @@ impl StageFilter { /// StageFilter.compare("_sum", ">", 0) /// ``` #[staticmethod] - fn compare(alias: String, op: &str, value: &Bound<'_, PyAny>) -> PyResult { + #[pyo3(signature = (alias, op, value, timestamp_unit="seconds"))] + fn compare( + alias: String, + op: &str, + value: &Bound<'_, PyAny>, + timestamp_unit: &str, + ) -> PyResult { + let clock = parse_index_clock(timestamp_unit)?; Ok(Self { inner: tt::StageFilter::Compare { alias, op: parse_compare_op(op)?, - rhs: tt::StageFilterRhs::Constant(extract_value(value)?), + rhs: tt::StageFilterRhs::Constant(extract_value(value, clock)?), }, }) } @@ -310,24 +331,51 @@ impl StageFilter { // ─── DiskGraphTTBuilder ────────────────────────────────────────────────────── +/// Keeps every plan time as given, with its unit; `build()` converts each one +/// once, into `timestamp_unit`. #[pyclass] pub struct DiskGraphTTBuilder { + timestamp_unit: TimeUnit, entity: Option>, - temporal: Option>, + temporal: Option, prediction_type: tt::PredictionType, is_prediction: bool, - split: Option>, - stages: Vec>, - current_ops: Vec>, + split: Option, + stages: Vec<(Vec, Option)>, + current_ops: Vec, target_stype: Option, fkey_mapping: Option>, } +struct TemporalArgs { + timestamps: Vec, + unit: tt::TimeUnit, + forecast_length: i64, + lag_timesteps: Option, +} + +/// Lag and horizon offsets are the only op fields typed by the clock, so they +/// wait for `build()`; every other op is complete when added. +enum BuilderOp { + Ready(tt::Op), + Autoregression(ShiftedOp), + FutureLabels(ShiftedOp), +} + +struct ShiftedOp { + alias: String, + target_op_json: String, + offsets: Vec, + unit: tt::TimeUnit, +} + #[pymethods] impl DiskGraphTTBuilder { #[new] - fn new() -> Self { - Self { + #[pyo3(signature = (timestamp_unit="seconds"))] + fn new(timestamp_unit: &str) -> PyResult { + Ok(Self { + timestamp_unit: parse_index_clock(timestamp_unit)?, entity: None, temporal: None, prediction_type: tt::PredictionType::Node, @@ -337,7 +385,7 @@ impl DiskGraphTTBuilder { current_ops: Vec::new(), target_stype: None, fkey_mapping: None, - } + }) } /// `FOR EACH {table}.{column}`. @@ -410,12 +458,9 @@ impl DiskGraphTTBuilder { lag_timesteps: Option, unit: &str, ) -> PyResult<()> { - let time_unit = parse_time_unit(unit)?; - self.temporal = Some(tt::TemporalSpec { - timestamps: timestamps - .into_iter() - .map(|v| abs_from_unit::(v, time_unit)) - .collect(), + self.temporal = Some(TemporalArgs { + timestamps, + unit: parse_time_unit(unit)?, forecast_length, lag_timesteps, }); @@ -434,7 +479,7 @@ impl DiskGraphTTBuilder { /// The unit rules live in [`make_split_spec`], shared with /// [`build_split_spec_json`] -- see there for why they are asymmetric. fn split_ratio(&mut self, spec: &Bound<'_, PyAny>) -> PyResult<()> { - self.split = Some(make_split_spec(extract_split_spec(spec)?)?); + self.split = Some(extract_split_spec(spec)?); Ok(()) } @@ -501,7 +546,7 @@ impl DiskGraphTTBuilder { let agg_func = parse_agg_func(func)?; let unit = window_unit.map(parse_time_unit).transpose()?; - self.current_ops.push(tt::Op::Aggregate { + self.current_ops.push(BuilderOp::Ready(tt::Op::Aggregate { alias, path: edge_path, func: agg_func, @@ -512,7 +557,7 @@ impl DiskGraphTTBuilder { row_filter: row_filter.map(|f| f.inner), limit, target_entity_table, - }); + })); Ok(()) } @@ -526,12 +571,12 @@ impl DiskGraphTTBuilder { path: Option<&Bound<'_, PyAny>>, ) -> PyResult<()> { let path = path.map(extract_edge_path).transpose()?.unwrap_or_default(); - self.current_ops.push(tt::Op::Column { + self.current_ops.push(BuilderOp::Ready(tt::Op::Column { alias, table, column, path, - }); + })); Ok(()) } @@ -543,12 +588,12 @@ impl DiskGraphTTBuilder { op: &str, value: &Bound<'_, PyAny>, ) -> PyResult<()> { - self.current_ops.push(tt::Op::Compare { + self.current_ops.push(BuilderOp::Ready(tt::Op::Compare { alias, lhs: lhs_alias, op: parse_compare_op(op)?, - rhs: extract_value(value)?, - }); + rhs: extract_value(value, self.timestamp_unit)?, + })); Ok(()) } @@ -566,11 +611,11 @@ impl DiskGraphTTBuilder { } }; - self.current_ops.push(tt::Op::LogicalOp { + self.current_ops.push(BuilderOp::Ready(tt::Op::LogicalOp { alias, op: op_kind, operands, - }); + })); Ok(()) } @@ -583,19 +628,8 @@ impl DiskGraphTTBuilder { lag_offsets: Vec, lag_offset_unit: &str, ) -> PyResult<()> { - let target_op: tt::Op = serde_json::from_str(&target_op_json).map_err(|e| { - crate::DiskgraphConfigError::new_err(format!("Invalid target op JSON: {}", e)) - })?; - - let unit = parse_time_unit(lag_offset_unit)?; - self.current_ops.push(tt::Op::Autoregression { - alias, - target_op: Box::new(target_op), - lag_offsets: lag_offsets - .into_iter() - .map(|v| dur_from_unit::(v, unit)) - .collect(), - }); + let op = ShiftedOp::new(alias, target_op_json, lag_offsets, lag_offset_unit)?; + self.current_ops.push(BuilderOp::Autoregression(op)); Ok(()) } @@ -607,19 +641,8 @@ impl DiskGraphTTBuilder { future_offsets: Vec, future_offset_unit: &str, ) -> PyResult<()> { - let target_op: tt::Op = serde_json::from_str(&target_op_json).map_err(|e| { - crate::DiskgraphConfigError::new_err(format!("Invalid target op JSON: {}", e)) - })?; - - let unit = parse_time_unit(future_offset_unit)?; - self.current_ops.push(tt::Op::FutureLabels { - alias, - target_op: Box::new(target_op), - future_offsets: future_offsets - .into_iter() - .map(|v| dur_from_unit::(v, unit)) - .collect(), - }); + let op = ShiftedOp::new(alias, target_op_json, future_offsets, future_offset_unit)?; + self.current_ops.push(BuilderOp::FutureLabels(op)); Ok(()) } @@ -635,14 +658,14 @@ impl DiskGraphTTBuilder { ) -> PyResult<()> { let path = extract_edge_path(path)?; let unit = window_unit.map(parse_time_unit).transpose()?; - self.current_ops.push(tt::Op::HasEdges { + self.current_ops.push(BuilderOp::Ready(tt::Op::HasEdges { alias, path, time_col, window, window_unit: unit, row_filter: row_filter.map(|f| f.inner.clone()), - }); + })); Ok(()) } @@ -653,38 +676,75 @@ impl DiskGraphTTBuilder { .current_ops .pop() .ok_or_else(|| crate::DiskgraphConfigError::new_err("No operations to pop"))?; - serde_json::to_string(&op).map_err(|e| { - crate::DiskgraphConfigError::new_err(format!("Failed to serialize op: {}", e)) - }) + with_index_clock!(self.timestamp_unit, U => serde_json::to_string(&op.in_clock::()?)) + .map_err(|e| { + crate::DiskgraphConfigError::new_err(format!("Failed to serialize op: {}", e)) + }) } /// Close current stage with a short-circuit filter, start new stage. fn filter(&mut self, expr: StageFilter) { let ops = std::mem::take(&mut self.current_ops); - self.stages.push(tt::Stage { - ops, - filter: Some(expr.inner), - }); + self.stages.push((ops, Some(expr.inner))); } /// Finalise and return `DiskGraphTTPlan` as JSON. fn build(&mut self) -> PyResult { if !self.current_ops.is_empty() { let ops = std::mem::take(&mut self.current_ops); - self.stages.push(tt::Stage { ops, filter: None }); + self.stages.push((ops, None)); } let entity = self.entity.take().ok_or_else(|| { crate::DiskgraphConfigError::new_err("call for_each() before build()") })?; - let plan: tt::DiskGraphTTPlan = tt::DiskGraphTTPlan { - entity, - temporal: self.temporal.take(), + with_index_clock!(self.timestamp_unit, U => self.plan::(entity)?.to_json()) + .map_err(|e| crate::DiskgraphConfigError::new_err(e.to_string())) + } +} + +impl DiskGraphTTBuilder { + fn plan( + &mut self, + entity: tt::EntitySpec, + ) -> PyResult> { + let stages = std::mem::take(&mut self.stages) + .into_iter() + .map(|(ops, filter)| { + let ops = ops + .iter() + .map(BuilderOp::in_clock) + .collect::>()?; + Ok(tt::Stage { ops, filter }) + }) + .collect::>()?; + + Ok(tt::DiskGraphTTPlan { + entity: tt::EntitySpec { + table: entity.table, + column: entity.column, + output_time_col_name: entity.output_time_col_name, + default_time_value: entity + .default_time_value + .map(|secs| abs_from_unit::(secs.raw(), TimeUnit::Seconds)), + time_col: entity.time_col, + valid_to_col: entity.valid_to_col, + filter: entity.filter, + }, + temporal: self.temporal.take().map(|temporal| tt::TemporalSpec { + timestamps: temporal + .timestamps + .into_iter() + .map(|v| abs_from_unit::(v, temporal.unit)) + .collect(), + forecast_length: temporal.forecast_length, + lag_timesteps: temporal.lag_timesteps, + }), prediction_type: std::mem::replace(&mut self.prediction_type, tt::PredictionType::Node), is_prediction: self.is_prediction, - split: self.split.take(), - stages: std::mem::take(&mut self.stages), + split: self.split.take().map(make_split_spec::).transpose()?, + stages, target_stype: self.target_stype.take(), // weight_col is set on custom-TT plans only, via the new // `build_train_table_artifacts_for_existing_table` entry @@ -695,13 +755,53 @@ impl DiskGraphTTBuilder { lp_resolution: None, lp_rhs_is_temporal: false, fact_pseudo_time: None, - }; + }) + } +} - serde_json::to_string(&plan) - .map_err(|e| crate::DiskgraphConfigError::new_err(e.to_string())) +impl BuilderOp { + fn in_clock(&self) -> PyResult> { + Ok(match self { + BuilderOp::Ready(op) => op + .clone() + .map_clock(&|offset| dur_from_unit(offset.raw(), TimeUnit::Seconds)), + BuilderOp::Autoregression(op) => tt::Op::Autoregression { + alias: op.alias.clone(), + target_op: Box::new(parse_op(&op.target_op_json)?), + lag_offsets: op.offsets_in_clock(), + }, + BuilderOp::FutureLabels(op) => tt::Op::FutureLabels { + alias: op.alias.clone(), + target_op: Box::new(parse_op(&op.target_op_json)?), + future_offsets: op.offsets_in_clock(), + }, + }) } } +impl ShiftedOp { + fn new(alias: String, target_op_json: String, offsets: Vec, unit: &str) -> PyResult { + Ok(Self { + alias, + target_op_json, + offsets, + unit: parse_time_unit(unit)?, + }) + } + + fn offsets_in_clock(&self) -> Vec> { + self.offsets + .iter() + .map(|&v| dur_from_unit::(v, self.unit)) + .collect() + } +} + +fn parse_op(json: &str) -> PyResult> { + serde_json::from_str(json) + .map_err(|e| crate::DiskgraphConfigError::new_err(format!("Invalid target op JSON: {}", e))) +} + // ─── Helpers ───────────────────────────────────────────────────────────────── fn parse_compare_op(op: &str) -> PyResult { @@ -835,7 +935,7 @@ fn parse_timestamp_unit(unit: &str) -> PyResult { } } -/// The single implementation of the kumo split config → plan `SplitSpec` +/// The single implementation of the kumo split config → plan `SplitSpec` /// contract. /// /// Both routes into the executor land here: the PQ path via @@ -847,13 +947,13 @@ fn parse_timestamp_unit(unit: &str) -> PyResult { /// /// One implementation is what keeps those two agreeing. The unit handling is /// asymmetric enough to invite drift: `offsets` must be *pre-scaled* into base -/// units here, because `ttgen_executor` reads them as `Dur` and never +/// units here, because `ttgen_executor` reads them as `Dur` and never /// rescales (`offset_unit` survives only as a label — no consumer reads it), /// whereas `timeframe` / `anchor_start_offset` / `timeframe_step` must cross /// *unscaled*, because the executor rescales them itself with /// `dur_from_unit(v, unit)` — and for `TimeRangeSplit` that unit may be /// calendar (`months` / `years`), which has no fixed-duration form. -fn make_split_spec(args: SplitSpecArgs) -> PyResult> { +fn make_split_spec(args: SplitSpecArgs) -> PyResult> { let SplitSpecArgs { kind, ratio, @@ -922,7 +1022,7 @@ fn make_split_spec(args: SplitSpecArgs) -> PyResult> { (_, Some(v)) => ( Some( v.into_iter() - .map(|o| dur_from_unit::(o, off_unit)) + .map(|o| dur_from_unit::(o, off_unit)) .collect::>(), ), None, @@ -951,13 +1051,13 @@ fn make_split_spec(args: SplitSpecArgs) -> PyResult> { v.into_iter() .map(|[lo, hi]| { ( - abs_from_unit::(lo, rng_unit), - abs_from_unit::(hi, rng_unit), + abs_from_unit::(lo, rng_unit), + abs_from_unit::(hi, rng_unit), ) }) .collect() }), - backtesting_timestamp: backtesting_timestamp.map(|v| abs_from_unit::(v, rng_unit)), + backtesting_timestamp: backtesting_timestamp.map(|v| abs_from_unit::(v, rng_unit)), split_column, split_output_column, manual_train_values, @@ -979,8 +1079,8 @@ fn make_split_spec(args: SplitSpecArgs) -> PyResult> { /// Marshal a `SplitSpec` into the standalone JSON that /// `build_train_table_artifacts_for_existing_table` accepts as -/// `split_spec_json`: a serialized `SplitSpec`, which `custom_tt_entry` -/// parses straight back with `serde_json::from_str`. Serialization and +/// `split_spec_json`: a serialized `SplitSpec` in `timestamp_unit`, which +/// `custom_tt_entry` parses straight back with `serde_json::from_str`. Serialization and /// deserialization are the derived impls on one struct, so the round-trip is /// symmetric by construction. /// @@ -989,12 +1089,25 @@ fn make_split_spec(args: SplitSpecArgs) -> PyResult> { /// the call itself -- two hand-maintained argument lists were the weaker /// version of the same failure mode this exists to prevent. #[pyfunction] -pub fn build_split_spec_json(spec: &Bound<'_, PyAny>) -> PyResult { - serde_json::to_string(&make_split_spec(extract_split_spec(spec)?)?).map_err(|e| { +#[pyo3(signature = (spec, timestamp_unit="seconds"))] +pub fn build_split_spec_json(spec: &Bound<'_, PyAny>, timestamp_unit: &str) -> PyResult { + let args = extract_split_spec(spec)?; + with_index_clock!(parse_index_clock(timestamp_unit)?, U => { + serde_json::to_string(&make_split_spec::(args)?) + }) + .map_err(|e| { crate::DiskgraphConfigError::new_err(format!("failed to serialize split spec: {e}")) }) } +pub(crate) fn parse_index_clock(unit: &str) -> PyResult { + TimeUnit::from_name(unit) + .filter(|unit| unit.is_index_clock()) + .ok_or_else(|| { + crate::DiskgraphConfigError::new_err(format!("unsupported timestamp_unit {unit:?}")) + }) +} + fn parse_time_unit(unit: &str) -> PyResult { match unit.to_lowercase().as_str() { "seconds" | "second" | "s" => Ok(tt::TimeUnit::Seconds), @@ -1038,7 +1151,7 @@ fn reject_calendar_date_offset_unit(field: &str, unit: tt::TimeUnit) -> PyResult Ok(()) } -fn extract_value(value: &Bound<'_, PyAny>) -> PyResult { +fn extract_value(value: &Bound<'_, PyAny>, clock: TimeUnit) -> PyResult { // Check bool before int — Python `bool` is a subclass of `int`. if let Ok(b) = value.extract::() { Ok(tt::Value::Bool(b)) @@ -1048,20 +1161,11 @@ fn extract_value(value: &Bound<'_, PyAny>) -> PyResult { Ok(tt::Value::Float(f)) } else if is_datetime(value) { // datetime.datetime (pd.Timestamp is a subclass, handled here too). - // Convert to epoch seconds via Python's .timestamp() method, then - // ceiling to match rowpack's div_ceil storage semantics. - let epoch_f: f64 = value - .call_method0("timestamp") - .map_err(|e| { - crate::DiskgraphConfigError::new_err(format!("datetime.timestamp() failed: {e}")) - })? - .extract::() - .map_err(|e| { - crate::DiskgraphConfigError::new_err(format!( - "datetime.timestamp() did not return float: {e}" - )) - })?; - Ok(tt::Value::Int(epoch_f.ceil() as i64)) + // Round up to the index clock to match rowpack's div_ceil storage semantics. + let (secs, nanos) = datetime_parts(value).map_err(|e| { + crate::DiskgraphConfigError::new_err(format!("cannot read datetime {value}: {e}")) + })?; + Ok(tt::Value::Int(ceil_parts(secs, nanos, clock))) } else if is_date(value) { // datetime.date — convert to midnight UTC epoch seconds. // date.toordinal() gives days since 0001-01-01; Unix epoch is ordinal 719163. @@ -1077,7 +1181,11 @@ fn extract_value(value: &Bound<'_, PyAny>) -> PyResult { )) })?; let days_since_epoch = ordinal - 719_163; - Ok(tt::Value::Int(days_since_epoch * 86_400)) + Ok(tt::Value::Int(convert_trunc( + days_since_epoch, + TimeUnit::Days, + clock, + ))) } else if let Ok(s) = value.extract::() { Ok(tt::Value::Str(s)) } else if let Ok(values) = value.cast::() { @@ -1086,7 +1194,7 @@ fn extract_value(value: &Bound<'_, PyAny>) -> PyResult { // instead of silently dropping a NULL member. let values = values .iter() - .map(|item| extract_value(&item)) + .map(|item| extract_value(&item, clock)) .collect::>>()?; Ok(tt::Value::List(values)) } else { @@ -1096,6 +1204,20 @@ fn extract_value(value: &Bound<'_, PyAny>) -> PyResult { } } +/// Exact epoch seconds and nanoseconds with `.timestamp()`'s time-zone rules +/// (a naive `datetime` is local time, a naive `pd.Timestamp` is UTC). The float +/// `.timestamp()` alone loses a `pd.Timestamp`'s nanoseconds. +fn datetime_parts(value: &Bound<'_, PyAny>) -> PyResult<(i64, u32)> { + if value.hasattr("nanosecond")? { + let nanos: i64 = value.getattr("value")?.extract()?; + return Ok(split_parts(nanos, TimeUnit::Nanoseconds)); + } + let epoch: f64 = value.call_method0("timestamp")?.extract()?; + let micros: u32 = value.getattr("microsecond")?.extract()?; + let secs = (epoch - f64::from(micros) / 1e6).round() as i64; + Ok((secs, micros * 1_000)) +} + /// Check if a Python object is an instance of `datetime.datetime`. /// Must be checked before `datetime.date` since datetime is a subclass of date. fn is_datetime(value: &Bound<'_, PyAny>) -> bool { @@ -1152,7 +1274,7 @@ mod tests { // `split_ratio` now takes a Python object, so exercise the shared // implementation directly -- it is what carries the semantics under // test, and it needs no interpreter. - let split = make_split_spec(SplitSpecArgs { + let split = make_split_spec::(SplitSpecArgs { kind: Some("TimeRangeSplit".to_string()), ratio: None, rnd_seed: None, diff --git a/packages/diskgraph/diskgraph-pybind/src/tt_executor.rs b/packages/diskgraph/diskgraph-pybind/src/tt_executor.rs index f4a7da840..fbe2bdae4 100644 --- a/packages/diskgraph/diskgraph-pybind/src/tt_executor.rs +++ b/packages/diskgraph/diskgraph-pybind/src/tt_executor.rs @@ -31,8 +31,9 @@ use std::path::{Component, Path}; use arrow::pyarrow::ToPyArrow; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::SchemaRef; -use diskgraph_core::time::Seconds; +use diskgraph_core::time::{TimeUnit, UnitSystem}; use diskgraph_core::tt_builder::{DiskGraphTTPlan, RefResolveKind, RefResolveSpec, TTFKeySpec}; +use diskgraph_core::with_index_clock; use diskgraph_train_table::ttgen_executor::{ execute_tt_pass, execute_tt_pass_with_source, CustomTtSource, StageGenInfo, TTPassOutput, STAGE_ARTIFACT_INVENTORY_SCHEMA_VERSION, @@ -109,7 +110,7 @@ pub(crate) struct TtIngestSpec { } impl TtIngestSpec { - pub(crate) fn from_plan(tt_id: &str, plan: &DiskGraphTTPlan) -> Self { + pub(crate) fn from_plan(tt_id: &str, plan: &DiskGraphTTPlan) -> Self { Self { schema_version: TT_INGEST_SPEC_SCHEMA_VERSION.to_string(), tt_id: tt_id.to_string(), @@ -144,6 +145,40 @@ pub(crate) fn rhs_node_type_from_resolution( }) } +/// The clock of the index at `index_path`. +pub(crate) fn index_clock(index_path: &str) -> PyResult { + let metadata = diskgraph_ingest::index_metadata::read_index_metadata(Path::new(index_path)) + .map_err(|e| { + crate::DiskgraphSystemError::new_err(format!("reading index metadata: {e}")) + })?; + let metadata: diskgraph_core::metadata::DatasetMetadata = serde_json::from_str(&metadata) + .map_err(|e| { + crate::DiskgraphConfigError::new_err(format!("invalid index metadata: {e}")) + })?; + if !metadata.timestamp_unit.is_index_clock() { + return Err(crate::DiskgraphConfigError::new_err(format!( + "unsupported index timestamp_unit {:?}", + metadata.timestamp_unit + ))); + } + Ok(metadata.timestamp_unit) +} + +/// Parses a plan, which must be written in the index clock `U`. +fn parse_plan(plan_json: &str) -> PyResult> { + let invalid = |e: serde_json::Error| { + crate::DiskgraphConfigError::new_err(format!("invalid plan JSON: {e}")) + }; + let plan_clock = diskgraph_core::tt_builder::plan_timestamp_unit(plan_json).map_err(invalid)?; + if plan_clock != U::UNIT { + return Err(crate::DiskgraphConfigError::new_err(format!( + "plan is written in {plan_clock:?} but the index clock is {:?}", + U::UNIT + ))); + } + serde_json::from_str(plan_json).map_err(invalid) +} + // ─── execute_tt_plan ───────────────────────────────────────────────────────── /// Run a TT-gen plan against a DiskGraph index. Writes TT artifacts and @@ -162,14 +197,24 @@ pub fn execute_tt_plan( plan_json: &str, index_path: &str, tt_id: &str, +) -> PyResult> { + with_index_clock!(index_clock(index_path)?, U => { + execute_tt_plan_in_clock::(py, plan_json, index_path, tt_id) + }) +} + +fn execute_tt_plan_in_clock( + py: Python<'_>, + plan_json: &str, + index_path: &str, + tt_id: &str, ) -> PyResult> { validate_tt_id(tt_id)?; info!( "execute_tt_plan started (TT-NG end-state path) tt_id={}", tt_id ); - let plan: DiskGraphTTPlan = serde_json::from_str(plan_json) - .map_err(|e| crate::DiskgraphConfigError::new_err(format!("invalid plan JSON: {e}")))?; + let plan: DiskGraphTTPlan = parse_plan(plan_json)?; let idx = index_path.to_owned(); let tt_id_owned = tt_id.to_owned(); @@ -719,8 +764,8 @@ pub(crate) fn resolve_ingest_stages( /// bundle the separate [`execute_tt_ingest`] job consumes. The producers /// differ only in how the plan is built, where the bundle lives, and their /// result-dict shapes — everything bundle-defining is here. -pub(crate) fn run_gen_bundle( - plan: &DiskGraphTTPlan, +pub(crate) fn run_gen_bundle( + plan: &DiskGraphTTPlan, index_root: &Path, out_dir: &Path, tt_id: &str, @@ -764,11 +809,21 @@ pub fn execute_tt_gen( plan_json: &str, index_path: &str, tt_id: &str, +) -> PyResult> { + with_index_clock!(index_clock(index_path)?, U => { + execute_tt_gen_in_clock::(py, plan_json, index_path, tt_id) + }) +} + +fn execute_tt_gen_in_clock( + py: Python<'_>, + plan_json: &str, + index_path: &str, + tt_id: &str, ) -> PyResult> { validate_tt_id(tt_id)?; info!("execute_tt_gen started (ttgen-only) tt_id={}", tt_id); - let plan: DiskGraphTTPlan = serde_json::from_str(plan_json) - .map_err(|e| crate::DiskgraphConfigError::new_err(format!("invalid plan JSON: {e}")))?; + let plan: DiskGraphTTPlan = parse_plan(plan_json)?; let idx = index_path.to_owned(); let tt_id_owned = tt_id.to_owned(); @@ -1499,8 +1554,18 @@ pub fn execute_bp_plan( index_path: &str, output_path: &str, ) -> PyResult { - let plan: DiskGraphTTPlan = serde_json::from_str(plan_json) - .map_err(|e| crate::DiskgraphConfigError::new_err(format!("invalid plan JSON: {e}")))?; + with_index_clock!(index_clock(index_path)?, U => { + execute_bp_plan_in_clock::(py, plan_json, index_path, output_path) + }) +} + +fn execute_bp_plan_in_clock( + py: Python<'_>, + plan_json: &str, + index_path: &str, + output_path: &str, +) -> PyResult { + let plan: DiskGraphTTPlan = parse_plan(plan_json)?; let idx = index_path.to_owned(); let out = output_path.to_owned(); diff --git a/packages/diskgraph/diskgraph-pybind/src/tt_executor/pseudo_time_tests.rs b/packages/diskgraph/diskgraph-pybind/src/tt_executor/pseudo_time_tests.rs index 334cbc137..b777e2865 100644 --- a/packages/diskgraph/diskgraph-pybind/src/tt_executor/pseudo_time_tests.rs +++ b/packages/diskgraph/diskgraph-pybind/src/tt_executor/pseudo_time_tests.rs @@ -146,6 +146,7 @@ fn build_index_order(root: &Path, reverse: bool) -> (PathBuf, Vec<(i64, i64)>) { ..Default::default() }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "static_lp".to_string(), tables: vec![ table( diff --git a/packages/diskgraph/diskgraph-pybind/src/tt_executor/tests.rs b/packages/diskgraph/diskgraph-pybind/src/tt_executor/tests.rs index 5248f27b0..ccb643101 100644 --- a/packages/diskgraph/diskgraph-pybind/src/tt_executor/tests.rs +++ b/packages/diskgraph/diskgraph-pybind/src/tt_executor/tests.rs @@ -708,6 +708,7 @@ unit: "s".to_string(), }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "predict_unify_test".to_string(), tables: vec![ TableSpec { diff --git a/packages/diskgraph/diskgraph-rowpack/Cargo.toml b/packages/diskgraph/diskgraph-rowpack/Cargo.toml index fc04b0ed1..fe23c6bd5 100644 --- a/packages/diskgraph/diskgraph-rowpack/Cargo.toml +++ b/packages/diskgraph/diskgraph-rowpack/Cargo.toml @@ -9,6 +9,7 @@ edition = "2021" license.workspace = true [dependencies] +diskgraph-core = { path = "../diskgraph-core" } anyhow = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/packages/diskgraph/diskgraph-rowpack/src/access.rs b/packages/diskgraph/diskgraph-rowpack/src/access.rs index 15372c987..8c5f65a48 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/access.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/access.rs @@ -12,6 +12,8 @@ use crate::{ use anyhow::{anyhow, Result}; use arrow_array::{Array, ArrayRef}; use arrow_schema::DataType; +use diskgraph_core::time::{ceil_to_clock, floor_parts, Seconds, TimeUnit, UnitSystem}; +use diskgraph_core::with_index_clock; use num_traits::ToPrimitive; use std::sync::Arc; @@ -63,8 +65,14 @@ pub struct RowAccess<'p> { /// Indexed by column index; `None` for columns that don't need it. /// kumo-ai/kumo#28478. tokenized_lists: Vec>, + /// The index clock that `get_timestamp` returns ticks of. + clock: TimeUnit, + /// `timestamp_in` for `clock`, chosen once rather than per value. + read_timestamp: ReadTimestamp<'p>, } +type ReadTimestamp<'p> = fn(&RowAccess<'p>, usize) -> Result>; + pub enum FloatListValues { F16(ArrayRef), F32(ArrayRef), @@ -778,9 +786,22 @@ impl<'p> RowAccess<'p> { hints, fast_parsers, tokenized_lists: vec![None; n], + clock: TimeUnit::Seconds, + read_timestamp: Self::timestamp_in::, } } + pub fn with_clock(mut self, clock: TimeUnit) -> Self { + self.clock = clock; + self.read_timestamp = + with_index_clock!(clock, U => Self::timestamp_in:: as ReadTimestamp<'p>); + self + } + + pub fn clock(&self) -> TimeUnit { + self.clock + } + pub fn set_cols(&mut self, cols: &[ArrayRef]) { self.cols.clear(); self.cols.extend(cols.iter().map(plain_array)); @@ -2067,7 +2088,7 @@ impl<'p> RowAccess<'p> { return Ok(Some(with_stack_str(value, trimmed, f)?)); } DataType::Timestamp(_, _) => { - let secs = timestamp_secs_from_array(array, self.row)?; + let secs = timestamp_from_array::(array, self.row)?; return Ok(Some(with_stack_str(secs, trimmed, f)?)); } DataType::Decimal128(_, scale) => { @@ -2392,7 +2413,7 @@ impl<'p> RowAccess<'p> { return Ok(Some(days)); } if let Some(value) = parse_i64(text) { - let secs = timestamp_from_int(value, unit); + let secs = in_clock(timestamp_from_int::(value, unit))?; return Ok(Some(days_from_secs(secs))); } Err(anyhow!("date parse failed")) @@ -2483,11 +2504,11 @@ impl<'p> RowAccess<'p> { } _ => return Err(anyhow!("date type mismatch")), }; - let secs = timestamp_from_int(value, unit); + let secs = in_clock(timestamp_from_int::(value, unit))?; Ok(Some(days_from_secs(secs))) } DataType::Timestamp(_, _) => { - let secs = timestamp_secs_from_array(array, self.row)?; + let secs = timestamp_from_array::(array, self.row)?; Ok(Some(days_from_secs(secs))) } _ => Err(anyhow!("date type mismatch")), @@ -2495,6 +2516,10 @@ impl<'p> RowAccess<'p> { } pub fn get_timestamp(&self, idx: usize) -> Result> { + (self.read_timestamp)(self, idx) + } + + pub fn timestamp_in(&self, idx: usize) -> Result> { let array = self.array(idx)?; if array.is_null(self.row) { return Ok(None); @@ -2504,13 +2529,9 @@ impl<'p> RowAccess<'p> { let unit = hint.timestamp_unit; let fast = self.fast_parsers.get(idx).and_then(|v| v.as_ref()); match array.data_type() { - DataType::Timestamp(_, _) => { - let secs = timestamp_secs_from_array(array, self.row)?; - Ok(Some(secs)) - } + DataType::Timestamp(_, _) => Ok(Some(timestamp_from_array::(array, self.row)?)), DataType::Time32(_) | DataType::Time64(_) => { - let secs = time_secs_from_array(array, self.row)?; - Ok(Some(secs)) + Ok(Some(time_from_array::(array, self.row)?)) } DataType::Utf8 | DataType::LargeUtf8 @@ -2552,18 +2573,43 @@ impl<'p> RowAccess<'p> { _ => return Err(anyhow!("timestamp type mismatch")), }; if let Some(parser) = fast { - if let Some(secs) = parser.parse_timestamp(text) { - return Ok(Some(secs)); + if let Some((secs, fraction)) = parser.parse_timestamp_parts(text) { + let ticks = floor_parts::(secs, || fraction_nanos(fraction)); + return Ok(Some(in_clock(ticks)?)); } } - if let Some(secs) = parse_timestamp_text(text, format) { - return Ok(Some(secs)); + if let Some((secs, nanos)) = parse_timestamp_text_parts(text, format) { + return Ok(Some(in_clock(floor_parts::(secs, || nanos))?)); } if let Some(value) = parse_i64(text) { - return Ok(Some(timestamp_from_int(value, unit))); + return Ok(Some(in_clock(timestamp_from_int::(value, unit))?)); } Err(anyhow!("timestamp parse failed")) } + DataType::Decimal128(_, scale) if format.is_none() => { + let value = array + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("decimal128 cast"))? + .value(self.row); + Ok(Some(in_clock(timestamp_from_decimal::( + *scale, + unit, + |scale| apply_decimal_scale_i128(value, scale).to_i64(), + ))?)) + } + DataType::Decimal256(_, scale) if format.is_none() => { + let value = array + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("decimal256 cast"))? + .value(self.row); + Ok(Some(in_clock(timestamp_from_decimal::( + *scale, + unit, + |scale| apply_decimal_scale_i256(value, scale).to_i64(), + ))?)) + } DataType::Int8 | DataType::Int16 | DataType::Int32 @@ -2655,11 +2701,11 @@ impl<'p> RowAccess<'p> { // integers like YYYYMMDD (20240301) that are NOT epoch values. if let Some(fmt) = format { let text = value.to_string(); - if let Some(secs) = parse_timestamp_text(&text, Some(fmt)) { - return Ok(Some(secs)); + if let Some((secs, nanos)) = parse_timestamp_text_parts(&text, Some(fmt)) { + return Ok(Some(in_clock(floor_parts::(secs, || nanos))?)); } } - Ok(Some(timestamp_from_int(value, unit))) + Ok(Some(in_clock(timestamp_from_int::(value, unit))?)) } DataType::Date32 => { let arr = array @@ -2675,9 +2721,8 @@ impl<'p> RowAccess<'p> { if !(-106752..=24855).contains(&days) { return Err(anyhow!("date out of range: day {days}")); } - // Convert days since epoch to seconds since epoch - let secs = days as i64 * 86400; - Ok(Some(secs)) + let ticks = ceil_to_clock::(i64::from(days), TimeUnit::Days); + Ok(Some(in_clock(ticks)?)) } _ => Err(anyhow!("timestamp type mismatch")), } @@ -3692,6 +3737,60 @@ mod tests { } } + #[test] + fn timestamps_round_like_main_in_every_index_clock() -> Result<()> { + use diskgraph_core::time::TimeUnit as Clock; + let fractional_text = row_access_for( + Arc::new(StringArray::from(vec!["2024-01-15 10:00:00.750"])), + RowAccessHints { + timestamp_format: Some("%Y-%m-%d %H:%M:%S%.f".into()), + ..RowAccessHints::default() + }, + ); + let epoch_millis = row_access_for( + Arc::new(Int64Array::from(vec![1_705_312_800_750])), + ts_hint(TimeUnit::Millisecond), + ); + let fractional_decimal = row_access_for( + Arc::new( + Decimal128Array::from(vec![170_531_280_075]) + .with_precision_and_scale(20, 2) + .unwrap(), + ), + RowAccessHints::default(), + ); + let expected = [ + ( + Clock::Nanoseconds, + 1_705_312_800_750_000_000, + 1_705_312_800_750_000_000, + 1_705_312_800_750_000_000, + ), + ( + Clock::Microseconds, + 1_705_312_800_750_000, + 1_705_312_800_750_000, + 1_705_312_800_750_000, + ), + ( + Clock::Milliseconds, + 1_705_312_800_750, + 1_705_312_800_750, + 1_705_312_800_750, + ), + (Clock::Seconds, 1_705_312_800, 1_705_312_801, 1_705_312_800), + (Clock::Hours, 473_698, 473_699, 473_698), + (Clock::Days, 19_737, 19_738, 19_738), + ]; + for (clock, text, millis, decimal) in expected { + let read = |access: &RowAccess<'static>| with_index_clock!(clock, U => access.timestamp_in::(0)); + assert_eq!(read(&fractional_text)?, Some(text), "{clock:?}"); + assert_eq!(read(&epoch_millis)?, Some(millis), "{clock:?}"); + assert_eq!(read(&fractional_decimal)?, Some(decimal), "{clock:?}"); + } + Ok(()) + } + #[test] fn i64_backed_conversions_reject_unsigned_and_decimal_overflow() { let overflow = i64::MAX as u64 + 1; diff --git a/packages/diskgraph/diskgraph-rowpack/src/lib.rs b/packages/diskgraph/diskgraph-rowpack/src/lib.rs index 533f687e0..49862ea0c 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/lib.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/lib.rs @@ -6,6 +6,7 @@ use anyhow::{anyhow, Result}; use arrow_array::{Array, RecordBatch}; use base64::Engine; use chrono::{DateTime, Utc}; +use diskgraph_core::time::{convert_trunc, split_parts, TimeUnit}; use gxhash::{HashMap, HashMapExt}; use half::f16; use half::slice::HalfFloatSliceExt; @@ -415,7 +416,7 @@ pub enum ValueRef<'a> { Uuid256([u8; 32]), Ip([u8; 16]), Date(i32), - /// Epoch seconds. + /// Ticks of the index clock. Timestamp(i64), Decimal128(i128), Decimal256([u8; 32]), @@ -5208,6 +5209,7 @@ fn estimate_bpe_bytes(samples: &[Vec], dict: &[Vec]) -> u64 { pub struct RowPackDecoder { ctx: RowPackContext, dict_values: Vec>>>, + clock: TimeUnit, } pub struct RowPackDecodeScratch { @@ -5258,7 +5260,20 @@ impl RowPackDecoder { _ => dict_values.push(None), } } - Ok(Self { ctx, dict_values }) + Ok(Self { + ctx, + dict_values, + clock: TimeUnit::Seconds, + }) + } + + pub fn with_clock(mut self, clock: TimeUnit) -> Self { + self.clock = clock; + self + } + + pub fn clock(&self) -> TimeUnit { + self.clock } pub fn context(&self) -> &RowPackContext { @@ -5306,6 +5321,7 @@ impl RowPackDecoder { tag, &mut data, self.dict_values[idx].as_ref(), + self.clock, )?; } Ok(out) @@ -5341,7 +5357,13 @@ impl RowPackDecoder { bool_idx += 1; JsonValue::Bool(v) } else { - decode_value_json(col, tag, &mut data, self.dict_values[idx].as_ref())? + decode_value_json( + col, + tag, + &mut data, + self.dict_values[idx].as_ref(), + self.clock, + )? }; map.insert(col.name.clone(), value); } @@ -5562,6 +5584,7 @@ impl RowPackDecoder { &mut data, self.dict_values[idx].as_ref(), scratch, + self.clock, )?; builder_cursor += 1; } else { @@ -5603,7 +5626,8 @@ impl RowPackDecoder { if !requested_mask[idx] { continue; } - let (field, builder) = make_arrow_builder(&col.name, col.logical_type, rows.len())?; + let (field, builder) = + make_arrow_builder(&col.name, col.logical_type, rows.len(), self.clock)?; decoded_indices.push(idx); fields.push(field); builders.push(builder); @@ -5693,6 +5717,7 @@ fn decode_value_msgpack( tag: u8, data: &mut &[u8], dict: Option<&Vec>>, + clock: TimeUnit, ) -> Result<()> { match &col.encoding { ColumnEncoding::IntDict { @@ -5777,8 +5802,8 @@ fn decode_value_msgpack( mp::write_sint(out, v)?; } ColumnEncoding::Timestamp => { - let secs = zigzag_decode(read_varint(data)?); - write_msgpack_timestamp(out, secs, 0)?; + let (secs, nanos) = split_parts(zigzag_decode(read_varint(data)?), clock); + write_msgpack_timestamp(out, secs, nanos)?; } ColumnEncoding::Decimal { scale, @@ -5874,6 +5899,7 @@ fn decode_value_json( tag: u8, data: &mut &[u8], dict: Option<&Vec>>, + clock: TimeUnit, ) -> Result { let value = match &col.encoding { ColumnEncoding::IntDict { @@ -5952,8 +5978,8 @@ fn decode_value_json( } ColumnEncoding::Date { varint, bits } => JsonValue::from(decode_int(data, *varint, *bits)?), ColumnEncoding::Timestamp => { - let secs = zigzag_decode(read_varint(data)?); - let formatted = format_timestamp_default(secs, 0); + let (secs, nanos) = split_parts(zigzag_decode(read_varint(data)?), clock); + let formatted = format_timestamp_default(secs, nanos); JsonValue::from(formatted) } ColumnEncoding::Decimal { @@ -7189,6 +7215,27 @@ fn append_null(builder: &mut dyn arrow_array::builder::ArrayBuilder) { b.append_null(); return; } + if let Some(b) = builder + .as_any_mut() + .downcast_mut::() + { + b.append_null(); + return; + } + if let Some(b) = builder + .as_any_mut() + .downcast_mut::() + { + b.append_null(); + return; + } + if let Some(b) = builder + .as_any_mut() + .downcast_mut::() + { + b.append_null(); + return; + } if let Some(b) = builder.as_any_mut().downcast_mut::() { b.append_null(); return; @@ -7226,12 +7273,13 @@ pub fn make_arrow_builder( name: &str, dtype: IngestedType, capacity: usize, + clock: TimeUnit, ) -> Result<( arrow_schema::Field, Box, )> { use arrow_array::builder::*; - use arrow_schema::{DataType, TimeUnit}; + use arrow_schema::DataType; let (arrow_type, builder): (DataType, Box) = match dtype { IngestedType::Bool => ( @@ -7291,10 +7339,25 @@ pub fn make_arrow_builder( DataType::Binary, Box::new(BinaryBuilder::with_capacity(capacity, capacity * 32)), ), - IngestedType::Timestamp => ( - DataType::Timestamp(TimeUnit::Second, None), - Box::new(TimestampSecondBuilder::with_capacity(capacity)), - ), + IngestedType::Timestamp => match clock { + TimeUnit::Nanoseconds => ( + DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, None), + Box::new(TimestampNanosecondBuilder::with_capacity(capacity)), + ), + TimeUnit::Microseconds => ( + DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None), + Box::new(TimestampMicrosecondBuilder::with_capacity(capacity)), + ), + TimeUnit::Milliseconds => ( + DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + Box::new(TimestampMillisecondBuilder::with_capacity(capacity)), + ), + TimeUnit::Seconds | TimeUnit::Hours | TimeUnit::Days => ( + DataType::Timestamp(arrow_schema::TimeUnit::Second, None), + Box::new(TimestampSecondBuilder::with_capacity(capacity)), + ), + other => unreachable!("{other:?} is not an index clock"), + }, IngestedType::Date => ( DataType::Date32, Box::new(Date32Builder::with_capacity(capacity)), @@ -7360,6 +7423,7 @@ fn decode_and_append_to_builder( data: &mut &[u8], dict: Option<&Vec>>, scratch: &mut RowPackDecodeScratch, + clock: TimeUnit, ) -> Result<()> { use arrow_array::builder::*; @@ -7491,12 +7555,27 @@ fn decode_and_append_to_builder( .append_value(v as i32); } ColumnEncoding::Timestamp => { - let secs = zigzag_decode(read_varint(data)?); - builder - .as_any_mut() - .downcast_mut::() - .unwrap() - .append_value(secs); + let ticks = zigzag_decode(read_varint(data)?); + let builder = builder.as_any_mut(); + match clock { + TimeUnit::Nanoseconds => builder + .downcast_mut::() + .unwrap() + .append_value(ticks), + TimeUnit::Microseconds => builder + .downcast_mut::() + .unwrap() + .append_value(ticks), + TimeUnit::Milliseconds => builder + .downcast_mut::() + .unwrap() + .append_value(ticks), + TimeUnit::Seconds | TimeUnit::Hours | TimeUnit::Days => builder + .downcast_mut::() + .unwrap() + .append_value(convert_trunc(ticks, clock, TimeUnit::Seconds)), + other => unreachable!("{other:?} is not an index clock"), + } } ColumnEncoding::Decimal { scale, diff --git a/packages/diskgraph/diskgraph-rowpack/src/parser.rs b/packages/diskgraph/diskgraph-rowpack/src/parser.rs index 9701ccb68..e838d4299 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/parser.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/parser.rs @@ -7,6 +7,7 @@ use anyhow::{anyhow, Result}; use arrow_array::{Array, ArrayRef}; use arrow_schema::DataType; use chrono::{NaiveDate, NaiveDateTime}; +use diskgraph_core::time::{ceil_to_clock, div_ceil, TimeUnit, UnitSystem}; use uuid_simd::{parse_hyphenated, parse_simple, AsOut}; pub(crate) fn parse_bool(value: &str) -> Option { @@ -242,6 +243,10 @@ fn parse_rfc3339_utc(value: &str) -> Option { } pub(crate) fn parse_timestamp_text(value: &str, format: Option<&str>) -> Option { + parse_timestamp_text_parts(value, format).map(|(secs, _)| secs) +} + +pub(crate) fn parse_timestamp_text_parts(value: &str, format: Option<&str>) -> Option<(i64, u32)> { let ts = if let Some(fmt) = format { if format_has_tz(fmt) { chrono::DateTime::parse_from_str(value, fmt) @@ -284,7 +289,7 @@ pub(crate) fn parse_timestamp_text(value: &str, format: Option<&str>) -> Option< })? }; let ts = ts.and_utc(); - Some(ts.timestamp()) + Some((ts.timestamp(), ts.timestamp_subsec_nanos())) } pub(crate) fn parse_date_text(value: &str, format: Option<&str>) -> Option { @@ -367,9 +372,13 @@ impl TimestampFormatParser { TimestampFormatKind::Ymd => parse_ymd_fast(value.as_bytes()), TimestampFormatKind::YmdHmSpace => parse_ymd_hm_space_fast(value.as_bytes()), TimestampFormatKind::YmdHmsSpace => parse_ymd_hms_space_fast(value.as_bytes()), - TimestampFormatKind::YmdHmsSpaceFrac => parse_ymd_hms_space_frac_fast(value.as_bytes()), + TimestampFormatKind::YmdHmsSpaceFrac => { + parse_ymd_hms_space_frac_fast(value.as_bytes()).map(|(secs, _)| secs) + } TimestampFormatKind::YmdHmsT => parse_ymd_hms_space_fast(value.as_bytes()), - TimestampFormatKind::YmdHmsTFrac => parse_ymd_hms_space_frac_fast(value.as_bytes()), + TimestampFormatKind::YmdHmsTFrac => { + parse_ymd_hms_space_frac_fast(value.as_bytes()).map(|(secs, _)| secs) + } TimestampFormatKind::YmdHmsSpaceOffset => { parse_ymd_hms_offset_nocolon_fast(value.as_bytes()) } @@ -393,6 +402,16 @@ impl TimestampFormatParser { } } + /// Whole epoch seconds and the digits of any fractional second. + pub(crate) fn parse_timestamp_parts<'a>(&self, value: &'a str) -> Option<(i64, &'a [u8])> { + match self.kind { + TimestampFormatKind::YmdHmsSpaceFrac | TimestampFormatKind::YmdHmsTFrac => { + parse_ymd_hms_space_frac_fast(value.as_bytes()) + } + _ => Some((self.parse_timestamp(value)?, &[])), + } + } + pub(crate) fn parse_date(&self, value: &str) -> Option { let secs = self.parse_timestamp(value)?; Some(days_from_secs(secs)) @@ -702,7 +721,7 @@ pub(crate) fn parse_ymd_hms_space_fast(bytes: &[u8]) -> Option { } // "%Y-%m-%d %H:%M:%S%.f" / "%Y-%m-%dT%H:%M:%S%.f" -pub(crate) fn parse_ymd_hms_space_frac_fast(bytes: &[u8]) -> Option { +pub(crate) fn parse_ymd_hms_space_frac_fast(bytes: &[u8]) -> Option<(i64, &[u8])> { if bytes.len() < 21 { return None; } @@ -714,9 +733,25 @@ pub(crate) fn parse_ymd_hms_space_frac_fast(bytes: &[u8]) -> Option { if next != bytes.len() { return None; } - // Sub-second fraction discarded — standardized to epoch seconds. let secs = timestamp_from_parts_fast(year as i32, month, day, hour, minute, second); - Some(secs) + Some((secs, &bytes[20..])) +} + +/// Nanoseconds in a fractional second written as `digits` (beyond nine truncate). +pub(crate) fn fraction_nanos(digits: &[u8]) -> u32 { + const NANOS_PER_LAST_DIGIT: [u32; 10] = [ + 0, + 100_000_000, + 10_000_000, + 1_000_000, + 100_000, + 10_000, + 1_000, + 100, + 10, + 1, + ]; + parse_fraction(digits, 0).map_or(0, |(value, digits, _)| value * NANOS_PER_LAST_DIGIT[digits]) } #[inline] @@ -850,29 +885,42 @@ pub(crate) fn days_from_civil_fast(year: i32, month: u32, day: u32) -> i64 { era * 146097 + doe - 719468 } -/// Ceiling division for signed integers with a positive divisor. -/// Rounds toward positive infinity — ensures sub-second timestamps are -/// never assigned to an earlier second, preventing future data leakage -/// in temporal aggregations. -#[inline] -fn div_ceil_i64(value: i64, divisor: i64) -> i64 { - let d = value.div_euclid(divisor); - let r = value.rem_euclid(divisor); - if r > 0 { d + 1 } else { d } +/// Fails for a value the index clock cannot represent; callers store a null. +pub(crate) fn in_clock(ticks: Option) -> Result { + ticks.ok_or_else(|| anyhow!("timestamp does not fit the index clock")) } -pub(crate) fn timestamp_from_int(value: i64, unit: Option) -> i64 { - let unit = unit.unwrap_or(arrow_schema::TimeUnit::Second); - match unit { - arrow_schema::TimeUnit::Second => value, - arrow_schema::TimeUnit::Millisecond => div_ceil_i64(value, 1_000), - arrow_schema::TimeUnit::Microsecond => div_ceil_i64(value, 1_000_000), - arrow_schema::TimeUnit::Nanosecond => div_ceil_i64(value, 1_000_000_000), +#[inline(always)] +pub(crate) fn timestamp_from_int( + value: i64, + unit: Option, +) -> Option { + match unit.unwrap_or(arrow_schema::TimeUnit::Second) { + arrow_schema::TimeUnit::Second => ceil_to_clock::(value, TimeUnit::Seconds), + arrow_schema::TimeUnit::Millisecond => ceil_to_clock::(value, TimeUnit::Milliseconds), + arrow_schema::TimeUnit::Microsecond => ceil_to_clock::(value, TimeUnit::Microseconds), + arrow_schema::TimeUnit::Nanosecond => ceil_to_clock::(value, TimeUnit::Nanoseconds), } } +/// A decimal epoch in `unit` as ticks of the clock `U`, where `scaled(s)` is +/// its unscaled value divided by `10^s` and truncated. It truncates to the +/// finer of `unit` and the tick, then rounds up like an integer epoch. +#[inline(always)] +pub(crate) fn timestamp_from_decimal( + scale: i8, + unit: Option, + scaled: impl FnOnce(i32) -> Option, +) -> Option { + let ticks_per_unit = timestamp_from_int::(1, unit)?; + if ticks_per_unit > 1 { + return scaled(i32::from(scale) - ticks_per_unit.ilog10() as i32); + } + timestamp_from_int::(scaled(i32::from(scale))?, unit) +} + pub(crate) fn days_from_secs(secs: i64) -> i32 { - div_ceil_i64(secs, 86_400) as i32 + div_ceil(secs, 86_400) as i32 } pub(crate) fn parse_decimal_i128(value: &str, scale: i32) -> Option { @@ -942,7 +990,8 @@ pub(crate) fn parse_decimal_i256(value: &str, scale: i32) -> Option Result { +#[inline] +pub(crate) fn timestamp_from_array(array: &ArrayRef, row: usize) -> Result { let unit = match array.data_type() { DataType::Timestamp(unit, _) => *unit, _ => return Err(anyhow!("timestamp type mismatch")), @@ -953,61 +1002,68 @@ pub(crate) fn timestamp_secs_from_array(array: &ArrayRef, row: usize) -> Result< .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("timestamp second cast"))?; - Ok(arr.value(row)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Seconds)) } arrow_schema::TimeUnit::Millisecond => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("timestamp millisecond cast"))?; - Ok(div_ceil_i64(arr.value(row), 1_000)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Milliseconds)) } arrow_schema::TimeUnit::Microsecond => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("timestamp microsecond cast"))?; - Ok(div_ceil_i64(arr.value(row), 1_000_000)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Microseconds)) } arrow_schema::TimeUnit::Nanosecond => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("timestamp nanosecond cast"))?; - Ok(div_ceil_i64(arr.value(row), 1_000_000_000)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) } } } -pub(crate) fn time_secs_from_array(array: &ArrayRef, row: usize) -> Result { +#[inline] +pub(crate) fn time_from_array(array: &ArrayRef, row: usize) -> Result { match array.data_type() { DataType::Time32(arrow_schema::TimeUnit::Second) => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time32 second cast"))?; - Ok(i64::from(arr.value(row))) + in_clock(ceil_to_clock::( + i64::from(arr.value(row)), + TimeUnit::Seconds, + )) } DataType::Time32(arrow_schema::TimeUnit::Millisecond) => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time32 millisecond cast"))?; - Ok(div_ceil_i64(i64::from(arr.value(row)), 1_000)) + in_clock(ceil_to_clock::( + i64::from(arr.value(row)), + TimeUnit::Milliseconds, + )) } DataType::Time64(arrow_schema::TimeUnit::Microsecond) => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time64 microsecond cast"))?; - Ok(div_ceil_i64(arr.value(row), 1_000_000)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Microseconds)) } DataType::Time64(arrow_schema::TimeUnit::Nanosecond) => { let arr = array .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time64 nanosecond cast"))?; - Ok(div_ceil_i64(arr.value(row), 1_000_000_000)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) } DataType::Time32(_) | DataType::Time64(_) => Err(anyhow!( "unsupported Arrow time unit for timestamp conversion: {:?}", diff --git a/packages/diskgraph/diskgraph-rowpack/src/tests_int_string_list.rs b/packages/diskgraph/diskgraph-rowpack/src/tests_int_string_list.rs index 272b838f6..9c148c737 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/tests_int_string_list.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/tests_int_string_list.rs @@ -1487,7 +1487,7 @@ fn test_int8_array_rejects_i16_overflow() -> anyhow::Result<()> { fn test_int32_array_arrow_output_schema_uses_int32_child() -> anyhow::Result<()> { use arrow_schema::DataType; let (field, _builder) = - crate::make_arrow_builder("tags", IngestedType::Int32Array, 16)?; + crate::make_arrow_builder("tags", IngestedType::Int32Array, 16, TimeUnit::Seconds)?; match field.data_type() { DataType::List(child) => assert_eq!(*child.data_type(), DataType::Int32), other => panic!("expected List, got {:?}", other), @@ -1498,7 +1498,8 @@ fn test_int32_array_arrow_output_schema_uses_int32_child() -> anyhow::Result<()> #[test] fn test_int8_array_arrow_output_schema_uses_int8_child() -> anyhow::Result<()> { use arrow_schema::DataType; - let (field, _) = crate::make_arrow_builder("tags", IngestedType::Int8Array, 16)?; + let (field, _) = + crate::make_arrow_builder("tags", IngestedType::Int8Array, 16, TimeUnit::Seconds)?; match field.data_type() { DataType::List(child) => assert_eq!(*child.data_type(), DataType::Int8), other => panic!("expected List, got {:?}", other), @@ -1509,7 +1510,8 @@ fn test_int8_array_arrow_output_schema_uses_int8_child() -> anyhow::Result<()> { #[test] fn test_int16_array_arrow_output_schema_uses_int16_child() -> anyhow::Result<()> { use arrow_schema::DataType; - let (field, _) = crate::make_arrow_builder("tags", IngestedType::Int16Array, 16)?; + let (field, _) = + crate::make_arrow_builder("tags", IngestedType::Int16Array, 16, TimeUnit::Seconds)?; match field.data_type() { DataType::List(child) => assert_eq!(*child.data_type(), DataType::Int16), other => panic!("expected List, got {:?}", other), diff --git a/packages/diskgraph/diskgraph-train-table/src/custom_tt_chunk_processor.rs b/packages/diskgraph/diskgraph-train-table/src/custom_tt_chunk_processor.rs index 2ce3c1e0e..06761a5be 100644 --- a/packages/diskgraph/diskgraph-train-table/src/custom_tt_chunk_processor.rs +++ b/packages/diskgraph/diskgraph-train-table/src/custom_tt_chunk_processor.rs @@ -50,7 +50,7 @@ use crossbeam_channel::Sender; use diskgraph_core::accumulators::{ HyperLogLog, LocalColumnCounts, SharedColumnCounts, TDigestEstimator, WelfordAccumulator, }; -use diskgraph_core::time::{Abs, TimeUnit, UnitSystem}; +use diskgraph_core::time::{convert_trunc, Abs, TimeUnit, UnitSystem}; use diskgraph_core::tt_builder::{PredictionType, Value}; use diskgraph_index::pkmap; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; @@ -386,44 +386,46 @@ pub(crate) fn cell_to_value(arr: &dyn Array, idx: usize) -> Option { None } -/// Read one cell of an Arrow timestamp column as epoch seconds. +/// Read one cell of an Arrow timestamp column as ticks of `clock`, truncating. /// /// Dispatches by the column's typed Arrow `DataType` — no inference of /// units from value magnitude. Supported: -/// * `Timestamp(Nanosecond)` → value / 1e9 (declared ignored; array self-describes) -/// * `Timestamp(Microsecond)` → value / 1e6 (declared ignored) -/// * `Timestamp(Millisecond)` → value / 1e3 (declared ignored) -/// * `Timestamp(Second)` → value as-is (declared ignored) -/// * `Int64` → divide by `declared`-derived factor +/// * `Timestamp(_)` → the array's own unit (declared ignored; array self-describes) +/// * `Int64` → the `declared` unit /// /// For `Int64` (no inherent unit), `declared` is the bridge-supplied /// `CustomTtSource::input_timestamp_unit`. Kumo's rehydration declares /// this as `Nanoseconds` because kumo's canonical TIMESTAMP is int64 ns. /// Non-sub-second variants (Minutes, Hours, ...) are not meaningful for /// epoch timestamps and return `None`. -fn read_timestamp_secs(arr: &dyn Array, idx: usize, declared: TimeUnit) -> Option { +fn read_timestamp_ticks( + arr: &dyn Array, + idx: usize, + declared: TimeUnit, + clock: TimeUnit, +) -> Option { if arr.is_null(idx) { return None; } if let Some(a) = arr.as_any().downcast_ref::() { - return Some(a.value(idx) / 1_000_000_000); + return Some(convert_trunc(a.value(idx), TimeUnit::Nanoseconds, clock)); } if let Some(a) = arr.as_any().downcast_ref::() { - return Some(a.value(idx) / 1_000_000); + return Some(convert_trunc(a.value(idx), TimeUnit::Microseconds, clock)); } if let Some(a) = arr.as_any().downcast_ref::() { - return Some(a.value(idx) / 1_000); + return Some(convert_trunc(a.value(idx), TimeUnit::Milliseconds, clock)); } if let Some(a) = arr.as_any().downcast_ref::() { - return Some(a.value(idx)); + return Some(convert_trunc(a.value(idx), TimeUnit::Seconds, clock)); } if let Some(a) = arr.as_any().downcast_ref::() { let v = a.value(idx); return match declared { - TimeUnit::Nanoseconds => Some(v / 1_000_000_000), - TimeUnit::Microseconds => Some(v / 1_000_000), - TimeUnit::Milliseconds => Some(v / 1_000), - TimeUnit::Seconds => Some(v), + TimeUnit::Nanoseconds + | TimeUnit::Microseconds + | TimeUnit::Milliseconds + | TimeUnit::Seconds => Some(convert_trunc(v, declared, clock)), // Non-sub-second units are not used for epoch timestamps in // kumo's data model; surface as None so the caller's "skip // null timestamp" path applies. @@ -1204,11 +1206,14 @@ fn process_batch>( sentinel_ts } Some(arr) => { - let ts_secs = read_timestamp_secs(arr.as_ref(), row_idx, input_timestamp_unit) - .ok_or_else(|| { - TTGenError::SchemaError("custom-TT timestamp cannot be read".to_string()) - })?; - Abs::::from_secs(ts_secs) + let ticks = + read_timestamp_ticks(arr.as_ref(), row_idx, input_timestamp_unit, U::UNIT) + .ok_or_else(|| { + TTGenError::SchemaError( + "custom-TT timestamp cannot be read".to_string(), + ) + })?; + Abs::::from_raw(ticks) } None => sentinel_ts, }; @@ -1280,7 +1285,7 @@ fn process_batch>( record_filter(RowFailureKind::NonCanonicalEntity, entity_column(0))?; continue; } - let ts_secs = ts.to_secs(); + let ts_ticks = ts.raw(); for (col_idx, name) in &sequence_feature_indices { let entry = processor @@ -1450,7 +1455,7 @@ fn process_batch>( { pk_table.contains_pk_value_temporal( primary, - ts_secs, + ts_ticks, plan.custom_target_temporal_create_col_idx, plan.custom_target_temporal_end_col_idx, &mut target_pk_scratch, @@ -1578,7 +1583,7 @@ fn process_batch>( { pk_table.contains_pk_value_temporal( &target, - ts_secs, + ts_ticks, plan.custom_target_temporal_create_col_idx, plan.custom_target_temporal_end_col_idx, &mut target_pk_scratch, @@ -1888,9 +1893,10 @@ mod contract_tests { for (token, array) in scalar_corpus() { let listed = token.is_some_and(|t| TIMESTAMP_TYPES.contains(&t)); assert_eq!( - read_timestamp_secs(array.as_ref(), 0, TimeUnit::Nanoseconds).is_some(), + read_timestamp_ticks(array.as_ref(), 0, TimeUnit::Nanoseconds, TimeUnit::Seconds) + .is_some(), listed, - "TIMESTAMP_TYPES disagrees with read_timestamp_secs for {}", + "TIMESTAMP_TYPES disagrees with read_timestamp_ticks for {}", array.data_type() ); } @@ -2021,7 +2027,7 @@ mod tests { /// Typed `Timestamp(unit)` arrays self-describe their unit; the /// `declared` arg is ignored for them. #[test] - fn read_timestamp_secs_typed_units_ignore_declared() { + fn read_timestamp_ticks_typed_units_ignore_declared() { // Year 2024 (~1.7e9 epoch seconds) expressed in each unit. let secs: i64 = 1_700_000_000; let nanos: i64 = secs * 1_000_000_000; @@ -2033,18 +2039,30 @@ mod tests { let wrong = TimeUnit::Seconds; let ns_arr = TimestampNanosecondArray::from(vec![Some(nanos), None]); - assert_eq!(read_timestamp_secs(&ns_arr, 0, wrong), Some(secs)); - assert_eq!(read_timestamp_secs(&ns_arr, 1, wrong), None); + assert_eq!( + read_timestamp_ticks(&ns_arr, 0, wrong, TimeUnit::Seconds), + Some(secs) + ); + assert_eq!( + read_timestamp_ticks(&ns_arr, 1, wrong, TimeUnit::Seconds), + None + ); let us_arr = TimestampMicrosecondArray::from(vec![Some(micros)]); - assert_eq!(read_timestamp_secs(&us_arr, 0, wrong), Some(secs)); + assert_eq!( + read_timestamp_ticks(&us_arr, 0, wrong, TimeUnit::Seconds), + Some(secs) + ); let ms_arr = TimestampMillisecondArray::from(vec![Some(millis)]); - assert_eq!(read_timestamp_secs(&ms_arr, 0, wrong), Some(secs)); + assert_eq!( + read_timestamp_ticks(&ms_arr, 0, wrong, TimeUnit::Seconds), + Some(secs) + ); let s_arr = TimestampSecondArray::from(vec![Some(secs)]); assert_eq!( - read_timestamp_secs(&s_arr, 0, TimeUnit::Nanoseconds), + read_timestamp_ticks(&s_arr, 0, TimeUnit::Nanoseconds, TimeUnit::Seconds), Some(secs) ); } @@ -2053,7 +2071,7 @@ mod tests { /// `CustomTtSource::input_timestamp_unit`. Cover all four sub-second /// declarations. #[test] - fn read_timestamp_secs_int64_uses_declared() { + fn read_timestamp_ticks_int64_uses_declared() { let secs: i64 = 1_700_000_000; let nanos: i64 = secs * 1_000_000_000; let micros: i64 = secs * 1_000_000; @@ -2061,35 +2079,47 @@ mod tests { let ns = Int64Array::from(vec![Some(nanos)]); assert_eq!( - read_timestamp_secs(&ns, 0, TimeUnit::Nanoseconds), + read_timestamp_ticks(&ns, 0, TimeUnit::Nanoseconds, TimeUnit::Seconds), Some(secs) ); let us = Int64Array::from(vec![Some(micros)]); assert_eq!( - read_timestamp_secs(&us, 0, TimeUnit::Microseconds), + read_timestamp_ticks(&us, 0, TimeUnit::Microseconds, TimeUnit::Seconds), Some(secs) ); let ms = Int64Array::from(vec![Some(millis)]); assert_eq!( - read_timestamp_secs(&ms, 0, TimeUnit::Milliseconds), + read_timestamp_ticks(&ms, 0, TimeUnit::Milliseconds, TimeUnit::Seconds), Some(secs) ); let s = Int64Array::from(vec![Some(secs)]); - assert_eq!(read_timestamp_secs(&s, 0, TimeUnit::Seconds), Some(secs)); + assert_eq!( + read_timestamp_ticks(&s, 0, TimeUnit::Seconds, TimeUnit::Seconds), + Some(secs) + ); } /// Non-sub-second declared units (Minutes, Hours, ...) are not /// meaningful for epoch timestamps — return None so the caller /// skips the row. #[test] - fn read_timestamp_secs_int64_non_sub_second_returns_none() { + fn read_timestamp_ticks_int64_non_sub_second_returns_none() { let int_arr = Int64Array::from(vec![Some(1_700_000_000_i64)]); - assert_eq!(read_timestamp_secs(&int_arr, 0, TimeUnit::Minutes), None); - assert_eq!(read_timestamp_secs(&int_arr, 0, TimeUnit::Hours), None); - assert_eq!(read_timestamp_secs(&int_arr, 0, TimeUnit::Days), None); + assert_eq!( + read_timestamp_ticks(&int_arr, 0, TimeUnit::Minutes, TimeUnit::Seconds), + None + ); + assert_eq!( + read_timestamp_ticks(&int_arr, 0, TimeUnit::Hours, TimeUnit::Seconds), + None + ); + assert_eq!( + read_timestamp_ticks(&int_arr, 0, TimeUnit::Days, TimeUnit::Seconds), + None + ); } #[test] diff --git a/packages/diskgraph/diskgraph-train-table/src/split_decider.rs b/packages/diskgraph/diskgraph-train-table/src/split_decider.rs index c6ece4d3c..922eec0bc 100644 --- a/packages/diskgraph/diskgraph-train-table/src/split_decider.rs +++ b/packages/diskgraph/diskgraph-train-table/src/split_decider.rs @@ -157,7 +157,7 @@ pub(crate) trait SplitDecider: Send + Sync { /// is a refcount bump, no data copy). Then re-downcasts on construction /// and stores both the keep-alive Arc plus the typed pointer-equivalent. /// We avoid self-referential structs by re-running the downcast inside -/// `ts_secs_at` — but that's a single `as_any()` + ptr-cast per call, +/// `ts_at` — but that's a single `as_any()` + ptr-cast per call, /// not a schema-name lookup, and inlines. pub(crate) struct TsBatchView<'a> { /// Reference to the timestamp column (TimestampNanosecondArray or Int64Array). @@ -187,13 +187,13 @@ impl<'a> TsBatchView<'a> { }) } - /// Read the timestamp at `row_idx` as epoch seconds. + /// Read the timestamp at `row_idx` as a typed `Abs`. #[inline] - pub(crate) fn ts_secs_at(&self, row_idx: usize) -> i64 { + pub(crate) fn ts_at(&self, row_idx: usize) -> Abs { if let Some(arr) = self.ts_col.as_any().downcast_ref::() { - arr.value(row_idx) / 1_000_000_000 + Abs::::from_nanos(arr.value(row_idx)) } else if let Some(arr) = self.ts_col.as_any().downcast_ref::() { - arr.value(row_idx) + Abs::::from_secs(arr.value(row_idx)) } else { panic!( "TsBatchView: TIMESTAMP column has unexpected dtype {:?}", @@ -201,12 +201,6 @@ impl<'a> TsBatchView<'a> { ); } } - - /// Read the timestamp at `row_idx` as a typed `Abs`. - #[inline] - pub(crate) fn ts_at(&self, row_idx: usize) -> Abs { - Abs::::from_secs(self.ts_secs_at(row_idx)) - } } /// BatchView that doesn't read any column — used by FixedDecider / @@ -285,7 +279,7 @@ impl TimeRangeDecider { pub(crate) fn decide_ts(&self, ts: Abs) -> u8 { for (idx, (start, end)) in self.ranges.iter().enumerate() { let in_range = match self.end_shift { - Some(shift) => ts >= *start && apply_offset(ts, &shift) <= *end, + Some(shift) => ts >= *start && apply_offset(ts, &shift).is_some_and(|t| t <= *end), None => ts > *start && ts <= *end, }; if in_range { @@ -719,7 +713,7 @@ impl SplitDecider for RandomDecider { #[inline] fn decide(&self, row: &TTRowRef<'_, U>) -> u8 { - self.decide_entity_ts(row.entity_pk, row.ts.raw()) + self.decide_entity_ts(row.entity_pk, row.ts.to_secs()) } #[inline(always)] @@ -1057,7 +1051,7 @@ impl OffsetUnit { fn difference_at(&self, other: &Self, anchor: Abs) -> Dur { match (self, other) { (OffsetUnit::Fixed(a), OffsetUnit::Fixed(b)) => *a - *b, - _ => apply_offset(anchor, self) - apply_offset(anchor, other), + _ => offset_boundary(anchor, self) - offset_boundary(anchor, other), } } } @@ -1073,14 +1067,27 @@ impl OffsetUnit { /// `+ Months(-1)`) = `2023-12-31` (no clamp needed). This matches /// `pd.DateOffset(months=N)` and the existing `add_calendar_months` /// semantics already shared with the temporal-aggregation window code. -pub(crate) fn apply_offset(anchor: Abs, offset: &OffsetUnit) -> Abs { +pub(crate) fn apply_offset( + anchor: Abs, + offset: &OffsetUnit, +) -> Option> { match offset { - OffsetUnit::Fixed(d) => anchor + *d, + OffsetUnit::Fixed(d) => anchor.checked_add(*d), OffsetUnit::Months(m) => add_calendar_months(anchor, *m), OffsetUnit::Years(y) => add_calendar_months(anchor, y.saturating_mul(12)), } } +/// [`apply_offset`], pinned to the end of the clock it runs past. +fn offset_boundary(anchor: Abs, offset: &OffsetUnit) -> Abs { + let backwards = match offset { + OffsetUnit::Fixed(d) => d.raw() < 0, + OffsetUnit::Months(n) | OffsetUnit::Years(n) => *n < 0, + }; + let clock_end = if backwards { i64::MIN } else { i64::MAX }; + apply_offset(anchor, offset).unwrap_or(Abs::from_raw(clock_end)) +} + /// Subtract a calendar-aware offset from an anchor. /// /// `subtract_calendar_offset(anchor, OffsetUnit::Months(m))` = @@ -1088,9 +1095,9 @@ pub(crate) fn apply_offset(anchor: Abs, offset: &OffsetUnit pub(crate) fn subtract_calendar_offset( anchor: Abs, offset: OffsetUnit, -) -> Abs { +) -> Option> { match offset { - OffsetUnit::Fixed(d) => anchor - d, + OffsetUnit::Fixed(d) => Some(anchor - d), OffsetUnit::Months(m) => add_calendar_months(anchor, -m), OffsetUnit::Years(y) => add_calendar_months(anchor, -y.saturating_mul(12)), } @@ -1246,7 +1253,9 @@ impl DateOffsetDecider { if self.boundaries.len() < 2 { return UNUSED_STAGE; } - let shifted = ts + self.shift; + let Some(shifted) = ts.checked_add(self.shift) else { + return UNUSED_STAGE; + }; let mut out = UNUSED_STAGE; // Match pybind: scan all (start, end) pairs; LAST match wins. for (split_idx, w) in self.boundaries.windows(2).enumerate() { @@ -1316,7 +1325,7 @@ impl SplitDecider for DateOffsetDecider { let mut boundaries: Vec> = Vec::with_capacity(self.offsets.len() + 2); boundaries.push(min_ts); for off in &self.offsets { - boundaries.push(apply_offset(max_ts, off)); + boundaries.push(offset_boundary(max_ts, off)); } boundaries.push(max_ts + tf + self.anchor_start_offset); if self.flexible_boundaries { @@ -1962,7 +1971,7 @@ mod tests { use super::*; use crate::errors::TTGenError; use crate::ttgen_stats::TTRowRef; - use diskgraph_core::time::{Abs, Seconds}; + use diskgraph_core::time::{Abs, Nanoseconds, Seconds}; use diskgraph_core::tt_builder::Value; fn sample_row<'a>( @@ -2122,6 +2131,27 @@ mod tests { assert_eq!(d.decide_ts(Abs::from_secs(40)), UNUSED_STAGE); } + #[test] + fn rows_shifted_past_the_clock_belong_to_no_stage() { + let end = Abs::::from_raw(i64::MAX); + let time_range = TimeRangeDecider::new(vec![(Abs::from_raw(0), end)]) + .with_end_shift(Some(OffsetUnit::Fixed(Dur::from_raw(10)))); + assert_eq!(time_range.decide_ts(Abs::from_raw(i64::MAX - 20)), 0); + assert_eq!( + time_range.decide_ts(Abs::from_raw(i64::MAX - 5)), + UNUSED_STAGE + ); + + let mut date_offset = DateOffsetDecider::new(vec![], Dur::zero(), None); + date_offset.boundaries = vec![Abs::from_raw(0), end]; + date_offset.shift = Dur::from_raw(10); + assert_eq!(date_offset.decide_ts(Abs::from_raw(i64::MAX - 20)), 0); + assert_eq!( + date_offset.decide_ts(Abs::from_raw(i64::MAX - 5)), + UNUSED_STAGE + ); + } + #[test] fn test_time_range_decider_static_shared_boundary_belongs_to_earlier_stage() { let d = TimeRangeDecider::::new(vec![ @@ -2988,7 +3018,7 @@ mod tests { // = 2024-02-29 (leap year, last valid day of Feb). This matches // `pd.DateOffset(months=-1)` and `add_calendar_months`'s behavior. let anchor = ymd_secs(2024, 3, 31); - let result = subtract_calendar_offset::(anchor, OffsetUnit::Months(1)); + let result = subtract_calendar_offset::(anchor, OffsetUnit::Months(1)).unwrap(); assert_eq!( ymd_to_str(result), "2024-02-29", @@ -2998,7 +3028,7 @@ mod tests { // Non-leap year: 2023-03-31 - 1 month = 2023-02-28. let anchor_nonleap = ymd_secs(2023, 3, 31); let result_nonleap = - subtract_calendar_offset::(anchor_nonleap, OffsetUnit::Months(1)); + subtract_calendar_offset::(anchor_nonleap, OffsetUnit::Months(1)).unwrap(); assert_eq!( ymd_to_str(result_nonleap), "2023-02-28", @@ -3011,7 +3041,7 @@ mod tests { // Sanity: OffsetUnit::Fixed(d) is identical to anchor + d. let anchor = ymd_secs(2024, 6, 15); let d = Dur::::from_secs(3600); - let via_offset = apply_offset(anchor, &OffsetUnit::Fixed(d)); + let via_offset = apply_offset(anchor, &OffsetUnit::Fixed(d)).unwrap(); let via_dur = anchor + d; assert_eq!(via_offset.raw(), via_dur.raw()); } diff --git a/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs index 5a4467ddd..d35f63ffa 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs @@ -22,7 +22,7 @@ use diskgraph_core::accumulators::{ for_each_split_ascii_byte, HyperLogLog, LocalColumnCounts, SharedColumnCounts, TDigestEstimator, WelfordAccumulator, }; -use diskgraph_core::time::{add_calendar_months, dur_from_unit, Abs, Dur, TimeUnit, UnitSystem}; +use diskgraph_core::time::{dur_from_unit, Abs, Dur, TimeUnit, UnitSystem}; use diskgraph_core::tt_builder::*; use diskgraph_index::graph::GraphIndexReader; use parquet::arrow::ArrowWriter; @@ -1860,11 +1860,11 @@ fn execute_entity>( let time_col_val: Option> = plan .entity_time_col_idx .and_then(|idx| read_column_i64(&plan.entity_table, entity_node, idx)) - .map(Abs::::from_secs); + .map(Abs::::from_raw); let valid_to_val: Option> = plan .entity_valid_to_col_idx .and_then(|idx| read_column_i64(&plan.entity_table, entity_node, idx)) - .map(Abs::::from_secs); + .map(Abs::::from_raw); // 2. Pre-resolve entity column values (for filter RHS EntityColumn) let entity_cols: Vec> = plan @@ -1906,6 +1906,9 @@ fn execute_entity>( (_, Some(t)) => t, (None, None) => default_ts, }; + if !plan.fits_clock(Abs::::from_raw(0)) { + return Ok(()); + } link_frequencies.clear(); link_split_values.clear(); if execute_stages( @@ -1941,7 +1944,7 @@ fn execute_entity>( // 4. Temporal: binary search for valid timestamp range let timestamps = &plan.timestamps; - // Timestamps are sorted ascending. Point anchors (epoch seconds). + // Timestamps are sorted ascending. Point anchors (index-clock ticks). // Entity excluded if time_col > ts (strict >) // Entity excluded if valid_to_col < ts (strict <, i.e. LEQ boundary) let first_valid = match time_col_val { @@ -2546,39 +2549,30 @@ fn execute_target_sequence( if !can_optimize { for &offset in offsets { - let shifted_ts = base_ts + offset; - let val_opt = execute_aggregate( - plan, - entity_node, - agg, - entity_cols, - Some(shifted_ts), - &mut Vec::new(), - &mut Vec::new(), - ); + let val_opt = base_ts.checked_add(offset).and_then(|shifted_ts| { + execute_aggregate( + plan, + entity_node, + agg, + entity_cols, + Some(shifted_ts), + &mut Vec::new(), + &mut Vec::new(), + ) + }); results.push(val_opt); } return results; } - let ResolvedWindow::Fixed { - offset_lo, - offset_hi, - encoder, - } = agg.window.as_ref().unwrap() + let window = agg.window.as_ref().unwrap(); + let Some(windows) = offsets + .iter() + .map(|&offset| window.sort_keys(base_ts.checked_add(offset)?)) + .collect::>>() else { - unreachable!() + return vec![None; offsets.len()]; }; - let windows: Vec = offsets - .iter() - .map(|&offset| { - let win_ts = base_ts + offset; - let abs_lo = offset_lo - .map(|offset_lo| encoder.encode_upper((win_ts + offset_lo).raw())); - let abs_hi = encoder.encode_upper((win_ts + *offset_hi).raw()); - (abs_lo, abs_hi) - }) - .collect(); let min_lo = windows.iter().map(|w| w.0).min().unwrap(); let max_hi = windows.iter().map(|w| w.1).max().unwrap(); let edge_ref = &agg.path[0]; @@ -2702,7 +2696,7 @@ fn execute_target_sequence( /// Execute all stages for an (entity, timestamp) pair. /// Returns true if all stages passed (output row should be emitted). /// -/// `observation_ts` is the current observation timestamp (epoch seconds) for +/// `observation_ts` is the current observation timestamp (index-clock ticks) for /// temporal plans. Used to compute absolute window bounds from relative offsets. /// `None` for non-temporal plans. /// @@ -2904,37 +2898,11 @@ fn execute_has_edges( entity_cols: &[Option], observation_ts: Option>, ) -> Option { - // SAFETY: encode_upper expects epoch-seconds (the index's internal unit). - // This is correct when U=Seconds. If the pipeline is ever monomorphized - // over a different unit (e.g. Nanoseconds), the encoder must be made - // unit-aware as part of the ingest/sampling unit awareness work. - let window = has_edges.window.as_ref().map(|w| { - let obs = observation_ts.unwrap_or(Abs::::from_raw(0)); - match w { - ResolvedWindow::Fixed { - offset_lo, - offset_hi, - encoder, - } => { - let lo = offset_lo - .map(|offset_lo| encoder.encode_upper((obs + offset_lo).raw())); - let hi = encoder.encode_upper((obs + *offset_hi).raw()); - (lo, hi) - } - ResolvedWindow::Calendar { - lo_months, - hi_months, - encoder, - } => { - let hi_abs = add_calendar_months(obs, *hi_months); - let lo = lo_months.map(|lo_months| { - encoder.encode_upper(add_calendar_months(obs, lo_months).raw()) - }); - let hi = encoder.encode_upper(hi_abs.raw()); - (lo, hi) - } - } - }); + let obs = observation_ts.unwrap_or(Abs::::from_raw(0)); + let window = match &has_edges.window { + Some(window) => Some(window.sort_keys(obs)?), + None => None, + }; let dst_table = &has_edges.path.last().unwrap().dst_table; let mut traversed = false; @@ -2975,33 +2943,11 @@ fn execute_aggregate( // timestamp bucket. Combined with the (lo, hi] search in iter_edge, this // ensures edges exactly at the lower boundary are excluded and edges exactly // at the upper boundary are included. - let window = agg.window.as_ref().map(|w| { - let obs = observation_ts.unwrap_or(Abs::::from_raw(0)); - match w { - ResolvedWindow::Fixed { - offset_lo, - offset_hi, - encoder, - } => { - let lo = offset_lo - .map(|offset_lo| encoder.encode_upper((obs + offset_lo).raw())); - let hi = encoder.encode_upper((obs + *offset_hi).raw()); - (lo, hi) - } - ResolvedWindow::Calendar { - lo_months, - hi_months, - encoder, - } => { - let hi_abs = add_calendar_months(obs, *hi_months); - let lo = lo_months.map(|lo_months| { - encoder.encode_upper(add_calendar_months(obs, lo_months).raw()) - }); - let hi = encoder.encode_upper(hi_abs.raw()); - (lo, hi) - } - } - }); + let obs = observation_ts.unwrap_or(Abs::::from_raw(0)); + let window = match &agg.window { + Some(window) => Some(window.sort_keys(obs)?), + None => None, + }; // Get the destination table for column reads let dst_table = &agg.path.last().unwrap().dst_table; @@ -3137,13 +3083,11 @@ fn execute_aggregate( if agg.target_temporal_create_col_idx.is_some() || agg.target_temporal_end_col_idx.is_some() => { - let create_ref = match agg.window.as_ref() { - Some(ResolvedWindow::Fixed { offset_hi, .. }) => ts + *offset_hi, - Some(ResolvedWindow::Calendar { hi_months, .. }) => { - add_calendar_months(ts, *hi_months) - } - None => ts, - }; + let create_ref = agg + .window + .as_ref() + .and_then(|w| w.bounds(ts)) + .map_or(ts, |(_, end)| end); Some((create_ref.raw().max(ts.raw()), ts.raw())) } _ => None, @@ -3324,6 +3268,7 @@ fn execute_aggregate( #[cfg(test)] mod tests { use super::*; + use diskgraph_core::time::add_calendar_months; #[test] fn pass_a_writer_rotates_into_ordered_parts() { @@ -3637,7 +3582,7 @@ mod tests { // 2024-01-01 00:00:00 UTC let jan1 = Abs::::from_raw(1704067200); // +3 months = 2024-04-01 (leap year, so Jan→Apr = 91 days) - let apr1 = add_calendar_months(jan1, 3); + let apr1 = add_calendar_months(jan1, 3).unwrap(); assert_eq!(apr1.raw(), 1704067200 + 91 * 86400); } @@ -3646,7 +3591,7 @@ mod tests { // 2023-01-01 00:00:00 UTC let jan1 = Abs::::from_raw(1672531200); // +3 months = 2023-04-01 (non-leap, Jan→Apr = 90 days) - let apr1 = add_calendar_months(jan1, 3); + let apr1 = add_calendar_months(jan1, 3).unwrap(); assert_eq!(apr1.raw(), 1672531200 + 90 * 86400); } @@ -3655,7 +3600,7 @@ mod tests { // 2023-04-01 00:00:00 UTC let apr1 = Abs::::from_raw(1680307200); // +3 months = 2023-07-01 (Apr→Jul = 91 days) - let jul1 = add_calendar_months(apr1, 3); + let jul1 = add_calendar_months(apr1, 3).unwrap(); assert_eq!(jul1.raw(), apr1.raw() + 91 * 86400); } @@ -3664,7 +3609,7 @@ mod tests { // 2023-07-01 00:00:00 UTC let jul1 = Abs::::from_raw(1688169600); // +3 months = 2023-10-01 (Jul→Oct = 92 days) - let oct1 = add_calendar_months(jul1, 3); + let oct1 = add_calendar_months(jul1, 3).unwrap(); assert_eq!(oct1.raw(), jul1.raw() + 92 * 86400); } @@ -3673,7 +3618,7 @@ mod tests { // 2023-10-01 00:00:00 UTC let oct1 = Abs::::from_raw(1696118400); // +3 months = 2024-01-01 (Oct→Jan = 92 days) - let jan1 = add_calendar_months(oct1, 3); + let jan1 = add_calendar_months(oct1, 3).unwrap(); assert_eq!(jan1.raw(), oct1.raw() + 92 * 86400); } @@ -3681,7 +3626,7 @@ mod tests { fn test_add_calendar_months_negative() { // 2023-04-01 → -3 months = 2023-01-01 let apr1 = Abs::::from_raw(1680307200); - let jan1 = add_calendar_months(apr1, -3); + let jan1 = add_calendar_months(apr1, -3).unwrap(); assert_eq!(jan1.raw(), 1672531200); } @@ -3689,7 +3634,7 @@ mod tests { fn test_add_calendar_months_years() { // 2023-01-01 + 12 months = 2024-01-01 let y2023 = Abs::::from_raw(1672531200); - let y2024 = add_calendar_months(y2023, 12); + let y2024 = add_calendar_months(y2023, 12).unwrap(); assert_eq!(y2024.raw(), y2023.raw() + 365 * 86400); } } diff --git a/packages/diskgraph/diskgraph-train-table/src/ttgen_output.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_output.rs index 81f6f8835..662561fb3 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_output.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_output.rs @@ -780,10 +780,10 @@ impl ChannelBatchWriter { columns.push(flush_entity_buf(&mut self.entity_bufs[slot], dtype)); } - // Convert epoch seconds → nanoseconds to match Kumo's datetime64[ns] convention. + // Convert index ticks → nanoseconds to match Kumo's datetime64[ns] convention. let ts_nanos: Vec = std::mem::take(&mut self.timestamp_buf) .into_iter() - .map(|s| s.saturating_mul(1_000_000_000)) + .map(|ticks| Abs::::from_raw(ticks).to_nanos()) .collect(); let ts_arr = TimestampNanosecondArray::from(ts_nanos); columns.push(Arc::new(ts_arr)); diff --git a/packages/diskgraph/diskgraph-train-table/src/ttgen_resolve.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_resolve.rs index 1cfd93064..c4955d945 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_resolve.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_resolve.rs @@ -11,9 +11,10 @@ use std::sync::Arc; use crate::errors::{TTGenError, TTGenResult}; use crate::ttgen_output::PRIMARY_TARGET_COL; +use crate::ttgen_traversal::SortKeyWindow; use arrow_schema::{DataType, Field}; use diskgraph_core::sortkey::SortKeyEncoder; -use diskgraph_core::time::{Abs, Dur, UnitSystem}; +use diskgraph_core::time::{add_calendar_months, Abs, Dur, UnitSystem}; use diskgraph_core::tt_builder::*; use diskgraph_index::graph::{GraphIndexReader, TableReader}; use diskgraph_index::layout::RelationLayout; @@ -188,6 +189,29 @@ impl ResolvedOp { ResolvedOp::FutureLabels(a) => a.alias_idx, } } + + /// Whether every window and offset this op evaluates at `obs` stays inside the clock. + pub fn fits_clock(&self, obs: Abs) -> bool { + let shifted_fit = |target: &ResolvedOp, offsets: &[Dur]| { + offsets.iter().all(|&offset| match obs.checked_add(offset) { + Some(ts) => target.fits_clock(ts), + None => false, + }) + }; + match self { + ResolvedOp::Aggregate(a) => window_fits(&a.window, obs), + ResolvedOp::HasEdges(h) => window_fits(&h.window, obs), + ResolvedOp::Autoregression(a) => { + a.target_op.fits_clock(obs) && shifted_fit(&a.target_op, &a.lag_offsets) + } + ResolvedOp::FutureLabels(f) => shifted_fit(&f.target_op, &f.future_offsets), + ResolvedOp::Column(_) | ResolvedOp::Compare(_) | ResolvedOp::LogicalOp(_) => true, + } + } +} + +fn window_fits(window: &Option>, obs: Abs) -> bool { + window.as_ref().is_none_or(|w| w.bounds(obs).is_some()) } pub struct ResolvedAggregate { @@ -274,6 +298,47 @@ pub enum ResolvedWindow { }, } +impl ResolvedWindow { + /// The window's exclusive start and inclusive end around `obs`, or `None` + /// when either falls outside the clock. + pub fn bounds(&self, obs: Abs) -> Option<(Option>, Abs)> { + match self { + ResolvedWindow::Fixed { + offset_lo, + offset_hi, + .. + } => { + let lo = match offset_lo { + Some(offset) => Some(obs.checked_add(*offset)?), + None => None, + }; + Some((lo, obs.checked_add(*offset_hi)?)) + } + ResolvedWindow::Calendar { + lo_months, + hi_months, + .. + } => { + let lo = match lo_months { + Some(months) => Some(add_calendar_months(obs, *months)?), + None => None, + }; + Some((lo, add_calendar_months(obs, *hi_months)?)) + } + } + } + + pub fn sort_keys(&self, obs: Abs) -> Option { + let (ResolvedWindow::Fixed { encoder, .. } | ResolvedWindow::Calendar { encoder, .. }) = + self; + let (lo, hi) = self.bounds(obs)?; + Some(( + lo.map(|lo| encoder.encode_upper(lo.raw())), + encoder.encode_upper(hi.raw()), + )) + } +} + // ─── Resolved filter types ─────────────────────────────────────────────────── pub enum ResolvedRowFilter { @@ -598,7 +663,7 @@ pub fn resolve_plan( .map(|t| t.temporal_col_indices()) .unwrap_or((None, None)); - Ok(ResolvedPlan { + let mut resolved = ResolvedPlan { entity_table_spec: plan.entity.clone(), entity_table, entity_time_col_idx, @@ -640,10 +705,24 @@ pub fn resolve_plan( custom_target_pk_table, custom_target_temporal_create_col_idx, custom_target_temporal_end_col_idx, - }) + }; + resolved.timestamps = resolved + .timestamps + .iter() + .copied() + .filter(|&ts| resolved.fits_clock(ts)) + .collect(); + Ok(resolved) } impl ResolvedPlan { + /// Whether the row anchored at `obs` can be evaluated without leaving the clock. + pub fn fits_clock(&self, obs: Abs) -> bool { + self.stages + .iter() + .all(|stage| stage.ops.iter().all(|op| op.fits_clock(obs))) + } + /// Convenience accessor: the primary (slot 0) ENTITY rowpack column index. /// Used by all single-entity call sites (PQ path, custom-TT lookup). #[inline] diff --git a/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs index d5eda42b0..b21427a79 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs @@ -53,7 +53,7 @@ use diskgraph_core::accumulators::{ build_multicat_emit_stats, kumo_entry, kumo_key, kumo_safe_div, HyperLogLog, LocalColumnCounts, TDigestEstimator, WelfordAccumulator, }; -use diskgraph_core::time::{Abs, UnitSystem}; +use diskgraph_core::time::{split_parts, Abs, UnitSystem}; use diskgraph_core::tt_builder::Value; use crate::split_decider::SplitDecider; @@ -187,7 +187,7 @@ pub(crate) struct ChunkStats { pub(crate) entity_hll: Option, pub(crate) entity_welford: Option, pub(crate) entity_tdigest: Option, - // Distribution stats for TIMESTAMP column (numerical, epoch seconds) + // Distribution stats for TIMESTAMP column (numerical, index-clock ticks) pub(crate) timestamp_welford: Option, pub(crate) timestamp_tdigest: Option, pub(crate) timestamp_hll: Option, @@ -265,7 +265,7 @@ pub(crate) struct StatsMerger { entity_hll: Option, entity_welford: Option, entity_tdigest: Option, - // Distribution stats for TIMESTAMP column (numerical, epoch seconds) + // Distribution stats for TIMESTAMP column (numerical, index-clock ticks) timestamp_welford: Option, timestamp_tdigest: Option, timestamp_hll: Option, @@ -512,11 +512,11 @@ impl StatsMerger { 0 }; - // Format epoch-second bounds as RFC3339 strings for Kumo's cache. + // Format timestamp bounds as RFC3339 strings for Kumo's cache. use chrono::DateTime; let fmt_ts = |t: Abs| -> JsonValue { - let secs = t.to_secs(); - DateTime::from_timestamp(secs, 0) + let (secs, nanos) = split_parts(t.raw(), U::UNIT); + DateTime::from_timestamp(secs, nanos) .map(|dt| json!(dt.to_rfc3339())) .unwrap_or_else(|| json!(secs)) }; @@ -1123,7 +1123,7 @@ pub fn build_inline_col_stats( } } - // TIMESTAMP numerical stats (epoch seconds — full stats including quantiles) + // TIMESTAMP numerical stats (index-clock ticks — full stats including quantiles) if let Some(w) = timestamp_welford { let count = w.count; let stype = "kumo.numerical"; @@ -1572,7 +1572,7 @@ fn merge_chunk_stats_into(dst: &mut ChunkStats, src: ChunkStat #[cfg(test)] mod tests { use super::*; - use diskgraph_core::time::Seconds; + use diskgraph_core::time::{Milliseconds, Seconds}; #[test] fn test_finalize_contains_exact_stats_payload() { @@ -1983,14 +1983,14 @@ mod tests { /// Helper: drive merger->finalize with the minimum knobs the kumo /// contract needs. All other args take their no-op defaults. #[allow(clippy::too_many_arguments)] - fn finalize_minimal( - chunk: ChunkStats, + fn finalize_minimal( + chunk: ChunkStats, num_total_rows: u64, is_link_prediction: bool, target_category_counts: Option>, target_stype: &str, ) -> TrainTableStatsPayload { - let mut merger: StatsMerger = StatsMerger::default(); + let mut merger: StatsMerger = StatsMerger::default(); let train_rows = chunk.num_train_rows; merger.merge(chunk); merger.finalize( @@ -2289,6 +2289,20 @@ mod tests { ); } + #[test] + fn top_level_dates_keep_the_fraction_of_a_second() { + let mut chunk: ChunkStats = ChunkStats { + num_train_rows: 2, + ..Default::default() + }; + chunk.timestamp_min = Some(Abs::from_raw(1_700_000_000_000)); + chunk.timestamp_max = Some(Abs::from_raw(1_700_000_000_250)); + + let kc = finalize_minimal(chunk, 2, false, None, "numerical").kumo_contract; + assert_eq!(kc.min_date, "2023-11-14T22:13:20+00:00"); + assert_eq!(kc.max_date, "2023-11-14T22:13:20.250+00:00"); + } + #[test] fn test_finalize_top_level_dates_null_when_no_data() { // No timestamps observed → date fields are null. diff --git a/packages/diskgraph/diskgraph-train-table/src/ttgen_traversal.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_traversal.rs index e20772981..deb1a8216 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_traversal.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_traversal.rs @@ -7,6 +7,7 @@ //! `iter_edge` reads edges from a relation segment. //! `traverse_and_accumulate` recursively walks multi-hop paths. +use diskgraph_core::time::{convert_trunc, TimeUnit}; use diskgraph_core::tt_builder::Value; use diskgraph_index::graph::TableReader; use diskgraph_index::layout::RelationLayout; @@ -206,7 +207,7 @@ pub fn read_column_value( match value_ref { ValueRef::Null => None, - _ => Some(value_ref_to_value(&value_ref)), + _ => Some(value_ref_to_value(&value_ref, decoder.clock())), } } @@ -232,7 +233,7 @@ pub fn read_column_value_nullable( match value_ref { ValueRef::Null => Some(None), - _ => Some(Some(value_ref_to_value(&value_ref))), + _ => Some(Some(value_ref_to_value(&value_ref, decoder.clock()))), } } @@ -385,7 +386,8 @@ where if matches!(value_ref, ValueRef::Null) { return SourceNull; } - return if observe(MulticatToken::Legacy(value_ref_to_value(&value_ref))) { + let value = value_ref_to_value(&value_ref, decoder.clock()); + return if observe(MulticatToken::Legacy(value)) { Continue } else { Stop @@ -477,8 +479,8 @@ where } } -/// Convert a rowpack ValueRef to our Value type. -fn value_ref_to_value(vr: &ValueRef) -> Value { +/// Convert a rowpack ValueRef to our Value type; dates become ticks of `clock`. +fn value_ref_to_value(vr: &ValueRef, clock: TimeUnit) -> Value { match vr { ValueRef::Null => Value::Int(0), ValueRef::Int64(v) => Value::Int(*v), @@ -490,8 +492,8 @@ fn value_ref_to_value(vr: &ValueRef) -> Value { ValueRef::Bool(b) => Value::Bool(*b), ValueRef::String(s) => Value::Str(s.to_string()), ValueRef::Bytes(b) => Value::Str(String::from_utf8_lossy(b).to_string()), - ValueRef::Date(d) => Value::Int(*d as i64 * 86400), - ValueRef::Timestamp(secs) => Value::Int(*secs), + ValueRef::Date(d) => Value::Int(convert_trunc(i64::from(*d), TimeUnit::Days, clock)), + ValueRef::Timestamp(ticks) => Value::Int(*ticks), ValueRef::Decimal128(v) => Value::Float(*v as f64), _ => Value::Int(0), // Uuid, Ip, arrays, etc. → default } diff --git a/packages/diskgraph/diskgraph-train-table/tests/custom_tt_integration.rs b/packages/diskgraph/diskgraph-train-table/tests/custom_tt_integration.rs index b7a5190a6..d2386f0ce 100644 --- a/packages/diskgraph/diskgraph-train-table/tests/custom_tt_integration.rs +++ b/packages/diskgraph/diskgraph-train-table/tests/custom_tt_integration.rs @@ -290,6 +290,7 @@ fn build_users_only_graph(root: &Path, n_users: i64) -> PathBuf { write_parquet(&data_dir.join("users.parquet"), &users_batch); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_custom_tt".to_string(), tables: vec![TableSpec { name: "users".to_string(), @@ -373,6 +374,7 @@ fn build_users_and_products_graph(root: &Path, n_users: i64, n_products: i64) -> }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_custom_tt_lp".to_string(), tables: vec![ mk_table("users", "user_id", "users.parquet"), diff --git a/packages/diskgraph/diskgraph-train-table/tests/ttgen_e2e.rs b/packages/diskgraph/diskgraph-train-table/tests/ttgen_e2e.rs index 666826ee7..768d95841 100644 --- a/packages/diskgraph/diskgraph-train-table/tests/ttgen_e2e.rs +++ b/packages/diskgraph/diskgraph-train-table/tests/ttgen_e2e.rs @@ -178,6 +178,7 @@ fn build_e2e_graph(root: &Path) -> PathBuf { unit: "s".to_string(), }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_e2e".to_string(), tables: vec![ TableSpec { @@ -304,6 +305,7 @@ fn build_users_only_graph(root: &Path, n_users: i64) -> PathBuf { write_parquet(&data_dir.join("users.parquet"), &users_batch); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_e2e_users_only".to_string(), tables: vec![TableSpec { name: "users".to_string(), diff --git a/packages/diskgraph/diskgraph-train-table/tests/ttgen_integration.rs b/packages/diskgraph/diskgraph-train-table/tests/ttgen_integration.rs index 1de7d9526..aeb6e8ae0 100644 --- a/packages/diskgraph/diskgraph-train-table/tests/ttgen_integration.rs +++ b/packages/diskgraph/diskgraph-train-table/tests/ttgen_integration.rs @@ -683,6 +683,7 @@ fn build_graph1(root: &Path) -> PathBuf { // ── GraphConfig ── let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_graph1".to_string(), tables: vec![ TableSpec { @@ -906,6 +907,7 @@ fn build_related_column_graph(root: &Path) -> PathBuf { write_parquet(&data_dir.join("targets.parquet"), &targets); let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "related_column_graph".to_string(), tables: vec![ TableSpec { @@ -2624,6 +2626,7 @@ fn build_graph1_with_create_time(root: &Path) -> PathBuf { // ── GraphConfig with create_time on users ── let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_graph1_create_time".to_string(), tables: vec![ TableSpec { @@ -3438,6 +3441,7 @@ fn build_graph_calendar(root: &Path) -> PathBuf { // ── GraphConfig ── let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_calendar".to_string(), tables: vec![ TableSpec { @@ -4804,6 +4808,7 @@ fn build_graph2(root: &Path) -> PathBuf { // ── GraphConfig ── let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_graph2".to_string(), tables: vec![ TableSpec { @@ -5326,6 +5331,7 @@ fn build_graph3(root: &Path) -> PathBuf { // ── GraphConfig ── let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_graph3".to_string(), tables: vec![ TableSpec { @@ -6724,6 +6730,7 @@ fn build_graph_orphan_fk(root: &Path) -> PathBuf { }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_orphan_fk".to_string(), tables: vec![ TableSpec { @@ -7125,6 +7132,7 @@ fn build_graph_orphan_fk_all_valid(root: &Path) -> PathBuf { }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_orphan_fk_all_valid".to_string(), tables: vec![ TableSpec { @@ -7308,6 +7316,7 @@ fn build_graph_orphan_fk_all_orphans(root: &Path) -> PathBuf { }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_orphan_fk_all_orphans".to_string(), tables: vec![ TableSpec { @@ -7610,6 +7619,7 @@ fn build_graph_temporal_orphan(root: &Path) -> PathBuf { }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_temporal_orphan".to_string(), tables: vec![ TableSpec { @@ -9773,6 +9783,7 @@ fn build_multicat_graph( }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "test_multicat".to_string(), tables: vec![ TableSpec { @@ -10142,6 +10153,7 @@ fn build_scalar_status_graph(root: &Path) -> PathBuf { multicategorical_delimiter: None, }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "scalar_contains".to_string(), tables: vec![ TableSpec { @@ -10345,6 +10357,7 @@ fn build_multicat_with_empty_list(root: &Path) -> PathBuf { multicategorical_delimiter: None, }; let config = GraphConfig { + timestamp_unit: Default::default(), dataset_name: "multicat_empty".to_string(), tables: vec![ TableSpec { diff --git a/packages/diskgraph/third_party/artifact-manifest.json b/packages/diskgraph/third_party/artifact-manifest.json index 1f62d43ef..edb6159c8 100644 --- a/packages/diskgraph/third_party/artifact-manifest.json +++ b/packages/diskgraph/third_party/artifact-manifest.json @@ -5202,6 +5202,6 @@ ] } ], - "cargo_lock_sha256": "297851eb08303a13a97f091c68cf3ce5ccac633ad3d7fe09724c1506e486cdaf", + "cargo_lock_sha256": "57a8bc550cb096e01ae14f139d3a2b8b2a12641b625b31044626e20a8d384e96", "schema": 1 } diff --git a/packages/kumo-api/kumoapi/typing.py b/packages/kumo-api/kumoapi/typing.py index a61842a63..eb5dc6f9a 100644 --- a/packages/kumo-api/kumoapi/typing.py +++ b/packages/kumo-api/kumoapi/typing.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import builtins -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, get_args import pydantic from pandas import DateOffset @@ -31,6 +31,16 @@ def _forbid_extra_config(**options: Any) -> Any: SUPPORTED_MULTICATEGORICAL_DELIMITERS = frozenset({',', '|', ';', '\t'}) MulticategoricalDelimiter = Literal[',', '|', ';', '\t'] +TimestampPrecision = Literal['nanoseconds', 'microseconds', 'milliseconds', + 'seconds', 'hours', 'days'] + + +def validate_timestamp_precision(timestamp_precision: str) -> None: + if timestamp_precision not in get_args(TimestampPrecision): + raise ValueError(f"timestamp_precision must be one of " + f"{get_args(TimestampPrecision)}, got " + f"{timestamp_precision!r}") + class Stype(StrEnum): r"""The semantic type of a column. diff --git a/packages/kumo-core/src/kumo_core/diskgraph/ingest.py b/packages/kumo-core/src/kumo_core/diskgraph/ingest.py index 97d0d72fe..df218bbbf 100644 --- a/packages/kumo-core/src/kumo_core/diskgraph/ingest.py +++ b/packages/kumo-core/src/kumo_core/diskgraph/ingest.py @@ -12,6 +12,8 @@ from pathlib import Path from typing import Any +from kumoapi.typing import TimestampPrecision + from ..types.schema import TS_FORMAT_METADATA from ._progress import log_progress from .runner import ( @@ -223,6 +225,12 @@ def validate_diskgraph_ingest_config(ingest_config: dict[str, Any]) -> None: "delimiter, or use a native stringlist/intlist dtype.") +def set_timestamp_precision(ingest_config: dict[str, Any], + timestamp_precision: str) -> None: + if timestamp_precision != "seconds": + ingest_config["timestamp_unit"] = timestamp_precision + + @dataclass(frozen=True) class DiskgraphIngestSpec: """Fully resolved inputs for one DiskGraph snapshot ingest.""" @@ -233,6 +241,7 @@ class DiskgraphIngestSpec: source_paths: Mapping[str, str] execution: IngestExecutionConfig = field( default_factory=IngestExecutionConfig) + timestamp_precision: TimestampPrecision = "seconds" IngestRunner = Callable[ @@ -262,6 +271,7 @@ def execute_diskgraph_ingest( dict(spec.table_data_versions), source_paths=dict(spec.source_paths), ).convert() + set_timestamp_precision(ingest_config, spec.timestamp_precision) validate_diskgraph_ingest_config(ingest_config) run_result = runner(ingest_config, output_path, spec.execution) diff --git a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py index 3c8e30cec..e07e7f258 100644 --- a/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py +++ b/packages/kumo-core/src/kumo_core/diskgraph/lightning_data.py @@ -1105,6 +1105,8 @@ def _num_nodes_for_type( 'nanoseconds': 'ns', 'nanos': 'ns', 'ns': 'ns', + 'hours': 'h', + 'days': 'D', } # numpy datetime64 unit -> nanoseconds-per-one-tick. Used to convert Kumo's @@ -1116,6 +1118,8 @@ def _num_nodes_for_type( 'ms': 1_000_000, 'us': 1_000, 'ns': 1, + 'h': 3_600_000_000_000, + 'D': 86_400_000_000_000, } diff --git a/packages/kumo-core/src/kumo_core/diskgraph/runtime_adapter.py b/packages/kumo-core/src/kumo_core/diskgraph/runtime_adapter.py index ebaffc906..bff62b138 100644 --- a/packages/kumo-core/src/kumo_core/diskgraph/runtime_adapter.py +++ b/packages/kumo-core/src/kumo_core/diskgraph/runtime_adapter.py @@ -64,6 +64,8 @@ "ms": 1_000_000, "us": 1_000, "ns": 1, + "h": 3_600_000_000_000, + "D": 86_400_000_000_000, } # --- FFI-boundary shapes ----------------------------------------------- diff --git a/packages/kumo-core/src/kumo_core/pquery.py b/packages/kumo-core/src/kumo_core/pquery.py index 6954e2e08..ef13410ae 100644 --- a/packages/kumo-core/src/kumo_core/pquery.py +++ b/packages/kumo-core/src/kumo_core/pquery.py @@ -14,6 +14,7 @@ import hashlib import json import os +import re from dataclasses import dataclass from dataclasses import field as dataclass_field from dataclasses import replace @@ -892,6 +893,9 @@ def _split_options( return split, family +_SECONDS_FRACTION = re.compile(r"[.,](\d+)") + + def _timestamp_ns(value: Any) -> int: if isinstance(value, bool): raise PredictiveQueryPlanError("boolean timestamp statistic is invalid") @@ -913,7 +917,11 @@ def _timestamp_ns(value: Any) -> int: f"invalid timestamp statistic {value!r}") from error if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) - return int(parsed.timestamp() * 1_000_000_000) + nanoseconds = Timestamp(parsed).value + fraction = _SECONDS_FRACTION.search(rendered) + if fraction and len(fraction.group(1)) > 6: # datetime keeps only 6 digits + nanoseconds += int(fraction.group(1)[6:9].ljust(3, "0")) + return nanoseconds def _stat_value(stats: Mapping[str, Any], column: str, kind: str) -> Any: diff --git a/packages/kumo-core/src/kumo_core/prediction.py b/packages/kumo-core/src/kumo_core/prediction.py index effb6b69e..ac89557c5 100644 --- a/packages/kumo-core/src/kumo_core/prediction.py +++ b/packages/kumo-core/src/kumo_core/prediction.py @@ -166,7 +166,7 @@ def generate_prediction_table( minimum, table_name) if error is not None: raise ValueError(f"{error.title}. {error.message}") - compiler = runtime.DiskGraphPlanCompiler() + compiler = runtime.DiskGraphPlanCompiler(index_path=request.index_path) # Future targets belong to training only. Prediction retains the fitted # historical feature layout without asking the compiler for labels. forecast_options.pop("future_offsets_secs", None) diff --git a/packages/kumo-core/src/kumo_core/train/shallow_rhs_link_gnn_module.py b/packages/kumo-core/src/kumo_core/train/shallow_rhs_link_gnn_module.py index f0fc726e5..857a8684f 100644 --- a/packages/kumo-core/src/kumo_core/train/shallow_rhs_link_gnn_module.py +++ b/packages/kumo-core/src/kumo_core/train/shallow_rhs_link_gnn_module.py @@ -250,7 +250,7 @@ def _get_valid_rhs_idx_for_batch(self, batch: HeteroData) -> Tensor | None: seed_time = batch[self.lhs_name].seed_time return self._get_valid_rhs_idx_for_seed_time( - seed_time=int(seed_time[0]), # Only contains shared seed times. + seed_time=int(seed_time.min()), # Anchors share a second. device=seed_time.device, ) diff --git a/packages/kumo-core/src/kumo_core/training/table.py b/packages/kumo-core/src/kumo_core/training/table.py index 558e5b273..65efdb95d 100644 --- a/packages/kumo-core/src/kumo_core/training/table.py +++ b/packages/kumo-core/src/kumo_core/training/table.py @@ -485,7 +485,7 @@ def generate_train_table( ) if split is None: raise RuntimeError("split_config has no resolvable split family") - compiler = runtime.DiskGraphPlanCompiler() + compiler = runtime.DiskGraphPlanCompiler(index_path=request.index_path) compile_options = dict(request.compiler_options or {}) plan_json = compiler.compile( query=request.query, diff --git a/packages/kumo-core/tests/test_diskgraph_runtime.py b/packages/kumo-core/tests/test_diskgraph_runtime.py index 457e43fc2..17162a9dc 100644 --- a/packages/kumo-core/tests/test_diskgraph_runtime.py +++ b/packages/kumo-core/tests/test_diskgraph_runtime.py @@ -11,6 +11,7 @@ import tempfile from pathlib import Path from types import SimpleNamespace +from typing import get_args import kumo_core.diskgraph.runner as diskgraph_runner import pyarrow as pa @@ -19,6 +20,7 @@ from kumo_core.diskgraph.ingest import ( DiskgraphIngestSpec, execute_diskgraph_ingest, + set_timestamp_precision, ) from kumo_core.diskgraph.snapshot import ( DiskgraphIngestResult, @@ -27,6 +29,7 @@ DiskgraphSnapshotHandle, LocalSnapshotStore, ) +from kumoapi.typing import TimestampPrecision ZERO_COUNTS = ("diskgraph ingest completed: 0 node tables, 0 edge relations, " "0 nodes, 0 edges") @@ -788,3 +791,11 @@ def test_list_ready_skips_invalid_snapshot_without_hiding_valid( with pytest.raises(DiskgraphSnapshotError, match="no usable rows"): store.verify(DiskgraphSnapshotHandle("graph", "invalid")) assert Path(store.path(DiskgraphSnapshotHandle("graph", "invalid"))).exists() + + +@pytest.mark.parametrize("clock", get_args(TimestampPrecision)) +def test_ingest_config_names_the_clock_unless_seconds(clock: str) -> None: + ingest_config: dict = {} + set_timestamp_precision(ingest_config, clock) + expected = {} if clock == "seconds" else {"timestamp_unit": clock} + assert ingest_config == expected diff --git a/packages/kumo-core/tests/test_pquery.py b/packages/kumo-core/tests/test_pquery.py index 1dd0e985f..2b50b80be 100644 --- a/packages/kumo-core/tests/test_pquery.py +++ b/packages/kumo-core/tests/test_pquery.py @@ -758,6 +758,39 @@ def test_latest_temporal_timestamp_uses_source_maximum(tmp_path, monkeypatch): assert latest_temporal_timestamp(query, str(tmp_path / "index")) == maximum +@pytest.mark.parametrize("maximum, expected", [ + ("2023-11-14T22:13:20.123456789", 1_700_000_000_123_456_789), + ("kumo_timestamp_2023-11-14 22:13:20.123456789+00:00", + 1_700_000_000_123_456_789), + ("2024-04-29 23:59:26.477365909", 1_714_435_166_477_365_909), + ("2024-04-29T23:59:26.477365909", 1_714_435_166_477_365_909), + ("2024-04-29T23:59:26.477365909+02:00", 1_714_427_966_477_365_909), + ("2024-04-29T23:59:26.477365909-05:30", 1_714_454_966_477_365_909), + ("20240429T235926.477365909", 1_714_435_166_477_365_909), + ("2024-04-29 23:59:26,477365909", 1_714_435_166_477_365_909), + ("2024-W18-1T23:59:26.477365909", 1_714_435_166_477_365_909), + ("2024-04-29 23:59:26.4773659091234", 1_714_435_166_477_365_909), + ("2024-W18-1", 1_714_348_800_000_000_000), + ("2024-04-29", 1_714_348_800_000_000_000), + ("2024-04-29 23:59:26.477", 1_714_435_166_477_000_000), + ("1969-12-31 23:59:59.999999999", -1), +]) +def test_latest_temporal_timestamp_keeps_subsecond_statistics( + monkeypatch, maximum, expected): + monkeypatch.setattr( + "kumo_core.pquery.read_col_stats", + lambda *args: json.dumps({ + "event_time__NA_FRACTION__False": 0.0, + "event_time__MIN__False": "1969-01-01T00:00:00", + "event_time__MAX__False": maximum, + })) + graph = normalize_graph(forecasting_graph_config()) + query = parse_and_validate_query( + "PREDICT COUNT(events.*, 0, 7, days) FOR EACH articles.aid", graph) + + assert latest_temporal_timestamp(query, "/index") == expected + + @pytest.mark.parametrize("order", [ ("first", "second"), ("second", "first"), diff --git a/packages/kumo-core/tests/test_prediction.py b/packages/kumo-core/tests/test_prediction.py index 779cc083f..8043f06ff 100644 --- a/packages/kumo-core/tests/test_prediction.py +++ b/packages/kumo-core/tests/test_prediction.py @@ -197,7 +197,8 @@ def __init__(self, tmp_path): self.tmp_path = tmp_path self.compile_options = None - def DiskGraphPlanCompiler(self): + def DiskGraphPlanCompiler(self, *, index_path): + assert index_path == str(self.tmp_path / "index") runtime = self class Compiler: diff --git a/packages/kumo-core/tests/test_train_table.py b/packages/kumo-core/tests/test_train_table.py index 9d156dc1d..f601f5d14 100644 --- a/packages/kumo-core/tests/test_train_table.py +++ b/packages/kumo-core/tests/test_train_table.py @@ -49,7 +49,8 @@ class FakeRuntime: def __init__(self, staging_dir): self.calls = [] - self.DiskGraphPlanCompiler = lambda: FakeCompiler(self.calls) + self.DiskGraphPlanCompiler = lambda *, index_path: FakeCompiler( + self.calls) self.staging_dir = staging_dir def execute_tt_gen(self, plan, index_path, tt_id): diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index b5ecf3299..d2388efd3 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -68,6 +68,7 @@ from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumoapi.typing import TimestampPrecision, validate_timestamp_precision from kumo.engine._scratch import ( fingerprint, @@ -345,7 +346,10 @@ def validate(self) -> None: """ self.native_graph - def snapshot(self) -> GraphSnapshot: + def snapshot( + self, + timestamp_precision: TimestampPrecision = 'seconds', + ) -> GraphSnapshot: r"""Materialize this graph: read each table's connector, run the result through `kumo-core` to build column stats and the on-disk index. The index is stored under a directory keyed by both `self.id` @@ -354,9 +358,16 @@ def snapshot(self) -> GraphSnapshot: and `GraphSnapshot`'s docstring for the resulting `id` shape. If a valid build already exists for that exact key, ingest is skipped and it's reused as is. + + `timestamp_precision` is the clock the index stores timestamps in: + `"nanoseconds"`, `"microseconds"`, `"milliseconds"`, `"seconds"` (the + default), `"hours"` or `"days"`. Window lengths truncate to the clock, + so a window shorter than one tick is empty. Each clock is cached as a + separate snapshot. """ from filelock import FileLock + validate_timestamp_precision(timestamp_precision) self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) @@ -383,9 +394,11 @@ def snapshot(self) -> GraphSnapshot: # processes, and a lock-file-existence mutex would leave a stale # lock behind after a crash. with FileLock(f'{persist_dir}.lock'): - return self._snapshot_no_lock(graph_key, persist_dir) + return self._snapshot_no_lock(graph_key, persist_dir, + timestamp_precision) - def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ + def _snapshot_no_lock(self, graph_key: str, persist_dir: Path, + timestamp_precision: TimestampPrecision = 'seconds') \ -> GraphSnapshot: r"""`snapshot()`'s body, run under its `FileLock`. Not meant to be called on its own -- `persist_dir`, unlike everything else here, @@ -466,8 +479,12 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ source_paths[alias] = str( prepared.source.get('path') or dest_dir) - data_version = fingerprint( - json.dumps(table_data_versions, sort_keys=True).encode()) + version_parts = [ + json.dumps(table_data_versions, sort_keys=True).encode() + ] + if timestamp_precision != 'seconds': + version_parts.append(timestamp_precision.encode()) + data_version = fingerprint(*version_parts) handle = DiskgraphSnapshotHandle(graph_key=graph_key, graph_data_version=data_version) output_dir = Path(store.path(handle)) @@ -494,6 +511,7 @@ def _snapshot_no_lock(self, graph_key: str, persist_dir: Path) \ table_data_versions=table_data_versions, output_dir=str(output_dir), source_paths=source_paths, + timestamp_precision=timestamp_precision, execution=IngestExecutionConfig( success_marker_content=predicted.to_json()), ), diff --git a/packages/kumo-engine/test/test_snapshot_clock.py b/packages/kumo-engine/test/test_snapshot_clock.py new file mode 100644 index 000000000..d5dbe5b61 --- /dev/null +++ b/packages/kumo-engine/test/test_snapshot_clock.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +from types import SimpleNamespace +from typing import Any, get_args + +import pytest +from kumoapi.typing import TimestampPrecision + +from kumo.engine._scratch import fingerprint +from kumo.engine.graph import Graph + +CLOCKS = get_args(TimestampPrecision) + + +@pytest.fixture +def ingests(tmp_path, monkeypatch): + import kumo_core.diskgraph.ingest as ingest + import kumo_core.diskgraph.snapshot as snapshots + + import kumo.engine._scratch as scratch + + monkeypatch.setattr(scratch, '_root', tmp_path) + monkeypatch.setattr(Graph, 'id', 'graph-key') + monkeypatch.setattr(Graph, '_graph_config', {'tables': []}) + monkeypatch.setattr(Graph, 'validate', lambda self: None) + monkeypatch.setattr(snapshots, 'check_diskgraph_output', lambda path: None) + monkeypatch.setattr(snapshots.DiskgraphIngestResult, 'for_index', + lambda path: SimpleNamespace(to_json=lambda: '{}')) + specs = [] + + def execute(graph_config, spec): + specs.append(spec) + return SimpleNamespace(index_path=spec.output_dir) + + monkeypatch.setattr(ingest, 'execute_diskgraph_ingest', execute) + return specs + + +def _graph(): + class Source: + def prepare_ingest(self, connector, table_name, destination, **_): + return SimpleNamespace(credentials=None, version='v1', + source={'path': destination}) + + table: Any = SimpleNamespace( + table_name='events', + definition=SimpleNamespace(cols=[], source_table=None), + connector=SimpleNamespace(source_backend=Source(), + validate_scratch_dir=lambda path: None)) + return Graph(tables=[table]) + + +@pytest.mark.parametrize('clock', CLOCKS) +def test_snapshot_ingests_in_the_requested_clock(ingests, clock): + _graph().snapshot(timestamp_precision=clock) + assert [spec.timestamp_precision for spec in ingests] == [clock] + + +def test_seconds_snapshot_keeps_the_existing_cache_key(ingests): + snapshot = _graph().snapshot() + data_version = fingerprint( + json.dumps({'events': 'v1'}, sort_keys=True).encode()) + assert snapshot.id == f'graph-key@{data_version[:16]}' + + +def test_each_clock_has_its_own_snapshot(ingests): + ids = {_graph().snapshot(timestamp_precision=clock).id for clock in CLOCKS} + assert len(ids) == len(CLOCKS) + + +def test_unknown_clock_is_rejected(): + with pytest.raises(ValueError, match='timestamp_precision'): + _graph().snapshot(timestamp_precision='months') diff --git a/packages/kumo-sdk/kumoai_studio/backend/base.py b/packages/kumo-sdk/kumoai_studio/backend/base.py index 33fb6f2a7..d4368129f 100644 --- a/packages/kumo-sdk/kumoai_studio/backend/base.py +++ b/packages/kumo-sdk/kumoai_studio/backend/base.py @@ -86,6 +86,7 @@ from kumoapi.source_table import LocalFileSourceTable, SourceTableType from kumoapi.task import TaskType from kumoapi.train import TrainingStage +from kumoapi.typing import TimestampPrecision from kumoai_studio.identity import Identity @@ -330,6 +331,7 @@ def create_snapshot( *, refresh_source: bool = False, project_id: Optional[str] = None, + timestamp_precision: TimestampPrecision = 'seconds', ) -> snapshot_api.GraphSnapshotID: ... diff --git a/packages/kumo-sdk/kumoai_studio/client/graph.py b/packages/kumo-sdk/kumoai_studio/client/graph.py index fed3e33db..1e7d7014e 100644 --- a/packages/kumo-sdk/kumoai_studio/client/graph.py +++ b/packages/kumo-sdk/kumoai_studio/client/graph.py @@ -14,6 +14,7 @@ GraphValidationRequest, ) from kumoapi.json_serde import to_json_dict +from kumoapi.typing import TimestampPrecision, validate_timestamp_precision from kumoai_studio.client.endpoints import GraphEndpoints from kumoai_studio.client.transport import Transport @@ -79,16 +80,28 @@ def create_snapshot( *, refresh_source: bool = False, project_id: Optional[str] = None, + timestamp_precision: TimestampPrecision = 'seconds', ) -> GraphSnapshotID: + validate_timestamp_precision(timestamp_precision) params: Dict[str, Any] = { 'graph_id': graph_id, 'refresh_source': refresh_source } if project_id is not None: params['project_id'] = project_id + if timestamp_precision != 'seconds': + params['timestamp_precision'] = timestamp_precision resp = self._client._request(GraphEndpoints.create_snapshot, params=params) raise_on_error(resp) + # Servers that predate index clocks send no header and build seconds. + built_precision = resp.headers.get('X-Kumo-Timestamp-Precision', + 'seconds') + if built_precision != timestamp_precision: + raise RuntimeError( + f"Requested a {timestamp_precision} snapshot but the server " + f"built a {built_precision} one; upgrade the server to use " + f"this clock.") return GraphSnapshotID(parse_id_response(resp)) def get_snapshot( diff --git a/packages/kumo-sdk/kumoai_studio/graph/graph.py b/packages/kumo-sdk/kumoai_studio/graph/graph.py index 6d98f4bcb..a3d324f2f 100644 --- a/packages/kumo-sdk/kumoai_studio/graph/graph.py +++ b/packages/kumo-sdk/kumoai_studio/graph/graph.py @@ -31,6 +31,7 @@ from kumoapi.common import JobStatus from kumoapi.data_snapshot import GraphSnapshotID from kumoapi.source_table import CustomSQLSourceTable +from kumoapi.typing import TimestampPrecision from typing_extensions import Self from kumoai_studio.backend import cloud_backend @@ -261,6 +262,7 @@ def __init__( # Cached from backend: self._graph_snapshot_id: Optional[GraphSnapshotID] = None + self._graph_snapshot_precision: TimestampPrecision = 'seconds' @property def kumo_project_id(self) -> str: @@ -532,6 +534,7 @@ def snapshot( *, force_refresh: bool = False, non_blocking: bool = False, + timestamp_precision: Optional[TimestampPrecision] = None, ) -> snapshot_api.GraphSnapshotID: r"""Takes a *snapshot* of this graph's underlying data, and returns a unique identifier for this snapshot. @@ -553,12 +556,20 @@ def snapshot( after creating the snapshot, or await completion of the snapshot. If :obj:`True`, the snapshot will proceed in the background, and will be used for any downstream job. + timestamp_precision: The clock the snapshot index stores timestamps + in: ``"nanoseconds"``, ``"microseconds"``, + ``"milliseconds"``, ``"seconds"``, ``"hours"`` or ``"days"``. + Each clock is a separate snapshot. Defaults to the clock of + this graph's current snapshot, or seconds if it has none. Raises: RuntimeError: if ``non_blocking`` is set to :obj:`False` and the graph snapshot fails. """ - if self._graph_snapshot_id is None or force_refresh: + if timestamp_precision is None: + timestamp_precision = self._graph_snapshot_precision + if (self._graph_snapshot_id is None or force_refresh + or timestamp_precision != self._graph_snapshot_precision): self.save() if not force_refresh: snapshotted_table_names: List[str] = [] @@ -576,7 +587,9 @@ def snapshot( graph_id=self.id, refresh_source=True, project_id=self.kumo_project_id, + timestamp_precision=timestamp_precision, )) + self._graph_snapshot_precision = timestamp_precision logger.info("Graph snapshot with identifier %s created.", self._graph_snapshot_id) diff --git a/packages/kumo-sdk/test/graph/test_snapshot_clock.py b/packages/kumo-sdk/test/graph/test_snapshot_clock.py new file mode 100644 index 000000000..19fe7fd07 --- /dev/null +++ b/packages/kumo-sdk/test/graph/test_snapshot_clock.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from typing import get_args +from unittest import mock + +import pytest +from kumoapi.data_snapshot import GraphSnapshotID +from kumoapi.typing import TimestampPrecision + +from kumoai_studio.client.client import KumoClient +from kumoai_studio.graph import Graph + +BASE_URL = 'https://kumo.ai' +SNAPSHOTS_URL = f'{BASE_URL}/v1/graphsnapshots' +CLOCKS = get_args(TimestampPrecision) + + +def _create_snapshot(requests_mock, echo, **kwargs): + headers = {} if echo is None else {'X-Kumo-Timestamp-Precision': echo} + requests_mock.post(SNAPSHOTS_URL, json={'id': 'graph@v1'}, headers=headers) + return KumoClient(BASE_URL, + 'api-key').graph_api.create_snapshot('graph', **kwargs) + + +@pytest.mark.parametrize('clock', CLOCKS) +def test_snapshot_in_a_clock_the_server_built(requests_mock, clock): + _create_snapshot(requests_mock, clock, timestamp_precision=clock) + sent = requests_mock.last_request.qs.get('timestamp_precision') + assert sent == (None if clock == 'seconds' else [clock]) + + +def test_server_without_clocks_still_builds_seconds(requests_mock): + assert _create_snapshot(requests_mock, None) == 'graph@v1' + + +def test_server_without_clocks_cannot_build_other_clocks(requests_mock): + with pytest.raises(RuntimeError, match='built a seconds one'): + _create_snapshot(requests_mock, None, + timestamp_precision='milliseconds') + + +def test_server_building_a_different_clock_is_an_error(requests_mock): + with pytest.raises(RuntimeError, match='built a seconds one'): + _create_snapshot(requests_mock, 'seconds', timestamp_precision='hours') + + +def test_unknown_clock_is_rejected_before_any_request(requests_mock): + with pytest.raises(ValueError, match='timestamp_precision'): + _create_snapshot(requests_mock, 'months', timestamp_precision='months') + assert not requests_mock.called + + +def _graph(monkeypatch, *snapshot_ids): + import kumoai_studio.graph.graph as graph_module + + backend = mock.Mock() + backend.graph.create_snapshot.side_effect = [ + GraphSnapshotID(snapshot_id) for snapshot_id in snapshot_ids + ] + backend.graph.get_snapshot.return_value = SimpleNamespace(table_ids={}) + monkeypatch.setattr(graph_module, 'cloud_backend', lambda: backend) + graph = Graph(project_id='project-test') + monkeypatch.setattr(graph, 'save', lambda: None) + return graph, backend + + +def test_changing_the_clock_takes_a_new_snapshot(monkeypatch): + graph, backend = _graph(monkeypatch, 'graph@seconds', 'graph@days') + + assert graph.snapshot(non_blocking=True) == 'graph@seconds' + assert graph.snapshot(non_blocking=True) == 'graph@seconds' + assert graph.snapshot(non_blocking=True, + timestamp_precision='days') == 'graph@days' + clocks = [ + call.kwargs['timestamp_precision'] + for call in backend.graph.create_snapshot.call_args_list + ] + assert clocks == ['seconds', 'days'] + + +def test_a_snapshot_without_a_clock_reuses_the_current_one(monkeypatch): + graph, backend = _graph(monkeypatch, 'graph@milliseconds') + + graph.snapshot(non_blocking=True, timestamp_precision='milliseconds') + + assert graph.snapshot(non_blocking=True) == 'graph@milliseconds' + assert backend.graph.create_snapshot.call_count == 1 diff --git a/pool/activities.py b/pool/activities.py index bf91014b5..489edaffb 100644 --- a/pool/activities.py +++ b/pool/activities.py @@ -1417,6 +1417,8 @@ def heartbeat() -> None: data_persist_dir=ingest_args.data_persist_dir, table_data_versions=( ingest_args.table_data_versions), + timestamp_precision=( + ingest_args.timestamp_precision), source_paths=(source_paths or None), source_specs=(source_specs or None), subprocess_environment=(subprocess_environment diff --git a/pool/diskgraph_args.py b/pool/diskgraph_args.py index 22bbe0e22..365768888 100644 --- a/pool/diskgraph_args.py +++ b/pool/diskgraph_args.py @@ -10,6 +10,8 @@ from dataclasses import dataclass from typing import Dict +from kumoapi.typing import TimestampPrecision + from rest.config.data_backend import ConnectorExecutionRef from rest.config.graph_v2 import GraphDefinition from rest.config.table_v2 import TableDefinition @@ -25,3 +27,4 @@ class DiskgraphIngestArgs: connector_specs: Dict[str, ConnectorExecutionRef] all_graph_tables: Dict[str, TableDefinition] source_size_bytes_by_table: Dict[str, int] + timestamp_precision: TimestampPrecision = 'seconds' diff --git a/rest/backends/data_backend/activities/data_backend_activities.py b/rest/backends/data_backend/activities/data_backend_activities.py index c6800c84d..367fba8a7 100644 --- a/rest/backends/data_backend/activities/data_backend_activities.py +++ b/rest/backends/data_backend/activities/data_backend_activities.py @@ -641,14 +641,18 @@ async def _refresh_graph_data_version( md5_hasher = md5() for table_name, table_version in sorted(table_versions.items()): md5_hasher.update(f'{table_name}={table_version}'.encode()) + if graph_input.timestamp_precision != 'seconds': + md5_hasher.update(('|timestamp_precision=' + f'{graph_input.timestamp_precision}').encode()) graph_data_version = md5_hasher.hexdigest() # Save graph and table version info. await data_backend_cache().graph_cache.save_latest_graph_version.\ coroutine(graph_input.graph_key, graph_data_version) - version_info = GraphDataVersionInfo(GraphVersion(graph_data_version), - table_versions) + version_info = GraphDataVersionInfo( + GraphVersion(graph_data_version), table_versions, + timestamp_precision=graph_input.timestamp_precision) await DataBackendActivities.save_graph_data_version( graph_input, version_info, job_id) return version_info diff --git a/rest/backends/data_backend/data_backend.py b/rest/backends/data_backend/data_backend.py index 7b3a1d439..d4d84006f 100644 --- a/rest/backends/data_backend/data_backend.py +++ b/rest/backends/data_backend/data_backend.py @@ -266,6 +266,7 @@ def _compute_bp_sample_col_stats( sample, stats_dict: Dict[BaseColumnConfig, List[ColStatType]], data_version: TableDataVersion, + timestamp_precision: str, ) -> Dict[BaseColumnConfig, Dict[ColStatType, Stats]]: """Reduce a bounded prediction sample, including SCORE grouped by CLASS.""" from kumo.connector.connector import _compute_col_stats_from_sample @@ -275,6 +276,7 @@ def _compute_bp_sample_col_stats( sample, approx=True, data_version=data_version, + timestamp_precision=timestamp_precision, ) score_config = next( (config @@ -1286,6 +1288,10 @@ async def start_bp_stats_activity( raise ValueError( 'Direct Databricks stats stage contains no Parquet files.') + timestamp_precision = (job_input.graph_version_info.timestamp_precision + if job_input.graph_version_info is not None else + 'seconds') + entity_col_name = job_input.predictive_query.get_bp_entity_col_name() (bp_connector_spec, bp_table_def, bp_data_version, bp_col_stats_set) = await generate_bp_stats_workflow_args( @@ -1363,6 +1369,7 @@ async def start_bp_stats_activity( sample, request, bp_data_version, + timestamp_precision, ) if not col_stat_dict_complete(sampled_stats, request): raise RuntimeError( @@ -1438,6 +1445,7 @@ async def start_bp_stats_activity( parquet_paths=parquet_paths, kumo_column_specs=kumo_column_specs, data_version=str(bp_data_version), + timestamp_precision=timestamp_precision, ) finally: if staging_dir is not None: diff --git a/rest/backends/data_backend/workflows/diskgraph_ingest_workflow.py b/rest/backends/data_backend/workflows/diskgraph_ingest_workflow.py index a6fb99eff..74b59ce62 100644 --- a/rest/backends/data_backend/workflows/diskgraph_ingest_workflow.py +++ b/rest/backends/data_backend/workflows/diskgraph_ingest_workflow.py @@ -62,6 +62,7 @@ async def run( source_size_bytes: int, resource_hint: ResourceHint | None, source_size_bytes_by_table: dict[str, int], + timestamp_precision: str = 'seconds', ) -> DiskgraphIngestResult: snapshot_id = DataSnapshot.GRAPH.canonicalize( f'{_graph_key_to_v2(graph_key)}@{graph_data_version}') @@ -83,6 +84,7 @@ async def run( source_size_bytes=source_size_bytes, source_size_bytes_by_table=source_size_bytes_by_table, resource_hint=current_resource_hint, + timestamp_precision=timestamp_precision, ) except _LocalNoSpaceAtTier as failure: failed_tier = failure.tier_ssd_gib @@ -126,6 +128,7 @@ async def _ingest_once( source_size_bytes: int, source_size_bytes_by_table: dict[str, int], resource_hint: ResourceHint | None, + timestamp_precision: str, ) -> DiskgraphIngestResult: # The controller computes the on-SSD estimate # (source × ingest_space_multiplier + buffer) from its own config, @@ -161,6 +164,7 @@ async def _ingest_once( all_graph_tables=all_graph_tables, source_size_bytes_by_table=( source_size_bytes_by_table), + timestamp_precision=timestamp_precision, ), ], schedule_to_close_timeout=timedelta( diff --git a/rest/backends/data_backend/workflows/graph_refresh_workflow.py b/rest/backends/data_backend/workflows/graph_refresh_workflow.py index e791a2c01..d3ee8504c 100644 --- a/rest/backends/data_backend/workflows/graph_refresh_workflow.py +++ b/rest/backends/data_backend/workflows/graph_refresh_workflow.py @@ -301,6 +301,10 @@ async def _run_diskgraph_colstats_pipeline( use_ingest_stats_v2 = workflow.patched('diskgraph-ingest-stats-v2') if use_ingest_stats_v2: wf_id = f'{wf_id}:stats-v2' + index_clock_args = [] + if workflow.patched('diskgraph-index-clock-v1'): + index_clock_args.append( + self._graph_version_info.timestamp_precision) child_task = await workflow_utils.start_child_data_wf( diskgraph_ingest_workflow.DiskgraphIngestWorkflow.run, self._ctx, @@ -317,6 +321,7 @@ async def _run_diskgraph_colstats_pipeline( source_size_bytes, self._resource_hint, source_size_bytes_by_table, + *index_clock_args, ], execution_timeout=timedelta( hours=DATA_BACKEND_INGEST_TIMEOUT_HOURS), diff --git a/rest/config/data_backend.py b/rest/config/data_backend.py index 3790b0a81..9a7219b5c 100644 --- a/rest/config/data_backend.py +++ b/rest/config/data_backend.py @@ -6,6 +6,8 @@ from dataclasses import dataclass, field from typing import Any, ClassVar, Dict, Mapping, Optional, Union +from kumoapi.typing import TimestampPrecision + from io_utils import kumo_dataclass from kumo.config.predictive_query import PredictiveQueryV2Config from kumo.data.cache.graph_cache import GraphVersion @@ -148,3 +150,4 @@ class GraphDataVersionInfo: # Name/alias of the graph at the time the snapshot was taken. # None if the graph had no alias set. graph_name: Optional[str] = None + timestamp_precision: TimestampPrecision = 'seconds' diff --git a/rest/config/graph_v2.py b/rest/config/graph_v2.py index 588675f1f..e6864b337 100644 --- a/rest/config/graph_v2.py +++ b/rest/config/graph_v2.py @@ -11,6 +11,7 @@ import kumoapi.data_snapshot as snapshot_api import kumoapi.graph as api from kumo_core.types.schema import FPID +from kumoapi.typing import TimestampPrecision from pydantic import ( BaseModel, ValidationError, @@ -438,6 +439,7 @@ class GraphDataComputeInput: table_aliases: Dict[TableID, TableName] = field(default_factory=dict) pinned_data_versions: Dict[TableID, TableDataVersion] = field(default_factory=dict) + timestamp_precision: TimestampPrecision = 'seconds' def __post_init__(self) -> None: for connector_id, reference in self.connectors.items(): diff --git a/rest/driver/data_snapshot_api_driver.py b/rest/driver/data_snapshot_api_driver.py index 0cff1a80a..1e90133d0 100644 --- a/rest/driver/data_snapshot_api_driver.py +++ b/rest/driver/data_snapshot_api_driver.py @@ -13,6 +13,7 @@ from kumo_core.types import ConnectorType, TableStatType from kumo_core.types.schema import FPID from kumoapi.common import JobStatus +from kumoapi.typing import TimestampPrecision from temporalio.client import ( WorkflowExecutionDescription, WorkflowExecutionStatus, @@ -344,6 +345,7 @@ async def create_graph_snapshot( base_graph_name: Optional[str] = None, *, project_id: ProjectID, + timestamp_precision: TimestampPrecision = 'seconds', ) -> snapshot_interfaces.GraphSnapshotID: """Create a graph data snapshot for an existing graph: the data snapshot corresponds to a specific graph data version. @@ -388,6 +390,7 @@ async def create_graph_snapshot( # immutable) TableDefinition. `refresh_graph_data_version` reads # from `pinned_data_versions` to override the latest-version fetch. data_compute_input.pinned_data_versions = dict(table_versions) + data_compute_input.timestamp_precision = timestamp_precision version_info = await DataBackend.refresh_graph_data_version_in_activity( self._data_backend_ctx, data_compute_input, diff --git a/rest/public/routers/data_snapshot.py b/rest/public/routers/data_snapshot.py index 134469b2c..301ff0a26 100644 --- a/rest/public/routers/data_snapshot.py +++ b/rest/public/routers/data_snapshot.py @@ -6,10 +6,11 @@ from typing import Callable import kumoapi.data_snapshot as snapshot_interfaces -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, Response from fastapi.params import Query from kumoapi.graph import EdgeHealthResponse from kumoapi.table import TableDefinition as ApiTableDefinition +from kumoapi.typing import TimestampPrecision from typing_extensions import Annotated from io_utils import logging @@ -18,7 +19,10 @@ from rest.config.status import IDResponse from rest.config.table_v2 import TableDefinition, TableSnapshotResource from rest.deps.authentication import Authentication -from rest.driver.data_snapshot_api_driver import DataSnapshotApiDriver +from rest.driver.data_snapshot_api_driver import ( + DataSnapshotApiDriver, + get_graph_snapshot_info, +) from rest.utils.public import wrap_driver_errors from rest.utils.routers import StrictQueryAPIRoute, wrap_validation_exceptions from statedb.config.common import GraphIDv2, ProjectID @@ -153,11 +157,14 @@ def __init__( async def create( self, + response: Response, graph_id: Annotated[str, Query(description='Graph resource identifier')], project_id: Annotated[str, Query()], refresh_source: Annotated[ bool, Query(description=_refresh_source_desc)] = False, + timestamp_precision: Annotated[TimestampPrecision, + Query()] = 'seconds', ) -> IDResponse: """Create a GraphSnapshotResource.""" with wrap_driver_errors(), wrap_validation_exceptions(): @@ -167,7 +174,15 @@ async def create( graph_fpid, refresh_source, project_id=ProjectID(project_id), + timestamp_precision=timestamp_precision, ) + built_precision = timestamp_precision + if timestamp_precision != 'seconds': + # A worker that predates index clocks records a seconds build. + _, version_info = await get_graph_snapshot_info( + graph_snapshot_id) + built_precision = version_info.timestamp_precision + response.headers['X-Kumo-Timestamp-Precision'] = built_precision return IDResponse(graph_snapshot_id) async def get( diff --git a/test/connector/local/test_timestamp_convert.py b/test/connector/local/test_timestamp_convert.py index 8a12b70e3..969ee065a 100644 --- a/test/connector/local/test_timestamp_convert.py +++ b/test/connector/local/test_timestamp_convert.py @@ -5,7 +5,7 @@ r"""Tests for the timestamp/date converter. Covers the epoch-domain arithmetic kumo still owns (``div_ceil`` / -``timestamp_from_int`` / ``days_from_secs`` / ``timestamp_secs_from_array`` / +``timestamp_from_int`` / ``days_from_secs`` / ``timestamp_ticks_from_array`` / Date32), and the DiskGraph-backed apply path: the exhaustive format-x-value matrix, the vectorized entry points, and MIN/MAX. @@ -261,31 +261,31 @@ def test_i64_min_magnitude_no_overflow(): # =========================================================================== # Native temporal path (no parsing: the unit rides the dtype) # =========================================================================== -def test_timestamp_secs_from_array_ceiling_ns(): +def test_timestamp_ticks_from_array_ceiling_ns(): base = int(pd.Timestamp('2020-01-01 00:00:00').timestamp()) ser = pd.Series([ pd.Timestamp('2020-01-01 00:00:00.500'), # sub-second -> ceil up pd.Timestamp('2020-01-01 00:00:01'), # exact second ]) - assert list(tc.timestamp_secs_from_array(ser)) == [base + 1, base + 1] + assert list(tc.timestamp_ticks_from_array(ser)) == [base + 1, base + 1] -def test_timestamp_secs_from_array_ceiling_ms_unit(): +def test_timestamp_ticks_from_array_ceiling_ms_unit(): base = int(pd.Timestamp('2020-01-01 00:00:00').timestamp()) ser = pd.Series([pd.Timestamp('2020-01-01 00:00:00.500')]).astype( 'datetime64[ms]') - assert list(tc.timestamp_secs_from_array(ser)) == [base + 1] + assert list(tc.timestamp_ticks_from_array(ser)) == [base + 1] -def test_timestamp_secs_from_array_nat_is_none(): +def test_timestamp_ticks_from_array_nat_is_none(): ser = pd.Series([pd.Timestamp('2020-01-01'), pd.NaT]) - assert tc.timestamp_secs_from_array(ser)[1] is None + assert tc.timestamp_ticks_from_array(ser)[1] is None -def test_timestamp_secs_from_array_tz_aware_normalized_to_utc(): +def test_timestamp_ticks_from_array_tz_aware_normalized_to_utc(): ser = pd.Series([pd.Timestamp('2020-01-01 00:00:00', tz='US/Eastern')]) want = int(pd.Timestamp('2020-01-01 00:00:00', tz='US/Eastern').timestamp()) - assert list(tc.timestamp_secs_from_array(ser)) == [want] + assert list(tc.timestamp_ticks_from_array(ser)) == [want] def _scalar_native_seconds(ser): @@ -308,14 +308,14 @@ def test_parity_native_vectorized_matches_scalar(resolution): pd.NaT, # null -> None, never epoch ]) ser = ts.astype(f'datetime64[{resolution}]') - out = tc.timestamp_secs_from_array(ser) + out = tc.timestamp_ticks_from_array(ser) assert list(out) == _scalar_native_seconds(ser) def test_parity_native_all_nat_and_empty(): - assert list(tc.timestamp_secs_from_array( + assert list(tc.timestamp_ticks_from_array( pd.Series([pd.NaT, pd.NaT]))) == [None, None] - assert list(tc.timestamp_secs_from_array( + assert list(tc.timestamp_ticks_from_array( pd.Series([], dtype='datetime64[ns]'))) == [] @@ -450,3 +450,46 @@ def test_hint_constructors_shape(): assert TemporalHint.of_format('%Y%m%d').kind is HintKind.FORMAT assert TemporalHint.native().kind is HintKind.NATIVE assert TemporalHint.unresolved().kind is HintKind.UNRESOLVED + + +_TYPED = pd.Series( + pd.to_datetime([ + '1969-12-31 23:59:59.5', '2024-01-01 00:00:00.000000001', + '1700-01-01 00:00:00.25', None + ], format='ISO8601')) +_CLOCKS = list(tc.PANDAS_UNIT_OF_INDEX_CLOCK) + + +@pytest.mark.parametrize('clock', _CLOCKS) +@pytest.mark.parametrize('series', [ + _TYPED.astype('datetime64[s]'), + _TYPED.astype('datetime64[ms]'), + _TYPED.astype('datetime64[us]'), + _TYPED, + _TYPED.dt.tz_localize('US/Pacific'), +], ids=['s', 'ms', 'us', 'ns', 'tz']) +def test_typed_columns_convert_like_diskgraph_in_every_clock(series, clock): + hint = TemporalHint.native() + assert list(tc.to_epoch_ticks(series, hint, clock)) == list( + tc._resolve(series, hint, days=False, timestamp_precision=clock)) + + +@pytest.mark.parametrize('clock,expected', [ + ('nanoseconds', '2023-11-14 22:13:20.123'), + ('microseconds', '2023-11-14 22:13:20.123'), + ('milliseconds', '2023-11-14 22:13:20.123'), + ('seconds', '2023-11-14 22:13:20'), + ('hours', '2023-11-14 22:00:00'), + ('days', '2023-11-14'), +]) +def test_text_statistics_truncate_into_the_snapshot_clock(clock, expected): + from kumoapi.typing import Stype + + from kumo.config import BaseColumnConfig + from kumo.connector.utils import coerce_timestamp_stype_series_for_stats + + config = BaseColumnConfig(name='ts', stype=Stype.timestamp) + coerced = coerce_timestamp_stype_series_for_stats( + pd.Series(['2023-11-14 22:13:20.123']), config, clock) + + assert coerced.tolist() == [pd.Timestamp(expected)] diff --git a/test/data/generic/test_diskgraph_batch_pred.py b/test/data/generic/test_diskgraph_batch_pred.py index bd2bdc5ba..288cb37b9 100644 --- a/test/data/generic/test_diskgraph_batch_pred.py +++ b/test/data/generic/test_diskgraph_batch_pred.py @@ -91,6 +91,8 @@ def setup_dirs(tmpdir): out_path = os.path.join(tmpdir, "output") os.makedirs(idx_path, exist_ok=True) os.makedirs(out_path, exist_ok=True) + with open(os.path.join(idx_path, "index_metadata.json"), "w") as f: + f.write("{}") return str(idx_path), str(out_path) diff --git a/test/data/generic/test_diskgraph_generate_train_table.py b/test/data/generic/test_diskgraph_generate_train_table.py index b842b93a2..2a3cc57a3 100644 --- a/test/data/generic/test_diskgraph_generate_train_table.py +++ b/test/data/generic/test_diskgraph_generate_train_table.py @@ -79,6 +79,9 @@ def test_returns_manifest_from_native_artifact_execution(self, tmp_path): class _FakeCompiler: + def __init__(self, *, index_path): + pass + def compile(self, **kwargs): return "{\"fake\": true}" diff --git a/test/data/generic/test_diskgraph_predict_plan_ttng.py b/test/data/generic/test_diskgraph_predict_plan_ttng.py index d6cc18037..26b95a54e 100644 --- a/test/data/generic/test_diskgraph_predict_plan_ttng.py +++ b/test/data/generic/test_diskgraph_predict_plan_ttng.py @@ -35,6 +35,9 @@ def stub_compiler(monkeypatch): captured = {} class _StubCompiler: + def __init__(self, *, index_path): + captured['index_path'] = index_path + def compile(self, **kwargs): captured['kwargs'] = kwargs plan = { @@ -62,6 +65,7 @@ def _build(task_type, *, list_array, stub_compiler): plan_json = bp.build_ttng_predict_plan( SimpleNamespace( entity_column='users.user_id'), # only entity_column read + index_path='/index', task_type=task_type, graph=SimpleNamespace(), timestamps=[], @@ -160,6 +164,9 @@ class _TimeUnit: Seconds = 'Seconds' class _Compiler: + def __init__(self, *, index_path): + captured['compiler_index_path'] = index_path + def compile(self, **kwargs): captured['compile'] = kwargs return json.dumps({'task_type': kwargs['task_type']}) @@ -228,6 +235,7 @@ def _read_tt_parquet_batches(path): assert (rows, parts) == (7, 1) assert captured['execute']['index_path'] == index_path + assert captured['compiler_index_path'] == index_path assert captured['execute']['tt_id'] == 'pred_bp-job' assert captured['compile']['is_prediction'] is True assert captured['compile']['task_type'] == TaskType.FORECASTING.value diff --git a/test/data/generic/test_diskgraph_runtime_adapter.py b/test/data/generic/test_diskgraph_runtime_adapter.py index e67337f35..1176fec90 100644 --- a/test/data/generic/test_diskgraph_runtime_adapter.py +++ b/test/data/generic/test_diskgraph_runtime_adapter.py @@ -21,6 +21,7 @@ from kumo_core.diskgraph.runtime_adapter import ( _expand_lp_exclude, attach_tensor_frames, + engine_int64_to_ns, runtime_batch_to_hetero, ) from kumo_core.encoders.target.categorical import CategoricalTargetEncoder @@ -1397,6 +1398,16 @@ def test_seed_time_ms_unit_scaled_to_ns() -> None: assert hetero["train_table"].seed_time.tolist() == [1700 * 1_000_000] +@pytest.mark.parametrize("unit,tick_ns", [ + ("h", 3_600_000_000_000), + ("D", 86_400_000_000_000), +]) +def test_hours_and_days_ticks_scale_to_ns(unit, tick_ns) -> None: + ticks = pa.array([2, -1, None], type=pa.int64()) + assert engine_int64_to_ns( + ticks, unit).to_pylist() == [2 * tick_ns, -tick_ns, None] + + def test_neighbor_node_time_normalized_to_ns_int64_seconds() -> None: """A bare int64 node ``time`` column (engine seconds) -> ns canonical.""" features = pa.RecordBatch.from_arrays( diff --git a/test/data/generic/test_diskgraph_train_table_forecasting.py b/test/data/generic/test_diskgraph_train_table_forecasting.py index beb86567b..dee1121b3 100644 --- a/test/data/generic/test_diskgraph_train_table_forecasting.py +++ b/test/data/generic/test_diskgraph_train_table_forecasting.py @@ -37,7 +37,7 @@ def compile(self, **kwargs): return '{"time_col_name": "TIMESTAMP"}' diskgraph = SimpleNamespace( - DiskGraphPlanCompiler=lambda: _Compiler(), + DiskGraphPlanCompiler=lambda *, index_path: _Compiler(), SplitSpec=_SplitSpec, execute_tt_plan_with_artifacts=lambda **_kwargs: { 'manifest_path': str(tmp_path / 'manifest.json'), @@ -162,7 +162,7 @@ def compile(self, **kwargs): return '{"time_col_name": "TIMESTAMP"}' diskgraph = SimpleNamespace( - DiskGraphPlanCompiler=lambda: _Compiler(), + DiskGraphPlanCompiler=lambda *, index_path: _Compiler(), SplitSpec=_SplitSpec, execute_tt_gen=lambda *_args: { 'staging_dir': str(tmp_path / 'staging'), @@ -273,7 +273,7 @@ def _execute_tt_gen(plan_json, *_args): } diskgraph = SimpleNamespace( - DiskGraphPlanCompiler=lambda: _Compiler(), + DiskGraphPlanCompiler=lambda *, index_path: _Compiler(), SplitSpec=_SplitSpec, execute_tt_gen=_execute_tt_gen, ) diff --git a/test/data/generic/test_ttng_cross_graph_contamination_e2e.py b/test/data/generic/test_ttng_cross_graph_contamination_e2e.py index d25a1af45..7862878da 100644 --- a/test/data/generic/test_ttng_cross_graph_contamination_e2e.py +++ b/test/data/generic/test_ttng_cross_graph_contamination_e2e.py @@ -386,6 +386,7 @@ def test_predict_table_cross_graph_contamination(contamination_env, caplog, env = contamination_env plan_json = build_ttng_predict_plan( env["validated_query"], + index_path=env["g1_index"], task_type=TaskType.REGRESSION, graph=env["graph"], timestamps=[0], diff --git a/test/rest/backends/test_data_backend.py b/test/rest/backends/test_data_backend.py index 6b4497689..a06f314e3 100644 --- a/test/rest/backends/test_data_backend.py +++ b/test/rest/backends/test_data_backend.py @@ -50,7 +50,9 @@ async def test_embeddings_only_bp_stats_never_opens_predictions() -> None: @pytest.mark.asyncio -async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path) -> None: +@pytest.mark.parametrize('clock', ['seconds', 'milliseconds']) +async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path, + clock) -> None: from types import SimpleNamespace from rest.utils.table import ColStatsSet @@ -71,6 +73,7 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path) -> None: ) job_input.predictive_query.get_bp_entity_col_name.return_value = 'ENTITY' job_input.predictive_query.query_type = mock.sentinel.query_type + job_input.graph_version_info = SimpleNamespace(timestamp_precision=clock) table_def = mock.Mock() table_def.fingerprint.return_value = 'fingerprint' @@ -131,6 +134,7 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path) -> None: kwargs = compute.coroutine.await_args.kwargs assert kwargs['parquet_paths'] == [str(path) for path in selected] assert str(artifact_root / 'foreign.parquet') not in kwargs['parquet_paths'] + assert kwargs['timestamp_precision'] == clock class _UnapprovedSingleActivityError(PublicException): @@ -877,6 +881,7 @@ async def test_direct_databricks_bp_stats_sample_stage_without_internal_copy( connector_id = ConnectorID('connector_' + '2' * 32) job_input = mock.Mock() job_input.id = 'bp-job-direct' + job_input.graph_version_info.timestamp_precision = 'seconds' job_input.config = mock.Mock( outputs=[ DatabricksPredictionOutput( diff --git a/test/rest/backends/test_diskgraph_graph_refresh_workflow.py b/test/rest/backends/test_diskgraph_graph_refresh_workflow.py index 4fc7227b7..c847de258 100644 --- a/test/rest/backends/test_diskgraph_graph_refresh_workflow.py +++ b/test/rest/backends/test_diskgraph_graph_refresh_workflow.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import inspect import json from types import SimpleNamespace from typing import Any, cast, get_type_hints @@ -183,6 +184,7 @@ def _setup_pipeline_test(monkeypatch): wf._graph_version_info = SimpleNamespace( graph_data_version="v1", table_data_versions={"users": "tv1"}, + timestamp_precision="seconds", ) monkeypatch.setattr(mod, "get_persist_dir", lambda: "/tmp/persist") @@ -246,6 +248,26 @@ async def fake_save(table_data_versions, _table_validations): return wf, activity_calls, child_data_wf_calls, save_calls, state +@pytest.mark.asyncio +@pytest.mark.parametrize('patched', [False, True]) +async def test_child_ingest_receives_the_snapshot_clock_once_patched( + monkeypatch, patched): + wf, _, child_data_wf_calls, _, _ = _setup_pipeline_test(monkeypatch) + wf._graph_version_info.timestamp_precision = 'milliseconds' + monkeypatch.setattr( + graph_refresh_workflow.workflow, 'patched', + lambda patch_id: patched and patch_id == 'diskgraph-index-clock-v1') + + await wf._run_diskgraph_colstats_pipeline(source_size_bytes_by_table={}, + source_size_bytes=0) + + child_args = inspect.signature( + diskgraph_ingest_workflow.DiskgraphIngestWorkflow.run).bind( + None, *child_data_wf_calls[0]['args']).arguments + assert child_args.get('timestamp_precision') == ('milliseconds' + if patched else None) + + @pytest.mark.asyncio @pytest.mark.parametrize( ('patched', 'expected_wf_id', 'expected_duplicate_counts'), diff --git a/test/rest/backends/test_snapshot_clock.py b/test/rest/backends/test_snapshot_clock.py new file mode 100644 index 000000000..fd8d7293d --- /dev/null +++ b/test/rest/backends/test_snapshot_clock.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from hashlib import md5 +from types import SimpleNamespace +from typing import Any, cast, get_args +from unittest import mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from kumoapi.typing import TimestampPrecision + +from io_utils.json import from_json_dict +from kumo.diskgraph.service import DiskgraphService +from rest.backends.data_backend.activities import \ + data_backend_activities as activities +from rest.config.data_backend import GraphDataVersionInfo +from rest.config.graph_v2 import GraphDataComputeInput +from rest.public.routers import data_snapshot +from rest.public.routers.data_snapshot import GraphSnapshotRouter + +CLOCKS = get_args(TimestampPrecision) + + +def _snapshot_version(monkeypatch, clock: str) -> GraphDataVersionInfo: + monkeypatch.setattr(activities.DataBackendActivities, + 'save_graph_data_version', mock.AsyncMock()) + monkeypatch.setattr( + activities, 'data_backend_cache', + lambda: SimpleNamespace(graph_cache=SimpleNamespace( + save_latest_graph_version=SimpleNamespace( + coroutine=mock.AsyncMock())))) + graph_input = GraphDataComputeInput( + 'graph', 'graph-key', cast(Any, {'table': None}), {}, {}, + cast(Any, None), table_aliases=cast(Any, {'table': 'events'}), + pinned_data_versions=cast(Any, {'table': 'v1'}), + timestamp_precision=clock) + return asyncio.run( + activities.DataBackendActivities._refresh_graph_data_version( + graph_input, {}, None)) + + +def test_seconds_snapshot_keeps_the_existing_identity(monkeypatch): + version = _snapshot_version(monkeypatch, 'seconds') + assert version.graph_data_version == md5(b'events=v1').hexdigest() + + +def test_each_clock_has_its_own_snapshot_identity(monkeypatch): + versions = [_snapshot_version(monkeypatch, clock) for clock in CLOCKS] + assert [version.timestamp_precision for version in versions] == list(CLOCKS) + assert len({version.graph_data_version for version in versions}) == 6 + + +def test_snapshots_recorded_before_clocks_are_seconds(): + version = from_json_dict( + { + 'graph_data_version': 'v1', + 'table_data_versions': {} + }, GraphDataVersionInfo) + assert version.timestamp_precision == 'seconds' + + +@pytest.fixture +def snapshot_api(monkeypatch): + driver = mock.AsyncMock() + driver.create_graph_snapshot.return_value = 'graph@v1' + recorded = SimpleNamespace(timestamp_precision='seconds') + monkeypatch.setattr(data_snapshot, 'get_graph_snapshot_info', + mock.AsyncMock(return_value=(None, recorded))) + monkeypatch.setattr(data_snapshot.GraphIDv2, 'from_public', + lambda graph_id: SimpleNamespace(fpid='graph')) + app = FastAPI() + app.include_router( + GraphSnapshotRouter(lambda: None, driver, mock.Mock()).router) + return SimpleNamespace(client=TestClient(app), driver=driver, + recorded=recorded) + + +def _create(snapshot_api, **params): + return snapshot_api.client.post('/graphsnapshots', params={ + 'graph_id': 'graph', + 'project_id': 'default', + **params + }) + + +def test_request_without_a_clock_builds_seconds(snapshot_api): + response = _create(snapshot_api) + kwargs = snapshot_api.driver.create_graph_snapshot.await_args.kwargs + assert kwargs['timestamp_precision'] == 'seconds' + assert response.headers['X-Kumo-Timestamp-Precision'] == 'seconds' + + +@pytest.mark.parametrize('clock', CLOCKS) +def test_snapshot_builds_the_requested_clock(snapshot_api, clock): + snapshot_api.recorded.timestamp_precision = clock + response = _create(snapshot_api, timestamp_precision=clock) + kwargs = snapshot_api.driver.create_graph_snapshot.await_args.kwargs + assert kwargs['timestamp_precision'] == clock + assert response.headers['X-Kumo-Timestamp-Precision'] == clock + + +def test_snapshot_echoes_the_recorded_clock_not_the_request(snapshot_api): + response = _create(snapshot_api, timestamp_precision='hours') + assert response.headers['X-Kumo-Timestamp-Precision'] == 'seconds' + + +def test_unknown_clock_is_rejected(snapshot_api): + assert _create(snapshot_api, + timestamp_precision='months').status_code == 422 + snapshot_api.driver.create_graph_snapshot.assert_not_awaited() + + +def test_prediction_statistics_use_the_snapshot_clock(monkeypatch): + import diskgraph + + binding = mock.Mock(return_value=('{}', '{}')) + monkeypatch.setattr(diskgraph, 'compute_col_stats_for_parquet', binding) + DiskgraphService(cast(Any, None)).compute_col_stats_for_parquet( + ['predictions.parquet'], [], 'v1', timestamp_precision='hours') + assert binding.call_args.kwargs['timestamp_unit'] == 'hours' diff --git a/test/train/datamodule/test_diskgraph_lightning_data.py b/test/train/datamodule/test_diskgraph_lightning_data.py index 9b8c68324..0e021b71b 100644 --- a/test/train/datamodule/test_diskgraph_lightning_data.py +++ b/test/train/datamodule/test_diskgraph_lightning_data.py @@ -3941,6 +3941,32 @@ def test_seed_time_to_index_unit_uses_manifest_unit(tmp_path) -> None: assert dm_ms._seed_time_to_index_unit(seed_ns) == 5_000 # 5e9 ns -> 5e3 ms +# 1 day, 1 hour, 1 minute, 1 second and 1 nanosecond after the epoch. +_SEED_NS = 90_061_000_000_001 + + +@pytest.mark.parametrize("clock,anchor", [ + ("nanoseconds", 90_061_000_000_001), + ("microseconds", 90_061_000_000), + ("milliseconds", 90_061_000), + ("seconds", 90_061), + ("hours", 25), + ("days", 1), +]) +def test_node_range_anchor_floors_into_the_index_clock(tmp_path, clock, + anchor) -> None: + index = tmp_path / "index" + index.mkdir() + (index / "index_metadata.json").write_text( + json.dumps({"timestamp_unit": clock})) + dm = _rhs_dm(_rhs_lp_manifest(tmp_path / "m", num_rhs=3, test_time_ns=1)) + dm._index_np_time_unit = dldm._index_np_time_unit( + {"graph_index_path": str(index)}) + + assert dm._seed_time_to_index_unit(torch.tensor([_SEED_NS])) == anchor + assert dm._seed_time_to_index_unit(torch.tensor([-1])) == -1 + + def test_train_dataloader_sets_fork_context_when_workers(tmp_path) -> None: """Python 3.14 workers retain fork to avoid pickling CUDA state.""" dm = DiskgraphLightningDataModule( diff --git a/test/train/test_shallow_rhs_link_gnn_module.py b/test/train/test_shallow_rhs_link_gnn_module.py index 30540f058..c15161146 100644 --- a/test/train/test_shallow_rhs_link_gnn_module.py +++ b/test/train/test_shallow_rhs_link_gnn_module.py @@ -5,7 +5,7 @@ from decimal import Decimal from pathlib import Path from types import SimpleNamespace -from unittest.mock import PropertyMock, patch +from unittest.mock import Mock, PropertyMock, patch import kumo_core.modeling.context as modeling_context import pandas as pd @@ -186,6 +186,22 @@ def test_ttng_rhs_temporal_validity_intersects_canonical_ids() -> None: assert valid.tolist() == [2, 4] +def test_temporal_batch_takes_valid_rhs_at_its_earliest_anchor() -> None: + model = _rhs_lookup_model(_ttng_rhs_index(['a', 'b', 'c'])) + lp_train_table = SimpleNamespace( + valid_rhs_idx_dict={}, + get_valid_rhs_idx=Mock(return_value=torch.tensor([0, 1, 2])), + ) + model._trainer = SimpleNamespace(datamodule=SimpleNamespace( + lp_train_table=lp_train_table)) + batch = HeteroData() + batch['LHS'].seed_time = torch.tensor([1_500, 1_250, 1_999]) + + model._get_valid_rhs_idx_for_batch(batch) + + lp_train_table.get_valid_rhs_idx.assert_called_once_with(1_250) + + def test_evaluation_end_restores_canonical_rhs_features_for_training() -> None: rhs_feats, rhs_encoders = get_tf_and_col_encoder_dict() model = ShallowRHSLinkGNNLightningModule(