From fcba4be6bd43852b24bf0d93978ed5174bbb5ed5 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:04 +0000 Subject: [PATCH 01/25] diskgraph-core: index clocks from nanoseconds to days Adds microsecond, millisecond, hour and day unit systems next to seconds and nanoseconds, the conversions main applies to timestamps (round up, floor, truncate; saturating when a value grows), a dispatch macro over the six index clocks, and 64-bit edge time keys. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph/diskgraph-core/src/metadata.rs | 1 + .../diskgraph/diskgraph-core/src/sortkey.rs | 34 ++ packages/diskgraph/diskgraph-core/src/time.rs | 355 +++++++++++++++++- .../diskgraph-index/src/bin/explore.rs | 2 + 4 files changed, 386 insertions(+), 6 deletions(-) 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..34c82a957 100644 --- a/packages/diskgraph/diskgraph-core/src/time.rs +++ b/packages/diskgraph/diskgraph-core/src/time.rs @@ -26,7 +26,36 @@ 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 fn tick_nanos(self) -> i64 { + TICK_NANOS[fixed_index(self)] + } + /// Parse a unit name string (used by builder methods that still accept strings). pub fn from_name(s: &str) -> Option { match s { @@ -45,6 +74,149 @@ 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. +// --------------------------------------------------------------------------- + +/// The tick in nanoseconds of each unit with a fixed length, by `fixed_index`. +const TICK_NANOS: [i64; 8] = [ + 1, + 1_000, + 1_000_000, + NS_PER_SEC, + 60 * NS_PER_SEC, + 3_600 * NS_PER_SEC, + 86_400 * NS_PER_SEC, + 604_800 * NS_PER_SEC, +]; + +/// Orders the fixed-length units from finest to coarsest. +fn fixed_index(unit: TimeUnit) -> usize { + match unit { + TimeUnit::Nanoseconds => 0, + TimeUnit::Microseconds => 1, + TimeUnit::Milliseconds => 2, + TimeUnit::Seconds => 3, + TimeUnit::Minutes => 4, + TimeUnit::Hours => 5, + TimeUnit::Days => 6, + TimeUnit::Weeks => 7, + other => panic!("unsupported unit for a fixed tick: {:?}", other), + } +} + +/// `FACTORS[i][j]` is how many ticks of the finer of units `i` and `j` make one +/// tick of the coarser, precomputed so that converting a value never divides +/// to find its factor. +const FACTORS: [[i64; 8]; 8] = { + let mut factors = [[0; 8]; 8]; + let mut i = 0; + while i < 8 { + let mut j = 0; + while j < 8 { + let (a, b) = (TICK_NANOS[i], TICK_NANOS[j]); + factors[i][j] = if a >= b { a / b } else { b / a }; + j += 1; + } + i += 1; + } + factors +}; + +enum Scale { + Same, + Up(i64), + Down(i64), +} + +#[inline] +fn scale(from: TimeUnit, to: TimeUnit) -> Scale { + if from == to { + return Scale::Same; + } + let (from, to) = (fixed_index(from), fixed_index(to)); + if from > to { + Scale::Up(FACTORS[from][to]) + } else { + Scale::Down(FACTORS[from][to]) + } +} + +/// Calls `round(value, divisor)` with a literal divisor for the steps between +/// sub-second units and seconds, so the compiler can turn the division into a +/// multiplication as it did when these conversions were hard-coded. +#[inline(always)] +fn with_literal_divisor(value: i64, divisor: i64, round: impl Fn(i64, i64) -> i64) -> i64 { + match divisor { + 1_000 => round(value, 1_000), + 1_000_000 => round(value, 1_000_000), + 1_000_000_000 => round(value, 1_000_000_000), + _ => round(value, divisor), + } +} + +pub fn convert_trunc(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { + match scale(from, to) { + Scale::Same => value, + Scale::Up(factor) => value.saturating_mul(factor), + Scale::Down(divisor) => with_literal_divisor(value, divisor, |v, d| v / d), + } +} + +pub fn convert_floor(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { + match scale(from, to) { + Scale::Same => value, + Scale::Up(factor) => value.saturating_mul(factor), + Scale::Down(divisor) => with_literal_divisor(value, divisor, i64::div_euclid), + } +} + +pub fn convert_ceil(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { + match scale(from, to) { + Scale::Same => value, + Scale::Up(factor) => value.saturating_mul(factor), + Scale::Down(divisor) => with_literal_divisor(value, divisor, |v, d| { + v.div_euclid(d) + i64::from(v.rem_euclid(d) != 0) + }), + } +} + +/// Floors `secs + nanos()` (with `0 <= nanos() < 1e9`) to a tick of `to`; +/// `nanos` is only evaluated when a tick is shorter than a second. +pub fn floor_parts(secs: i64, nanos: impl FnOnce() -> u32, to: TimeUnit) -> i64 { + let whole = convert_floor(secs, TimeUnit::Seconds, to); + if to.tick_nanos() < NS_PER_SEC { + whole.saturating_add(i64::from(nanos()) / to.tick_nanos()) + } else { + whole + } +} + +/// 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 +225,15 @@ 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; fn abs_from_secs(secs: i64) -> Abs; fn abs_from_millis(ms: i64) -> Abs; fn abs_from_micros(us: i64) -> Abs; @@ -243,13 +417,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 +472,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 +511,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. @@ -375,6 +638,7 @@ pub fn add_calendar_months(ts: Abs, months: i32) -> Abs { use chrono::{DateTime, Datelike, NaiveDate, Timelike}; let secs = ts.to_secs(); + let subsecond = ts - Abs::::from_secs(secs); let dt = DateTime::from_timestamp(secs, 0) .expect("timestamp out of range for calendar arithmetic") .naive_utc(); @@ -393,7 +657,7 @@ pub fn add_calendar_months(ts: Abs, months: i32) -> Abs { .and_hms_opt(dt.hour(), dt.minute(), dt.second()) .expect("invalid time in calendar arithmetic"); - Abs::::from_secs(result.and_utc().timestamp()) + Abs::::from_secs(result.and_utc().timestamp()) + subsecond } /// Add an offset to a timestamp. For fixed-size units (seconds through weeks) @@ -898,6 +1162,85 @@ 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, TimeUnit::Seconds), 10); + assert_eq!( + floor_parts(10, || 123_456_789, TimeUnit::Milliseconds), + 10_123 + ); + assert_eq!( + floor_parts(10, || 123_456_789, TimeUnit::Nanoseconds), + 10_123_456_789 + ); + assert_eq!(floor_parts(-1, || 500_000_000, TimeUnit::Hours), -1); + } + + #[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).raw(), 1_707_955_200_250); + } + #[test] fn time_unit_from_name_invalid() { assert_eq!(TimeUnit::from_name("invalid"), None); 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, } } From 31e0d2cbfad679484a8e36dcfa7398c0402081ca Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:19 +0000 Subject: [PATCH 02/25] diskgraph-rowpack: read timestamps in the index clock RowAccess carries the index clock and keeps main's parsing and null rules; only the final conversion now targets the clock instead of seconds. Typed timestamps and integer epochs round up, text fractions truncate, decimal epochs truncate to the finer of their unit and the tick and then round up, and Date32 multiplies. Seconds conversions are identities, so seconds rowpacks are unchanged. Decoders emit Arrow timestamps in the clock's unit (seconds for hours and days). Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/Cargo.lock | 1 + .../diskgraph/diskgraph-engine/src/engine.rs | 10 +- .../diskgraph/diskgraph-rowpack/Cargo.toml | 1 + .../diskgraph/diskgraph-rowpack/src/access.rs | 148 +++++++++++++++--- .../diskgraph/diskgraph-rowpack/src/lib.rs | 117 +++++++++++--- .../diskgraph/diskgraph-rowpack/src/parser.rs | 104 +++++++++--- .../src/tests_int_string_list.rs | 8 +- .../third_party/artifact-manifest.json | 2 +- 8 files changed, 321 insertions(+), 70 deletions(-) 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-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index fa4deee73..08d3ed1b2 100644 --- a/packages/diskgraph/diskgraph-engine/src/engine.rs +++ b/packages/diskgraph/diskgraph-engine/src/engine.rs @@ -2341,9 +2341,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); 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..a229a3dc3 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/access.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/access.rs @@ -12,6 +12,7 @@ use crate::{ use anyhow::{anyhow, Result}; use arrow_array::{Array, ArrayRef}; use arrow_schema::DataType; +use diskgraph_core::time::{convert_trunc, floor_parts, TimeUnit}; use num_traits::ToPrimitive; use std::sync::Arc; @@ -63,6 +64,8 @@ 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, } pub enum FloatListValues { @@ -778,9 +781,19 @@ impl<'p> RowAccess<'p> { hints, fast_parsers, tokenized_lists: vec![None; n], + 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 set_cols(&mut self, cols: &[ArrayRef]) { self.cols.clear(); self.cols.extend(cols.iter().map(plain_array)); @@ -2067,7 +2080,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, TimeUnit::Seconds)?; return Ok(Some(with_stack_str(secs, trimmed, f)?)); } DataType::Decimal128(_, scale) => { @@ -2392,7 +2405,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 = timestamp_from_int(value, unit, TimeUnit::Seconds); return Ok(Some(days_from_secs(secs))); } Err(anyhow!("date parse failed")) @@ -2483,11 +2496,11 @@ impl<'p> RowAccess<'p> { } _ => return Err(anyhow!("date type mismatch")), }; - let secs = timestamp_from_int(value, unit); + let secs = timestamp_from_int(value, unit, TimeUnit::Seconds); 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, TimeUnit::Seconds)?; Ok(Some(days_from_secs(secs))) } _ => Err(anyhow!("date type mismatch")), @@ -2495,6 +2508,10 @@ impl<'p> RowAccess<'p> { } pub fn get_timestamp(&self, idx: usize) -> Result> { + self.get_timestamp_in(idx, self.clock) + } + + pub fn get_timestamp_in(&self, idx: usize, clock: TimeUnit) -> Result> { let array = self.array(idx)?; if array.is_null(self.row) { return Ok(None); @@ -2504,13 +2521,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, clock)?)), DataType::Time32(_) | DataType::Time64(_) => { - let secs = time_secs_from_array(array, self.row)?; - Ok(Some(secs)) + Ok(Some(time_from_array(array, self.row, clock)?)) } DataType::Utf8 | DataType::LargeUtf8 @@ -2552,18 +2565,48 @@ 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) { + return Ok(Some(floor_parts(secs, || fraction_nanos(fraction), clock))); } } - 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(floor_parts(secs, || nanos, clock))); } if let Some(value) = parse_i64(text) { - return Ok(Some(timestamp_from_int(value, unit))); + return Ok(Some(timestamp_from_int(value, unit, clock))); } 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); + timestamp_from_decimal( + *scale, + unit, + clock, + |scale| apply_decimal_scale_i128(value, scale).to_i64(), + ) + .ok_or_else(|| anyhow!("decimal128 overflow")) + .map(Some) + } + DataType::Decimal256(_, scale) if format.is_none() => { + let value = array + .as_any() + .downcast_ref::() + .ok_or_else(|| anyhow!("decimal256 cast"))? + .value(self.row); + timestamp_from_decimal( + *scale, + unit, + clock, + |scale| apply_decimal_scale_i256(value, scale).to_i64(), + ) + .ok_or_else(|| anyhow!("decimal256 overflow")) + .map(Some) + } DataType::Int8 | DataType::Int16 | DataType::Int32 @@ -2655,11 +2698,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(floor_parts(secs, || nanos, clock))); } } - Ok(Some(timestamp_from_int(value, unit))) + Ok(Some(timestamp_from_int(value, unit, clock))) } DataType::Date32 => { let arr = array @@ -2675,9 +2718,7 @@ 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)) + Ok(Some(convert_trunc(i64::from(days), TimeUnit::Days, clock))) } _ => Err(anyhow!("timestamp type mismatch")), } @@ -3692,6 +3733,71 @@ 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 { + assert_eq!( + fractional_text.get_timestamp_in(0, clock)?, + Some(text), + "{clock:?}" + ); + assert_eq!( + epoch_millis.get_timestamp_in(0, clock)?, + Some(millis), + "{clock:?}" + ); + assert_eq!( + fractional_decimal.get_timestamp_in(0, clock)?, + 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..ff8e4590f 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::{convert_ceil, TimeUnit}; 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] @@ -861,14 +896,29 @@ fn div_ceil_i64(value: i64, divisor: i64) -> i64 { if r > 0 { d + 1 } else { d } } -pub(crate) fn timestamp_from_int(value: i64, unit: Option) -> i64 { +pub(crate) fn timestamp_from_int( + value: i64, + unit: Option, + clock: TimeUnit, +) -> 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), - } + convert_ceil(value, unit.into(), clock) +} + +/// A decimal epoch in `unit` as ticks of the clock `clock`, 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. +pub(crate) fn timestamp_from_decimal( + scale: i8, + unit: Option, + clock: TimeUnit, + scaled: impl FnOnce(i32) -> Option, +) -> Option { + let ticks_per_unit = timestamp_from_int(1, unit, clock); + if ticks_per_unit > 1 { + return scaled(i32::from(scale) - ticks_per_unit.ilog10() as i32); + } + Some(timestamp_from_int(scaled(i32::from(scale))?, unit, clock)) } pub(crate) fn days_from_secs(secs: i64) -> i32 { @@ -942,7 +992,7 @@ pub(crate) fn parse_decimal_i256(value: &str, scale: i32) -> Option Result { +pub(crate) fn timestamp_from_array(array: &ArrayRef, row: usize, clock: TimeUnit) -> Result { let unit = match array.data_type() { DataType::Timestamp(unit, _) => *unit, _ => return Err(anyhow!("timestamp type mismatch")), @@ -953,61 +1003,69 @@ 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)) + Ok(convert_ceil(arr.value(row), TimeUnit::Seconds, clock)) } 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)) + Ok(convert_ceil(arr.value(row), TimeUnit::Milliseconds, clock)) } 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)) + Ok(convert_ceil(arr.value(row), TimeUnit::Microseconds, clock)) } 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)) + Ok(convert_ceil(arr.value(row), TimeUnit::Nanoseconds, clock)) } } } -pub(crate) fn time_secs_from_array(array: &ArrayRef, row: usize) -> Result { +pub(crate) fn time_from_array(array: &ArrayRef, row: usize, clock: TimeUnit) -> 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))) + Ok(convert_ceil( + i64::from(arr.value(row)), + TimeUnit::Seconds, + clock, + )) } 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)) + Ok(convert_ceil( + i64::from(arr.value(row)), + TimeUnit::Milliseconds, + clock, + )) } 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)) + Ok(convert_ceil(arr.value(row), TimeUnit::Microseconds, clock)) } 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)) + Ok(convert_ceil(arr.value(row), TimeUnit::Nanoseconds, clock)) } 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/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 } From 5ad2dcbe5b6ecd1a956385c25f8ebab0776ce00b Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:36 +0000 Subject: [PATCH 03/25] diskgraph-ingest: ingest into the configured index clock The ingest config gains a top-level timestamp_unit that defaults to seconds when absent. The saved diskgraph_config.json and index_metadata.json always record it, and every row access reads it. An Encode marker written before the clock was recorded still resumes, read as seconds. Edge time keys stay 32-bit for seconds, keep their configured width for hours and days, and use the full 64-bit value for sub-second clocks; out-of-range edge times still clamp to the far-future key. TT-NG timestamp grouping stays one second wide, and stage extensions inherit the parent's clock. Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-infer/src/lib.rs | 1 + .../src/config_converter/convert.rs | 1 + .../src/config_converter/tests.rs | 7 ++ .../diskgraph-ingest/src/graph_config.rs | 42 +++++++++ .../diskgraph/diskgraph-ingest/src/lib.rs | 90 ++++++++++++++----- .../diskgraph-ingest/src/stage_extension.rs | 5 ++ .../diskgraph-ingest/src/tests_unit.rs | 40 +++++++++ .../diskgraph-ingest/tests/graph_config.rs | 4 + .../tests/resolution_ingest.rs | 3 + .../tests/resolution_runtime_batch.rs | 3 + .../tests/seed_table_ingest.rs | 6 ++ .../src/tt_executor/pseudo_time_tests.rs | 1 + .../diskgraph-pybind/src/tt_executor/tests.rs | 1 + .../tests/custom_tt_integration.rs | 2 + .../diskgraph-train-table/tests/ttgen_e2e.rs | 2 + .../tests/ttgen_integration.rs | 13 +++ 16 files changed, 199 insertions(+), 22 deletions(-) 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/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/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..a856d92e4 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, 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)) } @@ -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) + .get_timestamp_in(idx, TimeUnit::Seconds) .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/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-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 { From b862415640ff5d70a6c207dbeb75ce4ef88b2783 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:36 +0000 Subject: [PATCH 04/25] diskgraph: read and sample in the index clock Decoders, temporal validity checks and the sampler take the clock from index metadata. Validity columns keep main's conversions (an Int64 column is epoch nanoseconds, a Date32 whole days), seed timestamps still truncate into the index clock, and decoded seed TIMESTAMP columns are read in their own Arrow unit. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph/diskgraph-engine/src/engine.rs | 53 +++++++++++-------- .../diskgraph/diskgraph-index/src/graph.rs | 33 +++++++----- 2 files changed, 52 insertions(+), 34 deletions(-) diff --git a/packages/diskgraph/diskgraph-engine/src/engine.rs b/packages/diskgraph/diskgraph-engine/src/engine.rs index 08d3ed1b2..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); } } @@ -3388,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/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, } From 50e74023efce73a9cbebd33cce134d8709d8e7a8 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:45 +0000 Subject: [PATCH 05/25] diskgraph: generate training tables in the index clock Plan execution dispatches once on the index clock. The builder keeps plan times in nanoseconds and writes the plan in the clock it was created with, converting with main's truncation; DiskGraphPlanCompiler(index_path=...) uses the index's clock, and a non-seconds plan carries a timestamp_unit tag that must match the index. Datetime literals still round up and custom-table timestamps still truncate, now to the index clock. Outputs stay timestamp[ns], saturating, and random splits keep hashing whole seconds. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph-core/src/tt_builder.rs | 102 ++++++++++- .../python/diskgraph/compiler.py | 28 ++- .../python/diskgraph/tt_artifacts.py | 10 +- .../diskgraph-pybind/src/custom_tt_entry.rs | 21 ++- .../diskgraph-pybind/src/tt_builder.rs | 171 ++++++++++++------ .../diskgraph-pybind/src/tt_executor.rs | 85 ++++++++- .../src/custom_tt_chunk_processor.rs | 112 +++++++----- .../src/split_decider.rs | 18 +- .../src/ttgen_executor.rs | 13 +- .../diskgraph-train-table/src/ttgen_output.rs | 4 +- .../diskgraph-train-table/src/ttgen_stats.rs | 8 +- .../src/ttgen_traversal.rs | 16 +- 12 files changed, 424 insertions(+), 164 deletions(-) diff --git a/packages/diskgraph/diskgraph-core/src/tt_builder.rs b/packages/diskgraph/diskgraph-core/src/tt_builder.rs index ad317ca83..4d2ba80cd 100644 --- a/packages/diskgraph/diskgraph-core/src/tt_builder.rs +++ b/packages/diskgraph/diskgraph-core/src/tt_builder.rs @@ -51,7 +51,7 @@ pub const ENTITY_NODE_ID_COL: &str = "entity_node_id"; /// Resolved LP RHS positives (`list` node-ids). pub const TARGET_NODE_IDS_COL: &str = "target_node_ids"; -use crate::time::{Abs, Dur, UnitSystem}; +use crate::time::{convert_trunc, Abs, Dur, UnitSystem}; use serde::{Deserialize, Serialize}; // ─── Top-level plan ────────────────────────────────────────────────────────── @@ -537,6 +537,106 @@ 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). Converting between clocks truncates +// anchors, split bounds and offsets, as `main` does when it builds a plan. + +impl DiskGraphTTPlan { + pub fn to_clock(&self) -> serde_json::Result> { + let mut plan: DiskGraphTTPlan = serde_json::from_value(serde_json::to_value(self)?)?; + if let Some(time) = &mut plan.entity.default_time_value { + rescale_abs(time, U::UNIT); + } + if let Some(temporal) = &mut plan.temporal { + for time in &mut temporal.timestamps { + rescale_abs(time, U::UNIT); + } + } + if let Some(split) = &mut plan.split { + split.rescale_from(U::UNIT); + } + for stage in &mut plan.stages { + for op in &mut stage.ops { + op.rescale_from(U::UNIT); + } + } + Ok(plan) + } + + 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) + } +} + +/// 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) +} + +impl SplitSpec { + pub fn to_clock(&self) -> serde_json::Result> { + let mut split: SplitSpec = serde_json::from_value(serde_json::to_value(self)?)?; + split.rescale_from(U::UNIT); + Ok(split) + } + + fn rescale_from(&mut self, from: TimeUnit) { + for offset in self.offsets.iter_mut().flatten() { + rescale_dur(offset, from); + } + for (start, end) in self.ranges.iter_mut().flatten() { + rescale_abs(start, from); + rescale_abs(end, from); + } + if let Some(time) = &mut self.backtesting_timestamp { + rescale_abs(time, from); + } + } +} + +impl Op { + fn rescale_from(&mut self, from: TimeUnit) { + match self { + Op::Autoregression { + target_op, + lag_offsets: offsets, + .. + } + | Op::FutureLabels { + target_op, + future_offsets: offsets, + .. + } => { + target_op.rescale_from(from); + for offset in offsets { + rescale_dur(offset, from); + } + } + _ => {} + } + } +} + +fn rescale_abs(time: &mut Abs, from: TimeUnit) { + *time = Abs::from_raw(convert_trunc(time.raw(), from, U::UNIT)); +} + +fn rescale_dur(duration: &mut Dur, from: TimeUnit) { + *duration = Dur::from_raw(convert_trunc(duration.raw(), from, U::UNIT)); +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- 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/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/tt_builder.rs b/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs index 2b0f97e1f..219b3f85f 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, Nanoseconds, + TimeUnit, +}; 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,15 +331,18 @@ impl StageFilter { // ─── DiskGraphTTBuilder ────────────────────────────────────────────────────── +/// Holds plan times in nanoseconds; `build()` writes the plan in `timestamp_unit`. #[pyclass] pub struct DiskGraphTTBuilder { - entity: Option>, - temporal: Option>, + timestamp_unit: TimeUnit, + entity: Option>, + default_time_secs: Option, + temporal: Option>, prediction_type: tt::PredictionType, is_prediction: bool, - split: Option>, - stages: Vec>, - current_ops: Vec>, + split: Option>, + stages: Vec>, + current_ops: Vec>, target_stype: Option, fkey_mapping: Option>, } @@ -326,9 +350,12 @@ pub struct DiskGraphTTBuilder { #[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, + default_time_secs: None, temporal: None, prediction_type: tt::PredictionType::Node, is_prediction: false, @@ -337,10 +364,10 @@ impl DiskGraphTTBuilder { current_ops: Vec::new(), target_stype: None, fkey_mapping: None, - } + }) } - /// `FOR EACH {table}.{column}`. + /// `FOR EACH {table}.{column}`; `default_time_value` is in epoch seconds. #[pyo3(signature = (table, column, output_time_col_name=None, default_time_value=None))] fn for_each( &mut self, @@ -353,11 +380,12 @@ impl DiskGraphTTBuilder { table, column, output_time_col_name, - default_time_value: default_time_value.map(Abs::from_raw), + default_time_value: None, time_col: None, valid_to_col: None, filter: None, }); + self.default_time_secs = default_time_value; } /// Entity-level create-time column for temporal filtering. @@ -414,7 +442,7 @@ impl DiskGraphTTBuilder { self.temporal = Some(tt::TemporalSpec { timestamps: timestamps .into_iter() - .map(|v| abs_from_unit::(v, time_unit)) + .map(|v| abs_from_unit::(v, time_unit)) .collect(), forecast_length, lag_timesteps, @@ -547,7 +575,7 @@ impl DiskGraphTTBuilder { alias, lhs: lhs_alias, op: parse_compare_op(op)?, - rhs: extract_value(value)?, + rhs: extract_value(value, self.timestamp_unit)?, }); Ok(()) } @@ -583,9 +611,10 @@ 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 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 { @@ -593,7 +622,7 @@ impl DiskGraphTTBuilder { target_op: Box::new(target_op), lag_offsets: lag_offsets .into_iter() - .map(|v| dur_from_unit::(v, unit)) + .map(|v| dur_from_unit::(v, unit)) .collect(), }); Ok(()) @@ -607,9 +636,10 @@ 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 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 { @@ -617,7 +647,7 @@ impl DiskGraphTTBuilder { target_op: Box::new(target_op), future_offsets: future_offsets .into_iter() - .map(|v| dur_from_unit::(v, unit)) + .map(|v| dur_from_unit::(v, unit)) .collect(), }); Ok(()) @@ -678,7 +708,7 @@ impl DiskGraphTTBuilder { crate::DiskgraphConfigError::new_err("call for_each() before build()") })?; - let plan: tt::DiskGraphTTPlan = tt::DiskGraphTTPlan { + let plan: tt::DiskGraphTTPlan = tt::DiskGraphTTPlan { entity, temporal: self.temporal.take(), prediction_type: std::mem::replace(&mut self.prediction_type, tt::PredictionType::Node), @@ -697,8 +727,14 @@ impl DiskGraphTTBuilder { fact_pseudo_time: None, }; - serde_json::to_string(&plan) - .map_err(|e| crate::DiskgraphConfigError::new_err(e.to_string())) + let default_time_secs = self.default_time_secs.take(); + with_index_clock!(self.timestamp_unit, U => { + plan.to_clock::().and_then(|mut plan| { + plan.entity.default_time_value = default_time_secs.map(Abs::from_secs); + plan.to_json() + }) + }) + .map_err(|e| crate::DiskgraphConfigError::new_err(e.to_string())) } } @@ -835,7 +871,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 +883,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 +958,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 +987,14 @@ 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 +1016,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 +1026,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 split = make_split_spec(extract_split_spec(spec)?)?; + with_index_clock!(parse_index_clock(timestamp_unit)?, U => { + split.to_clock::().and_then(|split| serde_json::to_string(&split)) + }) + .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 +1088,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 +1098,9 @@ 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)?; + 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 +1116,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 +1129,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 +1139,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 { 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-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..9ac8ff71e 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 / @@ -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)] diff --git a/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs index 5a4467ddd..0ef933a6e 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_executor.rs @@ -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 @@ -1941,7 +1941,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 { @@ -2702,7 +2702,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,10 +2904,7 @@ 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. + // encode_upper expects index-clock ticks, which `U` is. let window = has_edges.window.as_ref().map(|w| { let obs = observation_ts.unwrap_or(Abs::::from_raw(0)); match w { 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_stats.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs index d5eda42b0..69196bde7 100644 --- a/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs +++ b/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs @@ -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,7 +512,7 @@ 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(); @@ -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"; 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 } From 5b2378ce9f37db8ac8523747d7e25bc67e3ac4a6 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:53 +0000 Subject: [PATCH 06/25] diskgraph: timestamp statistics in the index clock Timestamp column statistics hold index-clock ticks. col_stats.json keeps min_secs/max_secs for seconds and otherwise names the bounds after the clock (min_ms, min_hours, ...) and adds timestamp_unit. resolve_temporal and compute_col_stats_for_parquet take an optional timestamp_unit, default seconds. Date statistics and batch-prediction histograms are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph-ingest/src/col_stats.rs | 129 ++++++++++----- .../src/config_converter/inference.rs | 15 +- .../diskgraph/diskgraph-ingest/src/lib.rs | 2 +- .../diskgraph-pybind/src/bp_col_stats.rs | 155 +++++++++++++----- .../diskgraph/diskgraph-pybind/src/lib.rs | 8 +- .../diskgraph-pybind/src/temporal.rs | 20 ++- 6 files changed, 236 insertions(+), 93 deletions(-) diff --git a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs index b8bbd94b9..0aca584ed 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,12 +1583,26 @@ 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(); } + let (epoch_secs, _) = split_parts(ticks, clock); match chrono::DateTime::from_timestamp(epoch_secs, 0) { Some(dt) => format!("kumo_timestamp_{}", dt.format("%Y-%m-%d %H:%M:%S")), 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] 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/lib.rs b/packages/diskgraph/diskgraph-ingest/src/lib.rs index a856d92e4..1dd2ad0fa 100644 --- a/packages/diskgraph/diskgraph-ingest/src/lib.rs +++ b/packages/diskgraph/diskgraph-ingest/src/lib.rs @@ -7899,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 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/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, From 44ce189b212adf946e1917b69ca544b5c953a132 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 20:29:53 +0000 Subject: [PATCH 07/25] diskgraph: test that every index clock gives main's results On whole-day data every clock decodes, samples, generates the training table and computes statistics exactly like the seconds index. On a messy fixture each source type follows main's null and rounding rule in every clock. Co-Authored-By: Claude Opus 5.5 --- .../python/tests/test_index_clocks.py | 288 ++++++++++++++++++ 1 file changed, 288 insertions(+) create mode 100644 packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py 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..74b8d40b9 --- /dev/null +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py @@ -0,0 +1,288 @@ +# 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) -> 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, 7), 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 From 592456aef9d8b030d951f86f0e7cbd85e2832127 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:29:43 +0000 Subject: [PATCH 08/25] diskgraph: convert timestamps with clock-constant factors The row reader picks its generic timestamp_in:: once, and the ingest conversions are generic over the clock with one arm per source unit, so every factor is a compile-time constant. On the seconds clock this compiles to main's hard-coded divisions and the text fraction is never parsed. Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-core/src/time.rs | 188 ++++++++---------- .../diskgraph/diskgraph-ingest/src/lib.rs | 4 +- .../diskgraph/diskgraph-rowpack/src/access.rs | 65 +++--- .../diskgraph/diskgraph-rowpack/src/parser.rs | 61 +++--- 4 files changed, 148 insertions(+), 170 deletions(-) diff --git a/packages/diskgraph/diskgraph-core/src/time.rs b/packages/diskgraph/diskgraph-core/src/time.rs index 34c82a957..8ccdf141b 100644 --- a/packages/diskgraph/diskgraph-core/src/time.rs +++ b/packages/diskgraph/diskgraph-core/src/time.rs @@ -52,8 +52,18 @@ impl TimeUnit { Self::INDEX_CLOCKS.contains(&self) } - pub fn tick_nanos(self) -> i64 { - TICK_NANOS[fixed_index(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). @@ -79,117 +89,75 @@ impl TimeUnit { // `main` applies to some kind of value; coarse-to-fine conversions saturate. // --------------------------------------------------------------------------- -/// The tick in nanoseconds of each unit with a fixed length, by `fixed_index`. -const TICK_NANOS: [i64; 8] = [ - 1, - 1_000, - 1_000_000, - NS_PER_SEC, - 60 * NS_PER_SEC, - 3_600 * NS_PER_SEC, - 86_400 * NS_PER_SEC, - 604_800 * NS_PER_SEC, -]; - -/// Orders the fixed-length units from finest to coarsest. -fn fixed_index(unit: TimeUnit) -> usize { - match unit { - TimeUnit::Nanoseconds => 0, - TimeUnit::Microseconds => 1, - TimeUnit::Milliseconds => 2, - TimeUnit::Seconds => 3, - TimeUnit::Minutes => 4, - TimeUnit::Hours => 5, - TimeUnit::Days => 6, - TimeUnit::Weeks => 7, - other => panic!("unsupported unit for a fixed tick: {:?}", other), +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) } } -/// `FACTORS[i][j]` is how many ticks of the finer of units `i` and `j` make one -/// tick of the coarser, precomputed so that converting a value never divides -/// to find its factor. -const FACTORS: [[i64; 8]; 8] = { - let mut factors = [[0; 8]; 8]; - let mut i = 0; - while i < 8 { - let mut j = 0; - while j < 8 { - let (a, b) = (TICK_NANOS[i], TICK_NANOS[j]); - factors[i][j] = if a >= b { a / b } else { b / a }; - j += 1; - } - i += 1; +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) } - factors -}; - -enum Scale { - Same, - Up(i64), - Down(i64), } -#[inline] -fn scale(from: TimeUnit, to: TimeUnit) -> Scale { - if from == to { - return Scale::Same; - } - let (from, to) = (fixed_index(from), fixed_index(to)); - if from > to { - Scale::Up(FACTORS[from][to]) +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 { - Scale::Down(FACTORS[from][to]) + div_ceil(value, to / from) } } -/// Calls `round(value, divisor)` with a literal divisor for the steps between -/// sub-second units and seconds, so the compiler can turn the division into a -/// multiplication as it did when these conversions were hard-coded. -#[inline(always)] -fn with_literal_divisor(value: i64, divisor: i64, round: impl Fn(i64, i64) -> i64) -> i64 { - match divisor { - 1_000 => round(value, 1_000), - 1_000_000 => round(value, 1_000_000), - 1_000_000_000 => round(value, 1_000_000_000), - _ => round(value, divisor), - } +#[inline] +pub fn div_ceil(value: i64, divisor: i64) -> i64 { + value.div_euclid(divisor) + i64::from(value.rem_euclid(divisor) != 0) } -pub fn convert_trunc(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { - match scale(from, to) { - Scale::Same => value, - Scale::Up(factor) => value.saturating_mul(factor), - Scale::Down(divisor) => with_literal_divisor(value, divisor, |v, d| v / d), - } -} +// 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). -pub fn convert_floor(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { - match scale(from, to) { - Scale::Same => value, - Scale::Up(factor) => value.saturating_mul(factor), - Scale::Down(divisor) => with_literal_divisor(value, divisor, i64::div_euclid), +/// Converts `value` in `unit` to ticks of the clock `U`, rounding up. +#[inline(always)] +pub fn ceil_to_clock(value: i64, unit: TimeUnit) -> i64 { + 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()), } } -pub fn convert_ceil(value: i64, from: TimeUnit, to: TimeUnit) -> i64 { - match scale(from, to) { - Scale::Same => value, - Scale::Up(factor) => value.saturating_mul(factor), - Scale::Down(divisor) => with_literal_divisor(value, divisor, |v, d| { - v.div_euclid(d) + i64::from(v.rem_euclid(d) != 0) - }), +#[inline(always)] +fn ceil_from_tick(value: i64, tick: i64) -> i64 { + let clock_tick = U::TICK_NANOS; + if tick >= clock_tick { + value.saturating_mul(tick / clock_tick) + } else { + div_ceil(value, clock_tick / tick) } } -/// Floors `secs + nanos()` (with `0 <= nanos() < 1e9`) to a tick of `to`; -/// `nanos` is only evaluated when a tick is shorter than a second. -pub fn floor_parts(secs: i64, nanos: impl FnOnce() -> u32, to: TimeUnit) -> i64 { - let whole = convert_floor(secs, TimeUnit::Seconds, to); - if to.tick_nanos() < NS_PER_SEC { - whole.saturating_add(i64::from(nanos()) / to.tick_nanos()) +/// Floors `secs + nanos()` (with `0 <= nanos() < 1e9`) to a tick of the clock +/// `U`; `nanos` is only evaluated when a tick is shorter than a second. +#[inline(always)] +pub fn floor_parts(secs: i64, nanos: impl FnOnce() -> u32) -> i64 { + let clock_tick = U::TICK_NANOS; + if clock_tick < NS_PER_SEC { + let fraction = i64::from(nanos()) / clock_tick; + secs.saturating_mul(NS_PER_SEC / clock_tick) + .saturating_add(fraction) } else { - whole + secs.div_euclid(clock_tick / NS_PER_SEC) } } @@ -234,6 +202,7 @@ pub fn split_parts(value: i64, unit: TimeUnit) -> (i64, u32) { /// 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; @@ -1194,16 +1163,35 @@ mod tests { #[test] fn floor_parts_drops_the_fraction_below_one_tick() { - assert_eq!(floor_parts(10, || 999_999_999, TimeUnit::Seconds), 10); + assert_eq!(floor_parts::(10, || 999_999_999), 10); + assert_eq!(floor_parts::(10, || 123_456_789), 10_123); assert_eq!( - floor_parts(10, || 123_456_789, TimeUnit::Milliseconds), - 10_123 - ); - assert_eq!( - floor_parts(10, || 123_456_789, TimeUnit::Nanoseconds), + floor_parts::(10, || 123_456_789), 10_123_456_789 ); - assert_eq!(floor_parts(-1, || 500_000_000, TimeUnit::Hours), -1); + assert_eq!(floor_parts::(-1, || 500_000_000), -1); + } + + #[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)); + assert_eq!( + ticks, + convert_ceil(value, unit, clock), + "{unit:?} -> {clock:?}" + ); + } + } + } } #[test] diff --git a/packages/diskgraph/diskgraph-ingest/src/lib.rs b/packages/diskgraph/diskgraph-ingest/src/lib.rs index 1dd2ad0fa..bbc6e467d 100644 --- a/packages/diskgraph/diskgraph-ingest/src/lib.rs +++ b/packages/diskgraph/diskgraph-ingest/src/lib.rs @@ -83,7 +83,7 @@ use diskgraph_core::metadata::{ EdgeSortKeyWeightEncoding, EdgeSortKeyWeightKind, EdgeSortKeyWeightMetadata, EdgeStorageFormat, NodeDataCompression, NodeShardMetadata, NodeTypeMetadata, PrimaryKeyMap, }; -use diskgraph_core::time::{convert_floor, convert_trunc, TimeUnit}; +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, @@ -8897,7 +8897,7 @@ impl IngestContext { Some(idx) => Some( // Buckets stay one second wide whatever the index clock. access - .get_timestamp_in(idx, TimeUnit::Seconds) + .timestamp_in::(idx) .with_context(|| { format!( "TT-NG ts-grouping: failed to decode TIMESTAMP for table `{}` row {}", diff --git a/packages/diskgraph/diskgraph-rowpack/src/access.rs b/packages/diskgraph/diskgraph-rowpack/src/access.rs index a229a3dc3..459964007 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/access.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/access.rs @@ -12,7 +12,8 @@ use crate::{ use anyhow::{anyhow, Result}; use arrow_array::{Array, ArrayRef}; use arrow_schema::DataType; -use diskgraph_core::time::{convert_trunc, floor_parts, TimeUnit}; +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; @@ -66,8 +67,12 @@ pub struct RowAccess<'p> { 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), @@ -782,11 +787,14 @@ impl<'p> RowAccess<'p> { 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 } @@ -2080,7 +2088,7 @@ impl<'p> RowAccess<'p> { return Ok(Some(with_stack_str(value, trimmed, f)?)); } DataType::Timestamp(_, _) => { - let secs = timestamp_from_array(array, self.row, TimeUnit::Seconds)?; + let secs = timestamp_from_array::(array, self.row)?; return Ok(Some(with_stack_str(secs, trimmed, f)?)); } DataType::Decimal128(_, scale) => { @@ -2405,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, TimeUnit::Seconds); + let secs = timestamp_from_int::(value, unit); return Ok(Some(days_from_secs(secs))); } Err(anyhow!("date parse failed")) @@ -2496,11 +2504,11 @@ impl<'p> RowAccess<'p> { } _ => return Err(anyhow!("date type mismatch")), }; - let secs = timestamp_from_int(value, unit, TimeUnit::Seconds); + let secs = timestamp_from_int::(value, unit); Ok(Some(days_from_secs(secs))) } DataType::Timestamp(_, _) => { - let secs = timestamp_from_array(array, self.row, TimeUnit::Seconds)?; + let secs = timestamp_from_array::(array, self.row)?; Ok(Some(days_from_secs(secs))) } _ => Err(anyhow!("date type mismatch")), @@ -2508,10 +2516,10 @@ impl<'p> RowAccess<'p> { } pub fn get_timestamp(&self, idx: usize) -> Result> { - self.get_timestamp_in(idx, self.clock) + (self.read_timestamp)(self, idx) } - pub fn get_timestamp_in(&self, idx: usize, clock: TimeUnit) -> Result> { + pub fn timestamp_in(&self, idx: usize) -> Result> { let array = self.array(idx)?; if array.is_null(self.row) { return Ok(None); @@ -2521,9 +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(_, _) => Ok(Some(timestamp_from_array(array, self.row, clock)?)), + DataType::Timestamp(_, _) => Ok(Some(timestamp_from_array::(array, self.row)?)), DataType::Time32(_) | DataType::Time64(_) => { - Ok(Some(time_from_array(array, self.row, clock)?)) + Ok(Some(time_from_array::(array, self.row)?)) } DataType::Utf8 | DataType::LargeUtf8 @@ -2566,14 +2574,14 @@ impl<'p> RowAccess<'p> { }; if let Some(parser) = fast { if let Some((secs, fraction)) = parser.parse_timestamp_parts(text) { - return Ok(Some(floor_parts(secs, || fraction_nanos(fraction), clock))); + return Ok(Some(floor_parts::(secs, || fraction_nanos(fraction)))); } } if let Some((secs, nanos)) = parse_timestamp_text_parts(text, format) { - return Ok(Some(floor_parts(secs, || nanos, clock))); + return Ok(Some(floor_parts::(secs, || nanos))); } if let Some(value) = parse_i64(text) { - return Ok(Some(timestamp_from_int(value, unit, clock))); + return Ok(Some(timestamp_from_int::(value, unit))); } Err(anyhow!("timestamp parse failed")) } @@ -2583,10 +2591,9 @@ impl<'p> RowAccess<'p> { .downcast_ref::() .ok_or_else(|| anyhow!("decimal128 cast"))? .value(self.row); - timestamp_from_decimal( + timestamp_from_decimal::( *scale, unit, - clock, |scale| apply_decimal_scale_i128(value, scale).to_i64(), ) .ok_or_else(|| anyhow!("decimal128 overflow")) @@ -2598,10 +2605,9 @@ impl<'p> RowAccess<'p> { .downcast_ref::() .ok_or_else(|| anyhow!("decimal256 cast"))? .value(self.row); - timestamp_from_decimal( + timestamp_from_decimal::( *scale, unit, - clock, |scale| apply_decimal_scale_i256(value, scale).to_i64(), ) .ok_or_else(|| anyhow!("decimal256 overflow")) @@ -2699,10 +2705,10 @@ impl<'p> RowAccess<'p> { if let Some(fmt) = format { let text = value.to_string(); if let Some((secs, nanos)) = parse_timestamp_text_parts(&text, Some(fmt)) { - return Ok(Some(floor_parts(secs, || nanos, clock))); + return Ok(Some(floor_parts::(secs, || nanos))); } } - Ok(Some(timestamp_from_int(value, unit, clock))) + Ok(Some(timestamp_from_int::(value, unit))) } DataType::Date32 => { let arr = array @@ -2718,7 +2724,7 @@ impl<'p> RowAccess<'p> { if !(-106752..=24855).contains(&days) { return Err(anyhow!("date out of range: day {days}")); } - Ok(Some(convert_trunc(i64::from(days), TimeUnit::Days, clock))) + Ok(Some(ceil_to_clock::(i64::from(days), TimeUnit::Days))) } _ => Err(anyhow!("timestamp type mismatch")), } @@ -3779,21 +3785,12 @@ mod tests { (Clock::Days, 19_737, 19_738, 19_738), ]; for (clock, text, millis, decimal) in expected { - assert_eq!( - fractional_text.get_timestamp_in(0, clock)?, - Some(text), - "{clock:?}" - ); - assert_eq!( - epoch_millis.get_timestamp_in(0, clock)?, - Some(millis), - "{clock:?}" - ); - assert_eq!( - fractional_decimal.get_timestamp_in(0, clock)?, - Some(decimal), - "{clock:?}" - ); + 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(()) } diff --git a/packages/diskgraph/diskgraph-rowpack/src/parser.rs b/packages/diskgraph/diskgraph-rowpack/src/parser.rs index ff8e4590f..33a1d8ea3 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/parser.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/parser.rs @@ -7,7 +7,7 @@ use anyhow::{anyhow, Result}; use arrow_array::{Array, ArrayRef}; use arrow_schema::DataType; use chrono::{NaiveDate, NaiveDateTime}; -use diskgraph_core::time::{convert_ceil, TimeUnit}; +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 { @@ -885,44 +885,37 @@ 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 } -} - -pub(crate) fn timestamp_from_int( +#[inline(always)] +pub(crate) fn timestamp_from_int( value: i64, unit: Option, - clock: TimeUnit, ) -> i64 { - let unit = unit.unwrap_or(arrow_schema::TimeUnit::Second); - convert_ceil(value, unit.into(), clock) + 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 `clock`, where `scaled(s)` is +/// 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. -pub(crate) fn timestamp_from_decimal( +#[inline(always)] +pub(crate) fn timestamp_from_decimal( scale: i8, unit: Option, - clock: TimeUnit, scaled: impl FnOnce(i32) -> Option, ) -> Option { - let ticks_per_unit = timestamp_from_int(1, unit, clock); + 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); } - Some(timestamp_from_int(scaled(i32::from(scale))?, unit, clock)) + Some(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 { @@ -992,7 +985,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")), @@ -1003,43 +997,43 @@ pub(crate) fn timestamp_from_array(array: &ArrayRef, row: usize, clock: TimeUnit .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("timestamp second cast"))?; - Ok(convert_ceil(arr.value(row), TimeUnit::Seconds, clock)) + Ok(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(convert_ceil(arr.value(row), TimeUnit::Milliseconds, clock)) + Ok(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(convert_ceil(arr.value(row), TimeUnit::Microseconds, clock)) + Ok(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(convert_ceil(arr.value(row), TimeUnit::Nanoseconds, clock)) + Ok(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) } } } -pub(crate) fn time_from_array(array: &ArrayRef, row: usize, clock: TimeUnit) -> 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(convert_ceil( + Ok(ceil_to_clock::( i64::from(arr.value(row)), TimeUnit::Seconds, - clock, )) } DataType::Time32(arrow_schema::TimeUnit::Millisecond) => { @@ -1047,10 +1041,9 @@ pub(crate) fn time_from_array(array: &ArrayRef, row: usize, clock: TimeUnit) -> .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time32 millisecond cast"))?; - Ok(convert_ceil( + Ok(ceil_to_clock::( i64::from(arr.value(row)), TimeUnit::Milliseconds, - clock, )) } DataType::Time64(arrow_schema::TimeUnit::Microsecond) => { @@ -1058,14 +1051,14 @@ pub(crate) fn time_from_array(array: &ArrayRef, row: usize, clock: TimeUnit) -> .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time64 microsecond cast"))?; - Ok(convert_ceil(arr.value(row), TimeUnit::Microseconds, clock)) + Ok(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(convert_ceil(arr.value(row), TimeUnit::Nanoseconds, clock)) + Ok(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) } DataType::Time32(_) | DataType::Time64(_) => Err(anyhow!( "unsupported Arrow time unit for timestamp conversion: {:?}", From 6bf9e517106195ac8e9824ebb7fea50a0be61d7e Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:33:36 +0000 Subject: [PATCH 09/25] diskgraph-rowpack: null timestamps the index clock cannot hold A source value beyond the range of the index clock (for example year 9999 on a nanosecond index) now reads as an error, so ingest stores a null cell and drops the row when the column is a required time column, as main does for Date32 values outside its range. Data values are no longer saturated. Every real value fits the seconds clock, so seconds indexes are unchanged. Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-core/src/time.rs | 54 ++++++++++++------- .../python/tests/test_index_clocks.py | 33 ++++++++++++ .../diskgraph/diskgraph-rowpack/src/access.rs | 34 ++++++------ .../diskgraph/diskgraph-rowpack/src/parser.rs | 27 ++++++---- 4 files changed, 99 insertions(+), 49 deletions(-) diff --git a/packages/diskgraph/diskgraph-core/src/time.rs b/packages/diskgraph/diskgraph-core/src/time.rs index 8ccdf141b..38d4c2649 100644 --- a/packages/diskgraph/diskgraph-core/src/time.rs +++ b/packages/diskgraph/diskgraph-core/src/time.rs @@ -125,9 +125,10 @@ pub fn div_ceil(value: i64, divisor: i64) -> i64 { // 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. +/// 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) -> i64 { +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), @@ -138,26 +139,27 @@ pub fn ceil_to_clock(value: i64, unit: TimeUnit) -> i64 { } #[inline(always)] -fn ceil_from_tick(value: i64, tick: i64) -> i64 { +fn ceil_from_tick(value: i64, tick: i64) -> Option { let clock_tick = U::TICK_NANOS; if tick >= clock_tick { - value.saturating_mul(tick / clock_tick) + value.checked_mul(tick / clock_tick) } else { - div_ceil(value, clock_tick / tick) + Some(div_ceil(value, clock_tick / tick)) } } /// Floors `secs + nanos()` (with `0 <= nanos() < 1e9`) to a tick of the clock -/// `U`; `nanos` is only evaluated when a tick is shorter than a second. +/// `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) -> i64 { +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.saturating_mul(NS_PER_SEC / clock_tick) - .saturating_add(fraction) + secs.checked_mul(NS_PER_SEC / clock_tick)? + .checked_add(fraction) } else { - secs.div_euclid(clock_tick / NS_PER_SEC) + Some(secs.div_euclid(clock_tick / NS_PER_SEC)) } } @@ -1163,13 +1165,29 @@ mod tests { #[test] fn floor_parts_drops_the_fraction_below_one_tick() { - assert_eq!(floor_parts::(10, || 999_999_999), 10); - assert_eq!(floor_parts::(10, || 123_456_789), 10_123); + 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), - 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) ); - assert_eq!(floor_parts::(-1, || 500_000_000), -1); } #[test] @@ -1184,11 +1202,9 @@ mod tests { for clock in TimeUnit::INDEX_CLOCKS { let ticks = crate::with_index_clock!(clock, U => ceil_to_clock::(value, unit)); - assert_eq!( - ticks, - convert_ceil(value, unit, clock), - "{unit:?} -> {clock:?}" - ); + 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:?}"); } } } diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py index 74b8d40b9..8eed0daee 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py @@ -286,3 +286,36 @@ def test_messy_values_follow_mains_rules_in_every_clock(messy_data, tmp_path, cl 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] diff --git a/packages/diskgraph/diskgraph-rowpack/src/access.rs b/packages/diskgraph/diskgraph-rowpack/src/access.rs index 459964007..8c5f65a48 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/access.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/access.rs @@ -2413,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")) @@ -2504,7 +2504,7 @@ 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(_, _) => { @@ -2574,14 +2574,15 @@ impl<'p> RowAccess<'p> { }; if let Some(parser) = fast { if let Some((secs, fraction)) = parser.parse_timestamp_parts(text) { - return Ok(Some(floor_parts::(secs, || fraction_nanos(fraction)))); + let ticks = floor_parts::(secs, || fraction_nanos(fraction)); + return Ok(Some(in_clock(ticks)?)); } } if let Some((secs, nanos)) = parse_timestamp_text_parts(text, format) { - return Ok(Some(floor_parts::(secs, || nanos))); + 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")) } @@ -2591,13 +2592,11 @@ impl<'p> RowAccess<'p> { .downcast_ref::() .ok_or_else(|| anyhow!("decimal128 cast"))? .value(self.row); - timestamp_from_decimal::( + Ok(Some(in_clock(timestamp_from_decimal::( *scale, unit, |scale| apply_decimal_scale_i128(value, scale).to_i64(), - ) - .ok_or_else(|| anyhow!("decimal128 overflow")) - .map(Some) + ))?)) } DataType::Decimal256(_, scale) if format.is_none() => { let value = array @@ -2605,13 +2604,11 @@ impl<'p> RowAccess<'p> { .downcast_ref::() .ok_or_else(|| anyhow!("decimal256 cast"))? .value(self.row); - timestamp_from_decimal::( + Ok(Some(in_clock(timestamp_from_decimal::( *scale, unit, |scale| apply_decimal_scale_i256(value, scale).to_i64(), - ) - .ok_or_else(|| anyhow!("decimal256 overflow")) - .map(Some) + ))?)) } DataType::Int8 | DataType::Int16 @@ -2705,10 +2702,10 @@ impl<'p> RowAccess<'p> { if let Some(fmt) = format { let text = value.to_string(); if let Some((secs, nanos)) = parse_timestamp_text_parts(&text, Some(fmt)) { - return Ok(Some(floor_parts::(secs, || nanos))); + 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 @@ -2724,7 +2721,8 @@ impl<'p> RowAccess<'p> { if !(-106752..=24855).contains(&days) { return Err(anyhow!("date out of range: day {days}")); } - Ok(Some(ceil_to_clock::(i64::from(days), TimeUnit::Days))) + let ticks = ceil_to_clock::(i64::from(days), TimeUnit::Days); + Ok(Some(in_clock(ticks)?)) } _ => Err(anyhow!("timestamp type mismatch")), } @@ -3785,9 +3783,7 @@ mod tests { (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)) - }; + 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:?}"); diff --git a/packages/diskgraph/diskgraph-rowpack/src/parser.rs b/packages/diskgraph/diskgraph-rowpack/src/parser.rs index 33a1d8ea3..e838d4299 100644 --- a/packages/diskgraph/diskgraph-rowpack/src/parser.rs +++ b/packages/diskgraph/diskgraph-rowpack/src/parser.rs @@ -885,11 +885,16 @@ pub(crate) fn days_from_civil_fast(year: i32, month: u32, day: u32) -> i64 { era * 146097 + doe - 719468 } +/// 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")) +} + #[inline(always)] pub(crate) fn timestamp_from_int( value: i64, unit: Option, -) -> i64 { +) -> 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), @@ -907,11 +912,11 @@ pub(crate) fn timestamp_from_decimal( unit: Option, scaled: impl FnOnce(i32) -> Option, ) -> Option { - let ticks_per_unit = timestamp_from_int::(1, unit); + 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); } - Some(timestamp_from_int::(scaled(i32::from(scale))?, unit)) + timestamp_from_int::(scaled(i32::from(scale))?, unit) } pub(crate) fn days_from_secs(secs: i64) -> i32 { @@ -997,28 +1002,28 @@ pub(crate) fn timestamp_from_array(array: &ArrayRef, row: usize) .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("timestamp second cast"))?; - Ok(ceil_to_clock::(arr.value(row), TimeUnit::Seconds)) + 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(ceil_to_clock::(arr.value(row), TimeUnit::Milliseconds)) + 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(ceil_to_clock::(arr.value(row), TimeUnit::Microseconds)) + 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(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) } } } @@ -1031,7 +1036,7 @@ pub(crate) fn time_from_array(array: &ArrayRef, row: usize) -> Re .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time32 second cast"))?; - Ok(ceil_to_clock::( + in_clock(ceil_to_clock::( i64::from(arr.value(row)), TimeUnit::Seconds, )) @@ -1041,7 +1046,7 @@ pub(crate) fn time_from_array(array: &ArrayRef, row: usize) -> Re .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time32 millisecond cast"))?; - Ok(ceil_to_clock::( + in_clock(ceil_to_clock::( i64::from(arr.value(row)), TimeUnit::Milliseconds, )) @@ -1051,14 +1056,14 @@ pub(crate) fn time_from_array(array: &ArrayRef, row: usize) -> Re .as_any() .downcast_ref::() .ok_or_else(|| anyhow!("time64 microsecond cast"))?; - Ok(ceil_to_clock::(arr.value(row), TimeUnit::Microseconds)) + 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(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) + in_clock(ceil_to_clock::(arr.value(row), TimeUnit::Nanoseconds)) } DataType::Time32(_) | DataType::Time64(_) => Err(anyhow!( "unsupported Arrow time unit for timestamp conversion: {:?}", From 968e68ca42f1556e48ee79cc9bf582ad094ff839 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:44:38 +0000 Subject: [PATCH 10/25] diskgraph-train-table: skip rows whose windows leave the index clock Calendar arithmetic returns None instead of panicking, and a row whose window, lag or split shift runs past the clock is dropped rather than wrapped. Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-core/src/time.rs | 48 +++--- .../python/tests/test_index_clocks.py | 13 +- .../src/split_decider.rs | 60 ++++++-- .../src/ttgen_executor.rs | 140 ++++++------------ .../src/ttgen_resolve.rs | 85 ++++++++++- 5 files changed, 214 insertions(+), 132 deletions(-) diff --git a/packages/diskgraph/diskgraph-core/src/time.rs b/packages/diskgraph/diskgraph-core/src/time.rs index 38d4c2649..bf237799f 100644 --- a/packages/diskgraph/diskgraph-core/src/time.rs +++ b/packages/diskgraph/diskgraph-core/src/time.rs @@ -234,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) } @@ -605,14 +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 subsecond = ts - Abs::::from_secs(secs); - let dt = DateTime::from_timestamp(secs, 0) - .expect("timestamp out of range for calendar arithmetic") - .naive_utc(); + let dt = DateTime::from_timestamp(secs, 0)?.naive_utc(); let total_months = dt.month0() as i32 + months; // Euclidean division to handle negative months correctly @@ -623,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()) + subsecond + 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) @@ -637,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)), } } @@ -944,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); } @@ -953,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); } @@ -962,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); } @@ -971,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); } @@ -1093,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); } @@ -1101,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()); @@ -1111,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()); @@ -1242,7 +1242,17 @@ mod tests { #[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).raw(), 1_707_955_200_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] diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py index 8eed0daee..44acd6124 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py @@ -130,7 +130,7 @@ def _sampled_edges(index: str) -> dict: return sorted(edges) -def _training_table(index: str, clock: str) -> list: +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") @@ -140,7 +140,7 @@ def _training_table(index: str, clock: str) -> list: builder.set_target_stype("numerical") builder.aggregate( "orders_next_week", EdgeRef("users", "user_id", "orders", "user_id"), "count", - time_col="ordered_at", window=(0, 7), window_unit="days", + 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), @@ -319,3 +319,12 @@ def test_values_beyond_the_clock_are_null_and_drop_required_rows(tmp_path): 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)} diff --git a/packages/diskgraph/diskgraph-train-table/src/split_decider.rs b/packages/diskgraph/diskgraph-train-table/src/split_decider.rs index 9ac8ff71e..922eec0bc 100644 --- a/packages/diskgraph/diskgraph-train-table/src/split_decider.rs +++ b/packages/diskgraph/diskgraph-train-table/src/split_decider.rs @@ -279,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 { @@ -1051,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), } } } @@ -1067,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))` = @@ -1082,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)), } @@ -1240,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() { @@ -1310,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 { @@ -1956,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>( @@ -2116,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![ @@ -2982,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", @@ -2992,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", @@ -3005,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 0ef933a6e..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; @@ -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( @@ -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]; @@ -2904,34 +2898,11 @@ fn execute_has_edges( entity_cols: &[Option], observation_ts: Option>, ) -> Option { - // encode_upper expects index-clock ticks, which `U` is. - 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; @@ -2972,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; @@ -3134,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, @@ -3321,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() { @@ -3634,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); } @@ -3643,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); } @@ -3652,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); } @@ -3661,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); } @@ -3670,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); } @@ -3678,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); } @@ -3686,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_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] From 8fd395b5d157a3f6cc50ed034a305cd520ee7663 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:46:51 +0000 Subject: [PATCH 11/25] diskgraph: keep sub-second fractions in timestamp statistics Kumo timestamp strings and training-table dates print the fraction a finer clock stores, and nothing extra when it is zero, so seconds output is unchanged. Co-Authored-By: Claude Opus 5.5 --- .../diskgraph-ingest/src/col_stats.rs | 27 ++++++++++++++++-- .../diskgraph-train-table/src/ttgen_stats.rs | 28 ++++++++++++++----- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs index 0aca584ed..9212c54fc 100644 --- a/packages/diskgraph/diskgraph-ingest/src/col_stats.rs +++ b/packages/diskgraph/diskgraph-ingest/src/col_stats.rs @@ -1602,9 +1602,9 @@ fn format_kumo_timestamp(ticks: i64, clock: TimeUnit) -> String { if ticks == i64::MAX || ticks == i64::MIN { return "kumo_timestamp_NaT".to_string(); } - let (epoch_secs, _) = split_parts(ticks, clock); - 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(), } } @@ -3970,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-train-table/src/ttgen_stats.rs b/packages/diskgraph/diskgraph-train-table/src/ttgen_stats.rs index 69196bde7..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; @@ -515,8 +515,8 @@ impl StatsMerger { // 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)) }; @@ -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. From 5f9b9f5df98709272922f1a67fec9a033bafd618 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:59:16 +0000 Subject: [PATCH 12/25] diskgraph-pybind: convert builder times once, at build() DiskGraphTTBuilder held plan times in nanoseconds, so a seconds or millisecond plan clamped anchors outside 1677-2262. It now keeps each anchor, split bound and offset as given and converts it straight into the plan clock in build(). Co-Authored-By: Claude Opus 5.5 --- .../diskgraph-core/src/tt_builder.rs | 171 ++++++------ .../python/tests/test_index_clocks.py | 21 ++ .../diskgraph-pybind/src/tt_builder.rs | 243 +++++++++++------- 3 files changed, 271 insertions(+), 164 deletions(-) diff --git a/packages/diskgraph/diskgraph-core/src/tt_builder.rs b/packages/diskgraph/diskgraph-core/src/tt_builder.rs index 4d2ba80cd..35153c750 100644 --- a/packages/diskgraph/diskgraph-core/src/tt_builder.rs +++ b/packages/diskgraph/diskgraph-core/src/tt_builder.rs @@ -51,7 +51,7 @@ pub const ENTITY_NODE_ID_COL: &str = "entity_node_id"; /// Resolved LP RHS positives (`list` node-ids). pub const TARGET_NODE_IDS_COL: &str = "target_node_ids"; -use crate::time::{convert_trunc, Abs, Dur, UnitSystem}; +use crate::time::{Abs, Dur, UnitSystem}; use serde::{Deserialize, Serialize}; // ─── Top-level plan ────────────────────────────────────────────────────────── @@ -540,31 +540,9 @@ pub enum StageFilter { // ─── Index clocks ──────────────────────────────────────────────────────────── // // A serialized plan stores raw ticks of one clock, named by a top-level -// `timestamp_unit` (omitted for seconds). Converting between clocks truncates -// anchors, split bounds and offsets, as `main` does when it builds a plan. +// `timestamp_unit` (omitted for seconds). impl DiskGraphTTPlan { - pub fn to_clock(&self) -> serde_json::Result> { - let mut plan: DiskGraphTTPlan = serde_json::from_value(serde_json::to_value(self)?)?; - if let Some(time) = &mut plan.entity.default_time_value { - rescale_abs(time, U::UNIT); - } - if let Some(temporal) = &mut plan.temporal { - for time in &mut temporal.timestamps { - rescale_abs(time, U::UNIT); - } - } - if let Some(split) = &mut plan.split { - split.rescale_from(U::UNIT); - } - for stage in &mut plan.stages { - for op in &mut stage.ops { - op.rescale_from(U::UNIT); - } - } - Ok(plan) - } - pub fn to_json(&self) -> serde_json::Result { if U::UNIT == TimeUnit::Seconds { return serde_json::to_string(self); @@ -575,66 +553,109 @@ impl DiskGraphTTPlan { } } -/// 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) -} - -impl SplitSpec { - pub fn to_clock(&self) -> serde_json::Result> { - let mut split: SplitSpec = serde_json::from_value(serde_json::to_value(self)?)?; - split.rescale_from(U::UNIT); - Ok(split) - } - - fn rescale_from(&mut self, from: TimeUnit) { - for offset in self.offsets.iter_mut().flatten() { - rescale_dur(offset, from); - } - for (start, end) in self.ranges.iter_mut().flatten() { - rescale_abs(start, from); - rescale_abs(end, from); - } - if let Some(time) = &mut self.backtesting_timestamp { - rescale_abs(time, from); - } - } -} - impl Op { - fn rescale_from(&mut self, from: TimeUnit) { + /// 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: offsets, - .. - } - | Op::FutureLabels { + 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: offsets, - .. - } => { - target_op.rescale_from(from); - for offset in offsets { - rescale_dur(offset, from); - } - } - _ => {} + 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, + }, } } } -fn rescale_abs(time: &mut Abs, from: TimeUnit) { - *time = Abs::from_raw(convert_trunc(time.raw(), from, U::UNIT)); -} - -fn rescale_dur(duration: &mut Dur, from: TimeUnit) { - *duration = Dur::from_raw(convert_trunc(duration.raw(), from, U::UNIT)); +/// 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) } // --------------------------------------------------------------------------- diff --git a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py index 44acd6124..3123fe0ab 100644 --- a/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py +++ b/packages/diskgraph/diskgraph-pybind/python/tests/test_index_clocks.py @@ -328,3 +328,24 @@ def test_rows_whose_window_leaves_the_clock_are_skipped(whole_day_data, tmp_path 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/tt_builder.rs b/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs index 219b3f85f..81255e69e 100644 --- a/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs +++ b/packages/diskgraph/diskgraph-pybind/src/tt_builder.rs @@ -18,8 +18,8 @@ use pyo3::prelude::*; use pyo3::types::PyList; use diskgraph_core::time::{ - abs_from_unit, ceil_parts, convert_trunc, dur_from_unit, split_parts, Abs, Nanoseconds, - TimeUnit, + 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; @@ -331,22 +331,44 @@ impl StageFilter { // ─── DiskGraphTTBuilder ────────────────────────────────────────────────────── -/// Holds plan times in nanoseconds; `build()` writes the plan in `timestamp_unit`. +/// 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>, - default_time_secs: Option, - temporal: Option>, + entity: 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] @@ -355,7 +377,6 @@ impl DiskGraphTTBuilder { Ok(Self { timestamp_unit: parse_index_clock(timestamp_unit)?, entity: None, - default_time_secs: None, temporal: None, prediction_type: tt::PredictionType::Node, is_prediction: false, @@ -367,7 +388,7 @@ impl DiskGraphTTBuilder { }) } - /// `FOR EACH {table}.{column}`; `default_time_value` is in epoch seconds. + /// `FOR EACH {table}.{column}`. #[pyo3(signature = (table, column, output_time_col_name=None, default_time_value=None))] fn for_each( &mut self, @@ -380,12 +401,11 @@ impl DiskGraphTTBuilder { table, column, output_time_col_name, - default_time_value: None, + default_time_value: default_time_value.map(Abs::from_raw), time_col: None, valid_to_col: None, filter: None, }); - self.default_time_secs = default_time_value; } /// Entity-level create-time column for temporal filtering. @@ -438,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, }); @@ -462,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(()) } @@ -529,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, @@ -540,7 +557,7 @@ impl DiskGraphTTBuilder { row_filter: row_filter.map(|f| f.inner), limit, target_entity_table, - }); + })); Ok(()) } @@ -554,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(()) } @@ -571,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, self.timestamp_unit)?, - }); + })); Ok(()) } @@ -594,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(()) } @@ -611,20 +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(()) } @@ -636,20 +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(()) } @@ -665,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(()) } @@ -683,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 @@ -725,19 +755,53 @@ impl DiskGraphTTBuilder { lp_resolution: None, lp_rhs_is_temporal: false, fact_pseudo_time: None, - }; + }) + } +} - let default_time_secs = self.default_time_secs.take(); - with_index_clock!(self.timestamp_unit, U => { - plan.to_clock::().and_then(|mut plan| { - plan.entity.default_time_value = default_time_secs.map(Abs::from_secs); - plan.to_json() - }) +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(), + }, }) - .map_err(|e| crate::DiskgraphConfigError::new_err(e.to_string())) } } +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 { @@ -889,7 +953,7 @@ fn parse_timestamp_unit(unit: &str) -> PyResult { /// *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, @@ -958,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, @@ -987,14 +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, @@ -1028,9 +1091,9 @@ fn make_split_spec(args: SplitSpecArgs) -> PyResult> #[pyfunction] #[pyo3(signature = (spec, timestamp_unit="seconds"))] pub fn build_split_spec_json(spec: &Bound<'_, PyAny>, timestamp_unit: &str) -> PyResult { - let split = make_split_spec(extract_split_spec(spec)?)?; + let args = extract_split_spec(spec)?; with_index_clock!(parse_index_clock(timestamp_unit)?, U => { - split.to_clock::().and_then(|split| serde_json::to_string(&split)) + serde_json::to_string(&make_split_spec::(args)?) }) .map_err(|e| { crate::DiskgraphConfigError::new_err(format!("failed to serialize split spec: {e}")) @@ -1099,7 +1162,9 @@ fn extract_value(value: &Bound<'_, PyAny>, clock: TimeUnit) -> PyResult(SplitSpecArgs { kind: Some("TimeRangeSplit".to_string()), ratio: None, rnd_seed: None, From 27bb6b2f8cec8f5658b42c4512d84001ed3a8dc9 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 23:07:41 +0000 Subject: [PATCH 13/25] diskgraph-ingest: document index clocks Co-Authored-By: Claude Opus 5.5 --- packages/diskgraph/diskgraph-ingest/README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) 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: From 8056db46b992babaead7f3a764e61c26cc08710a Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 18:55:40 +0000 Subject: [PATCH 14/25] kumoapi: add TimestampPrecision, the six index clock names Co-Authored-By: Claude Opus 5.5 --- packages/kumo-api/kumoapi/typing.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/kumo-api/kumoapi/typing.py b/packages/kumo-api/kumoapi/typing.py index a61842a63..81eed315c 100644 --- a/packages/kumo-api/kumoapi/typing.py +++ b/packages/kumo-api/kumoapi/typing.py @@ -31,6 +31,9 @@ def _forbid_extra_config(**options: Any) -> Any: SUPPORTED_MULTICATEGORICAL_DELIMITERS = frozenset({',', '|', ';', '\t'}) MulticategoricalDelimiter = Literal[',', '|', ';', '\t'] +TimestampPrecision = Literal['nanoseconds', 'microseconds', 'milliseconds', + 'seconds', 'hours', 'days'] + class Stype(StrEnum): r"""The semantic type of a column. From f5ae8df7111417c95b87a9aea3aa7caa742a92d7 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 18:56:43 +0000 Subject: [PATCH 15/25] kumo-core: carry the index clock in DiskgraphIngestSpec Written into the ingest config only when it is not seconds, so seconds configs stay identical to main. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-core/src/kumo_core/diskgraph/ingest.py | 10 ++++++++++ 1 file changed, 10 insertions(+) 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) From 048d1dca7035fdac571d89442e41ded4d24d2246 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 18:56:43 +0000 Subject: [PATCH 16/25] kumo-engine: Graph.snapshot(timestamp_precision) Seconds keeps main's snapshot cache key; every other clock gets its own. Co-Authored-By: Claude Opus 5.5 --- packages/kumo-engine/kumo/engine/graph.py | 33 ++++++++++++++++++----- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/packages/kumo-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index b5ecf3299..a8834d6eb 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -64,10 +64,11 @@ from dataclasses import dataclass, field from functools import cached_property from pathlib import Path -from typing import Any, Dict, Mapping, Sequence +from typing import Any, Dict, Mapping, Sequence, get_args from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge +from kumoapi.typing import TimestampPrecision 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,19 @@ 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 + 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}") self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) @@ -383,9 +397,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 +482,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 +514,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()), ), From 0188311ca84126d2bc9b951c9b144d7a9bd3442d Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 19:00:17 +0000 Subject: [PATCH 17/25] kumo-core: compile plans against the index clock Training and prediction pass index_path to DiskGraphPlanCompiler, the runtime adapter and NodeRange anchors understand hours and days indexes, and timestamp statistics keep every nanosecond digit of their fraction. A temporal-LP training batch takes its valid RHS items at its earliest anchor, since its anchors share a second rather than a tick. Co-Authored-By: Claude Opus 5.5 --- .../src/kumo_core/diskgraph/lightning_data.py | 4 ++++ .../src/kumo_core/diskgraph/runtime_adapter.py | 2 ++ packages/kumo-core/src/kumo_core/pquery.py | 10 +++++++++- packages/kumo-core/src/kumo_core/prediction.py | 2 +- .../train/shallow_rhs_link_gnn_module.py | 2 +- .../kumo-core/src/kumo_core/training/table.py | 2 +- packages/kumo-core/tests/test_prediction.py | 3 ++- packages/kumo-core/tests/test_train_table.py | 3 ++- test/train/test_shallow_rhs_link_gnn_module.py | 18 +++++++++++++++++- 9 files changed, 39 insertions(+), 7 deletions(-) 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_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/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( From 6c06ce88dda3717fb56a15319659b46b47b6ceb1 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 19:00:56 +0000 Subject: [PATCH 18/25] kumoai_studio: Graph.snapshot(timestamp_precision) The clock is sent as a query parameter only when it is not seconds, and the SDK requires the server to echo the clock it built; a missing echo means an older server, which builds seconds. Without a clock, snapshot() keeps the clock of the graph's current snapshot, so the implicit calls in fit and predict reuse it instead of rebuilding in seconds. Co-Authored-By: Claude Opus 5.5 --- .../kumo-sdk/kumoai_studio/backend/base.py | 2 ++ .../kumo-sdk/kumoai_studio/client/graph.py | 18 +++++++++++++++++- packages/kumo-sdk/kumoai_studio/graph/graph.py | 15 ++++++++++++++- 3 files changed, 33 insertions(+), 2 deletions(-) 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..c7812d0d2 100644 --- a/packages/kumo-sdk/kumoai_studio/client/graph.py +++ b/packages/kumo-sdk/kumoai_studio/client/graph.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from http import HTTPStatus -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, get_args from kumoapi.common import ValidationResponse from kumoapi.data_snapshot import GraphSnapshotID, GraphSnapshotResource @@ -14,6 +14,7 @@ GraphValidationRequest, ) from kumoapi.json_serde import to_json_dict +from kumoapi.typing import TimestampPrecision from kumoai_studio.client.endpoints import GraphEndpoints from kumoai_studio.client.transport import Transport @@ -79,16 +80,31 @@ def create_snapshot( *, refresh_source: bool = False, project_id: Optional[str] = None, + timestamp_precision: TimestampPrecision = 'seconds', ) -> GraphSnapshotID: + 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}") 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) From 8c5674953f98b97d532380c55c6578e085dea0d2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 19:05:31 +0000 Subject: [PATCH 19/25] hosted: build snapshots in the requested index clock The clock travels from the snapshot request through GraphDataComputeInput, GraphDataVersionInfo and DiskgraphIngestArgs into the ingest config. Seconds keeps main's snapshot identity; other clocks hash their own. The refresh workflow passes the clock to the ingest child only behind the diskgraph-index-clock-v1 patch. The router echoes the clock recorded for the built snapshot in X-Kumo-Timestamp-Precision. Co-Authored-By: Claude Opus 5.5 --- kumo/diskgraph/ingest.py | 3 +++ kumo/diskgraph/service.py | 2 ++ pool/activities.py | 2 ++ pool/diskgraph_args.py | 3 +++ .../activities/data_backend_activities.py | 8 ++++++-- .../workflows/diskgraph_ingest_workflow.py | 4 ++++ .../workflows/graph_refresh_workflow.py | 5 +++++ rest/config/data_backend.py | 3 +++ rest/config/graph_v2.py | 2 ++ rest/driver/data_snapshot_api_driver.py | 3 +++ rest/public/routers/data_snapshot.py | 19 +++++++++++++++++-- .../test_diskgraph_graph_refresh_workflow.py | 1 + 12 files changed, 51 insertions(+), 4 deletions(-) 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..e2372f284 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, 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/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/rest/backends/test_diskgraph_graph_refresh_workflow.py b/test/rest/backends/test_diskgraph_graph_refresh_workflow.py index 4fc7227b7..25adcf16d 100644 --- a/test/rest/backends/test_diskgraph_graph_refresh_workflow.py +++ b/test/rest/backends/test_diskgraph_graph_refresh_workflow.py @@ -183,6 +183,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") From 2cea1085693f7b1bb6fa271b6ca41087a01c3817 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 19:05:31 +0000 Subject: [PATCH 20/25] hosted: compile training and prediction plans against the index clock Co-Authored-By: Claude Opus 5.5 --- kumo/data/generic/diskgraph_batch_pred.py | 9 ++++++--- kumo/data/generic/diskgraph_train_table.py | 2 +- test/data/generic/test_diskgraph_predict_plan_ttng.py | 8 ++++++++ .../generic/test_ttng_cross_graph_contamination_e2e.py | 1 + 4 files changed, 16 insertions(+), 4 deletions(-) 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/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_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], From 2e4bf20b052968fe70a4387a63eed5f9fdb899a2 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 19:05:31 +0000 Subject: [PATCH 21/25] stats: compute prediction statistics in the snapshot's index clock Batch-prediction statistics read the clock from the snapshot's index_metadata.json when the snapshot was not built in seconds, falling back to seconds when the file is missing or unreadable. Python sample statistics convert through resolve_temporal in that clock; typed seconds columns keep their vectorized path, which returns the same values. Co-Authored-By: Claude Opus 5.5 --- kumo/connector/connector.py | 4 ++- kumo/connector/timestamp_convert.py | 41 +++++++++++++++++----- kumo/connector/utils.py | 23 +++++++----- kumo/diskgraph/service.py | 17 +++++++++ rest/backends/data_backend/data_backend.py | 18 ++++++++-- test/rest/backends/test_data_backend.py | 17 +++++++-- 6 files changed, 98 insertions(+), 22 deletions(-) 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..861469532 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 @@ -190,9 +202,20 @@ 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. """ - if hint.kind is HintKind.NATIVE or is_datetime64_any_dtype(series): + 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. Typed columns on the seconds clock take a vectorized path + that returns exactly what :func:`_resolve` does. + """ + native = hint.kind is HintKind.NATIVE or is_datetime64_any_dtype(series) + if native and timestamp_precision == 'seconds': return timestamp_secs_from_array(series) - return _resolve(series, hint, days=False) + return _resolve(series, hint, days=False, + timestamp_precision=timestamp_precision) def to_epoch_days(series: pd.Series, hint: TemporalHint) -> np.ndarray: 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/diskgraph/service.py b/kumo/diskgraph/service.py index e2372f284..c96232cc3 100644 --- a/kumo/diskgraph/service.py +++ b/kumo/diskgraph/service.py @@ -1717,12 +1717,27 @@ def cleanup_staging(self, job_ids: List[str]) -> None: if os.path.isdir(staging_dir): shutil.rmtree(staging_dir, ignore_errors=True) + @blocking_method() + def snapshot_timestamp_unit(self, handle: DiskgraphHandle, + dataplane_root: str) -> str: + """The clock in the snapshot's ``index_metadata.json``, or seconds + when the file is missing or unreadable.""" + from io_utils import fsutil + snapshot_dir = self.dataplane_snapshot_path(handle, dataplane_root) + path = os.path.join(snapshot_dir, 'index_metadata.json') + try: + return json.loads(fsutil.read_text(path)).get( + 'timestamp_precision', 'seconds') + except Exception: + return 'seconds' + @blocking_method() def compute_col_stats_for_parquet( self, 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. @@ -1735,6 +1750,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. @@ -1744,6 +1760,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/rest/backends/data_backend/data_backend.py b/rest/backends/data_backend/data_backend.py index 7b3a1d439..bf275a77f 100644 --- a/rest/backends/data_backend/data_backend.py +++ b/rest/backends/data_backend/data_backend.py @@ -32,7 +32,7 @@ from io_utils import fsutil, json, logging from io_utils.json import fast_from_json, load_many_json, save_json from kumo.connector import MAX_SUPPORTED_FILES -from kumo.data import get_active_cache +from kumo.data import get_active_cache, get_persist_dir from kumo.data.backend_cache import ( AsyncDataTaskStatus, BackendCacheManager, @@ -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 @@ -1221,7 +1223,7 @@ async def start_bp_stats_activity( """ from kumoapi.pquery import QueryType - from kumo.diskgraph import get_service + from kumo.diskgraph import DiskgraphHandle, get_service from kumo.diskgraph.service import _to_kumo_col_specs from rest.backends.ml_backend.data_util import ( get_job_base_dir, @@ -1286,6 +1288,16 @@ async def start_bp_stats_activity( raise ValueError( 'Direct Databricks stats stage contains no Parquet files.') + timestamp_precision = 'seconds' + version_info = job_input.graph_version_info + if (version_info is not None + and version_info.timestamp_precision != 'seconds'): + snapshot = DiskgraphHandle( + graph_key=job_input.graph_data_compute_input.graph_key, + graph_data_version=version_info.graph_data_version) + timestamp_precision = await get_service( + ).snapshot_timestamp_unit.coroutine(snapshot, get_persist_dir()) + 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 +1375,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 +1451,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/test/rest/backends/test_data_backend.py b/test/rest/backends/test_data_backend.py index 6b4497689..7f1dd05e6 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,9 @@ 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_data_compute_input.graph_key = 'graph' + job_input.graph_version_info = SimpleNamespace(graph_data_version='v1', + timestamp_precision=clock) table_def = mock.Mock() table_def.fingerprint.return_value = 'fingerprint' @@ -90,7 +95,9 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path) -> None: table_refresh_result=mock.Mock(save=refresh_save), ) compute = mock.Mock(coroutine=mock.AsyncMock(return_value=('{}', '{}'))) - service = mock.Mock(compute_col_stats_for_parquet=compute) + snapshot_clock = mock.Mock(coroutine=mock.AsyncMock(return_value=clock)) + service = mock.Mock(compute_col_stats_for_parquet=compute, + snapshot_timestamp_unit=snapshot_clock) resolved = SimpleNamespace( root=str(artifact_root), files=tuple(str(path) for path in selected), @@ -122,6 +129,9 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path) -> None: mock.patch( 'rest.backends.data_backend.data_backend.get_active_cache', return_value=cache), mock.patch( + 'rest.backends.data_backend.data_backend.' + 'get_persist_dir', return_value=str(tmp_path)), \ + mock.patch( 'rest.backends.data_backend.data_backend.' '_list_parquet_files', side_effect=AssertionError( @@ -131,6 +141,8 @@ 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 + assert snapshot_clock.coroutine.await_count == (clock != 'seconds') class _UnapprovedSingleActivityError(PublicException): @@ -877,6 +889,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( From 4fbef12d5f3a41f77312bc45514e6c94e083dbfc Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 19:18:51 +0000 Subject: [PATCH 22/25] test: index clocks across the SDKs, hosted snapshots and statistics Six-clock acceptance in kumo-engine and kumoai_studio, the unchanged seconds cache key and hosted identity, the clock echo for old and new SDKs and servers, the patched ingest workflow argument, statistics in the snapshot's clock, and NodeRange and seed times on hours and days indexes. Co-Authored-By: Claude Opus 5.5 --- .../kumo-core/tests/test_diskgraph_runtime.py | 11 ++ packages/kumo-core/tests/test_pquery.py | 33 +++++ .../kumo-engine/test/test_snapshot_clock.py | 76 ++++++++++ .../test/graph/test_snapshot_clock.py | 90 +++++++++++ .../connector/local/test_timestamp_convert.py | 41 +++++ .../data/generic/test_diskgraph_batch_pred.py | 2 + .../test_diskgraph_generate_train_table.py | 3 + .../generic/test_diskgraph_runtime_adapter.py | 11 ++ .../test_diskgraph_train_table_forecasting.py | 6 +- .../test_diskgraph_graph_refresh_workflow.py | 21 +++ test/rest/backends/test_snapshot_clock.py | 140 ++++++++++++++++++ .../test_diskgraph_lightning_data.py | 26 ++++ 12 files changed, 457 insertions(+), 3 deletions(-) create mode 100644 packages/kumo-engine/test/test_snapshot_clock.py create mode 100644 packages/kumo-sdk/test/graph/test_snapshot_clock.py create mode 100644 test/rest/backends/test_snapshot_clock.py 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-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/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/test/connector/local/test_timestamp_convert.py b/test/connector/local/test_timestamp_convert.py index 8a12b70e3..d96823267 100644 --- a/test/connector/local/test_timestamp_convert.py +++ b/test/connector/local/test_timestamp_convert.py @@ -450,3 +450,44 @@ 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')) + + +@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_on_the_seconds_clock_match_diskgraph(series): + hint = TemporalHint.native() + assert list(tc.to_epoch_ticks(series, hint, 'seconds')) == list( + tc._resolve(series, hint, days=False, timestamp_precision='seconds')) + + +def test_statistics_convert_in_the_snapshot_clock(monkeypatch): + from kumoapi.typing import Stype + + from kumo.config import BaseColumnConfig + from kumo.connector.utils import coerce_timestamp_stype_series_for_stats + + clocks = [] + + def resolve_temporal(values, fmt, unit, as_days, *, timestamp_precision): + clocks.append(timestamp_precision) + return [1_700_000_000_123], fmt, unit + + monkeypatch.setattr(tc, 'resolve_temporal', resolve_temporal) + config = BaseColumnConfig(name='ts', stype=Stype.timestamp) + coerced = coerce_timestamp_stype_series_for_stats( + pd.Series(['2023-11-14 22:13:20.123']), config, 'milliseconds') + + assert clocks == ['milliseconds'] + assert coerced.tolist() == [pd.Timestamp('2023-11-14 22:13:20.123')] 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_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/rest/backends/test_diskgraph_graph_refresh_workflow.py b/test/rest/backends/test_diskgraph_graph_refresh_workflow.py index 25adcf16d..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 @@ -247,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..415de8504 --- /dev/null +++ b/test/rest/backends/test_snapshot_clock.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import json +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 DiskgraphHandle, 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() + + +@pytest.mark.parametrize('metadata,clock', [ + ({'timestamp_precision': 'milliseconds'}, 'milliseconds'), + ({'dataset_name': 'graph'}, 'seconds'), + (None, 'seconds'), +]) +def test_snapshot_clock_comes_from_index_metadata(tmp_path, metadata, clock): + service = DiskgraphService(cast(Any, None)) + handle = DiskgraphHandle(graph_key='graph', graph_data_version='v1') + if metadata is not None: + snapshot = tmp_path / 'diskgraph_snapshots' / 'graph' / 'v1' + snapshot.mkdir(parents=True) + (snapshot / 'index_metadata.json').write_text(json.dumps(metadata)) + assert service.snapshot_timestamp_unit(handle, str(tmp_path)) == clock + + +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( From c70bc20f0a54cc3cdd7c0ca1ebb956a0fd1736eb Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:16:27 +0000 Subject: [PATCH 23/25] kumoapi: validate the index clock in one shared helper Co-Authored-By: Claude Opus 5.5 --- packages/kumo-api/kumoapi/typing.py | 9 ++++++++- packages/kumo-engine/kumo/engine/graph.py | 9 +++------ packages/kumo-sdk/kumoai_studio/client/graph.py | 9 +++------ 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/packages/kumo-api/kumoapi/typing.py b/packages/kumo-api/kumoapi/typing.py index 81eed315c..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 @@ -35,6 +35,13 @@ def _forbid_extra_config(**options: Any) -> Any: '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-engine/kumo/engine/graph.py b/packages/kumo-engine/kumo/engine/graph.py index a8834d6eb..d2388efd3 100644 --- a/packages/kumo-engine/kumo/engine/graph.py +++ b/packages/kumo-engine/kumo/engine/graph.py @@ -64,11 +64,11 @@ from dataclasses import dataclass, field from functools import cached_property from pathlib import Path -from typing import Any, Dict, Mapping, Sequence, get_args +from typing import Any, Dict, Mapping, Sequence from kumo_core.pquery import NativeGraph, normalize_graph from kumoapi.graph import Edge -from kumoapi.typing import TimestampPrecision +from kumoapi.typing import TimestampPrecision, validate_timestamp_precision from kumo.engine._scratch import ( fingerprint, @@ -367,10 +367,7 @@ def snapshot( """ from filelock import FileLock - 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}") + validate_timestamp_precision(timestamp_precision) self.validate() for table in self._named_tables.values(): table.connector.validate_scratch_dir(get_scratch_dir()) diff --git a/packages/kumo-sdk/kumoai_studio/client/graph.py b/packages/kumo-sdk/kumoai_studio/client/graph.py index c7812d0d2..1e7d7014e 100644 --- a/packages/kumo-sdk/kumoai_studio/client/graph.py +++ b/packages/kumo-sdk/kumoai_studio/client/graph.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from http import HTTPStatus -from typing import Any, Dict, Optional, get_args +from typing import Any, Dict, Optional from kumoapi.common import ValidationResponse from kumoapi.data_snapshot import GraphSnapshotID, GraphSnapshotResource @@ -14,7 +14,7 @@ GraphValidationRequest, ) from kumoapi.json_serde import to_json_dict -from kumoapi.typing import TimestampPrecision +from kumoapi.typing import TimestampPrecision, validate_timestamp_precision from kumoai_studio.client.endpoints import GraphEndpoints from kumoai_studio.client.transport import Transport @@ -82,10 +82,7 @@ def create_snapshot( project_id: Optional[str] = None, timestamp_precision: TimestampPrecision = 'seconds', ) -> GraphSnapshotID: - 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}") + validate_timestamp_precision(timestamp_precision) params: Dict[str, Any] = { 'graph_id': graph_id, 'refresh_source': refresh_source From dc525627af93c7f3c6ddfa7c1b4af2b36d230ac4 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:16:45 +0000 Subject: [PATCH 24/25] stats: take the prediction statistics clock from the recorded snapshot version Co-Authored-By: Claude Opus 5.5 --- kumo/diskgraph/service.py | 14 -------------- rest/backends/data_backend/data_backend.py | 16 +++++----------- test/rest/backends/test_data_backend.py | 12 ++---------- test/rest/backends/test_snapshot_clock.py | 18 +----------------- 4 files changed, 8 insertions(+), 52 deletions(-) diff --git a/kumo/diskgraph/service.py b/kumo/diskgraph/service.py index c96232cc3..f230bbd49 100644 --- a/kumo/diskgraph/service.py +++ b/kumo/diskgraph/service.py @@ -1717,20 +1717,6 @@ def cleanup_staging(self, job_ids: List[str]) -> None: if os.path.isdir(staging_dir): shutil.rmtree(staging_dir, ignore_errors=True) - @blocking_method() - def snapshot_timestamp_unit(self, handle: DiskgraphHandle, - dataplane_root: str) -> str: - """The clock in the snapshot's ``index_metadata.json``, or seconds - when the file is missing or unreadable.""" - from io_utils import fsutil - snapshot_dir = self.dataplane_snapshot_path(handle, dataplane_root) - path = os.path.join(snapshot_dir, 'index_metadata.json') - try: - return json.loads(fsutil.read_text(path)).get( - 'timestamp_precision', 'seconds') - except Exception: - return 'seconds' - @blocking_method() def compute_col_stats_for_parquet( self, diff --git a/rest/backends/data_backend/data_backend.py b/rest/backends/data_backend/data_backend.py index bf275a77f..d4d84006f 100644 --- a/rest/backends/data_backend/data_backend.py +++ b/rest/backends/data_backend/data_backend.py @@ -32,7 +32,7 @@ from io_utils import fsutil, json, logging from io_utils.json import fast_from_json, load_many_json, save_json from kumo.connector import MAX_SUPPORTED_FILES -from kumo.data import get_active_cache, get_persist_dir +from kumo.data import get_active_cache from kumo.data.backend_cache import ( AsyncDataTaskStatus, BackendCacheManager, @@ -1223,7 +1223,7 @@ async def start_bp_stats_activity( """ from kumoapi.pquery import QueryType - from kumo.diskgraph import DiskgraphHandle, get_service + from kumo.diskgraph import get_service from kumo.diskgraph.service import _to_kumo_col_specs from rest.backends.ml_backend.data_util import ( get_job_base_dir, @@ -1288,15 +1288,9 @@ async def start_bp_stats_activity( raise ValueError( 'Direct Databricks stats stage contains no Parquet files.') - timestamp_precision = 'seconds' - version_info = job_input.graph_version_info - if (version_info is not None - and version_info.timestamp_precision != 'seconds'): - snapshot = DiskgraphHandle( - graph_key=job_input.graph_data_compute_input.graph_key, - graph_data_version=version_info.graph_data_version) - timestamp_precision = await get_service( - ).snapshot_timestamp_unit.coroutine(snapshot, get_persist_dir()) + 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, diff --git a/test/rest/backends/test_data_backend.py b/test/rest/backends/test_data_backend.py index 7f1dd05e6..a06f314e3 100644 --- a/test/rest/backends/test_data_backend.py +++ b/test/rest/backends/test_data_backend.py @@ -73,9 +73,7 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path, ) 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_data_compute_input.graph_key = 'graph' - job_input.graph_version_info = SimpleNamespace(graph_data_version='v1', - timestamp_precision=clock) + job_input.graph_version_info = SimpleNamespace(timestamp_precision=clock) table_def = mock.Mock() table_def.fingerprint.return_value = 'fingerprint' @@ -95,9 +93,7 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path, table_refresh_result=mock.Mock(save=refresh_save), ) compute = mock.Mock(coroutine=mock.AsyncMock(return_value=('{}', '{}'))) - snapshot_clock = mock.Mock(coroutine=mock.AsyncMock(return_value=clock)) - service = mock.Mock(compute_col_stats_for_parquet=compute, - snapshot_timestamp_unit=snapshot_clock) + service = mock.Mock(compute_col_stats_for_parquet=compute) resolved = SimpleNamespace( root=str(artifact_root), files=tuple(str(path) for path in selected), @@ -129,9 +125,6 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path, mock.patch( 'rest.backends.data_backend.data_backend.get_active_cache', return_value=cache), mock.patch( - 'rest.backends.data_backend.data_backend.' - 'get_persist_dir', return_value=str(tmp_path)), \ - mock.patch( 'rest.backends.data_backend.data_backend.' '_list_parquet_files', side_effect=AssertionError( @@ -142,7 +135,6 @@ async def test_bp_stats_reads_only_manifest_selected_parts(tmp_path, 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 - assert snapshot_clock.coroutine.await_count == (clock != 'seconds') class _UnapprovedSingleActivityError(PublicException): diff --git a/test/rest/backends/test_snapshot_clock.py b/test/rest/backends/test_snapshot_clock.py index 415de8504..fd8d7293d 100644 --- a/test/rest/backends/test_snapshot_clock.py +++ b/test/rest/backends/test_snapshot_clock.py @@ -3,7 +3,6 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio -import json from hashlib import md5 from types import SimpleNamespace from typing import Any, cast, get_args @@ -15,7 +14,7 @@ from kumoapi.typing import TimestampPrecision from io_utils.json import from_json_dict -from kumo.diskgraph.service import DiskgraphHandle, DiskgraphService +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 @@ -115,21 +114,6 @@ def test_unknown_clock_is_rejected(snapshot_api): snapshot_api.driver.create_graph_snapshot.assert_not_awaited() -@pytest.mark.parametrize('metadata,clock', [ - ({'timestamp_precision': 'milliseconds'}, 'milliseconds'), - ({'dataset_name': 'graph'}, 'seconds'), - (None, 'seconds'), -]) -def test_snapshot_clock_comes_from_index_metadata(tmp_path, metadata, clock): - service = DiskgraphService(cast(Any, None)) - handle = DiskgraphHandle(graph_key='graph', graph_data_version='v1') - if metadata is not None: - snapshot = tmp_path / 'diskgraph_snapshots' / 'graph' / 'v1' - snapshot.mkdir(parents=True) - (snapshot / 'index_metadata.json').write_text(json.dumps(metadata)) - assert service.snapshot_timestamp_unit(handle, str(tmp_path)) == clock - - def test_prediction_statistics_use_the_snapshot_clock(monkeypatch): import diskgraph From f42fb95c10cb975f640ee46b4451954463336143 Mon Sep 17 00:00:00 2001 From: Vedaant Jain Date: Wed, 28 Oct 2026 22:18:55 +0000 Subject: [PATCH 25/25] stats: convert typed columns into any index clock on the vectorized path Typed columns take the same vectorized conversion in every clock, which matches resolve_temporal; text statistics truncate into the snapshot clock. Co-Authored-By: Claude Opus 5.5 --- kumo/connector/timestamp_convert.py | 42 +++++++++------ .../connector/local/test_timestamp_convert.py | 54 ++++++++++--------- 2 files changed, 53 insertions(+), 43 deletions(-) diff --git a/kumo/connector/timestamp_convert.py b/kumo/connector/timestamp_convert.py index 861469532..2e68162ef 100644 --- a/kumo/connector/timestamp_convert.py +++ b/kumo/connector/timestamp_convert.py @@ -158,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 @@ -174,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. @@ -183,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 @@ -208,12 +218,10 @@ def to_epoch_seconds(series: pd.Series, hint: TemporalHint) -> np.ndarray: 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. Typed columns on the seconds clock take a vectorized path - that returns exactly what :func:`_resolve` does. + index clock. """ - native = hint.kind is HintKind.NATIVE or is_datetime64_any_dtype(series) - if native and timestamp_precision == 'seconds': - return timestamp_secs_from_array(series) + if hint.kind is HintKind.NATIVE or is_datetime64_any_dtype(series): + return timestamp_ticks_from_array(series, timestamp_precision) return _resolve(series, hint, days=False, timestamp_precision=timestamp_precision) @@ -224,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/test/connector/local/test_timestamp_convert.py b/test/connector/local/test_timestamp_convert.py index d96823267..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]'))) == [] @@ -457,8 +457,10 @@ def test_hint_constructors_shape(): '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]'), @@ -466,28 +468,28 @@ def test_hint_constructors_shape(): _TYPED, _TYPED.dt.tz_localize('US/Pacific'), ], ids=['s', 'ms', 'us', 'ns', 'tz']) -def test_typed_columns_on_the_seconds_clock_match_diskgraph(series): +def test_typed_columns_convert_like_diskgraph_in_every_clock(series, clock): hint = TemporalHint.native() - assert list(tc.to_epoch_ticks(series, hint, 'seconds')) == list( - tc._resolve(series, hint, days=False, timestamp_precision='seconds')) + assert list(tc.to_epoch_ticks(series, hint, clock)) == list( + tc._resolve(series, hint, days=False, timestamp_precision=clock)) -def test_statistics_convert_in_the_snapshot_clock(monkeypatch): +@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 - clocks = [] - - def resolve_temporal(values, fmt, unit, as_days, *, timestamp_precision): - clocks.append(timestamp_precision) - return [1_700_000_000_123], fmt, unit - - monkeypatch.setattr(tc, 'resolve_temporal', resolve_temporal) config = BaseColumnConfig(name='ts', stype=Stype.timestamp) coerced = coerce_timestamp_stype_series_for_stats( - pd.Series(['2023-11-14 22:13:20.123']), config, 'milliseconds') + pd.Series(['2023-11-14 22:13:20.123']), config, clock) - assert clocks == ['milliseconds'] - assert coerced.tolist() == [pd.Timestamp('2023-11-14 22:13:20.123')] + assert coerced.tolist() == [pd.Timestamp(expected)]