fix(array): accept timestamps beyond jiff's range

encodings/datetime-parts/src/compress.rs

@@ -6,6 +6,9 @@ use std::mem::MaybeUninit;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
+use vortex_array::aggregate_fn::NumericalAggregateOpts;
+use vortex_array::aggregate_fn::fns::min_max::MinMaxResult;
+use vortex_array::aggregate_fn::fns::min_max::min_max;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::TemporalArray;
@@ -18,9 +21,9 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
 use crate::timestamp;
-/// All parts are stored as i32. seconds and subseconds always fit in i32.
-/// For days we have a upper day limitation due to "jiff" validation so days
-/// also fit in i32.
+
+/// All parts are stored as i32. Second and millisecond timestamps can lie more than `i32::MAX`
+/// days from the epoch, so [`split_temporal`] checks the days.
 pub struct TemporalParts {
     pub days: ArrayRef,
     pub seconds: ArrayRef,
@@ -31,11 +34,21 @@ pub struct TemporalParts {
 ///
 /// Splitting the components by granularity creates more small values, which enables better
 /// cascading compression.
+///
+/// # Errors
+///
+/// Returns an error if the time unit is days, or if [`can_split_temporal`] is false.
 pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexResult<TemporalParts> {
     let time_unit = array.temporal_metadata().time_unit();
     if matches!(time_unit, TimeUnit::Days) {
         vortex_bail!("Cannot handle day-level data");
     }
+    if !days_fit(array.temporal_values(), time_unit, ctx)? {
+        vortex_bail!(
+            "{time_unit} timestamps lie more than i32::MAX days from the epoch and cannot be \
+             split into days"
+        );
+    }
 
     let temporal_values = array
         .temporal_values()
@@ -109,6 +122,31 @@ pub fn split_temporal(array: TemporalArray, ctx: &mut ExecutionCtx) -> VortexRes
     })
 }
 
+/// Whether every valid timestamp's days since the epoch fit the `i32` days part.
+pub fn can_split_temporal(array: &TemporalArray, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
+    let time_unit = array.temporal_metadata().time_unit();
+    if matches!(time_unit, TimeUnit::Days) {
+        return Ok(false);
+    }
+    days_fit(array.temporal_values(), time_unit, ctx)
+}
+
+fn days_fit(values: &ArrayRef, time_unit: TimeUnit, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
+    if matches!(time_unit, TimeUnit::Microseconds | TimeUnit::Nanoseconds) {
+        return Ok(true);
+    }
+    let Some(MinMaxResult { min, max }) = min_max(values, ctx, NumericalAggregateOpts::default())?
+    else {
+        return Ok(true);
+    };
+    Ok([min, max].iter().all(|bound| {
+        bound
+            .as_primitive()
+            .as_::<i64>()
+            .is_some_and(|ts| timestamp::days_fit(ts, time_unit))
+    }))
+}
+
 #[inline]
 fn split_slice<const DIVISOR: i64>(
     days: &mut [MaybeUninit<i32>],
@@ -177,13 +215,17 @@ mod tests {
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::TemporalArray;
+    use vortex_array::assert_arrays_eq;
     use vortex_array::extension::datetime::TimeUnit;
     use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
     use vortex_error::VortexExpect;
+    use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
+    use crate::DateTimeParts;
     use crate::TemporalParts;
+    use crate::can_split_temporal;
     use crate::split_temporal;
 
     static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
@@ -235,4 +277,35 @@ mod tests {
             Validity::NonNullable
         ));
     }
+
+    #[rstest]
+    #[case::seconds(TimeUnit::Seconds, false)]
+    #[case::milliseconds(TimeUnit::Milliseconds, false)]
+    #[case::microseconds(TimeUnit::Microseconds, true)]
+    #[case::nanoseconds(TimeUnit::Nanoseconds, true)]
+    fn extreme_timestamps_split_exactly_or_fail(
+        #[case] unit: TimeUnit,
+        #[case] splittable: bool,
+    ) -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let temporal = TemporalArray::new_timestamp(
+            PrimitiveArray::from_option_iter([Some(i64::MAX), None, Some(i64::MIN), Some(0)])
+                .into_array(),
+            unit,
+            Some("UTC".into()),
+        );
+        assert_eq!(can_split_temporal(&temporal, &mut ctx)?, splittable);
+
+        match DateTimeParts::try_from_temporal(temporal.clone(), &mut ctx) {
+            Ok(parts) => {
+                assert_arrays_eq!(parts, temporal, &mut ctx);
+                assert!(splittable, "{unit} must be rejected");
+            }
+            Err(err) => {
+                assert!(!splittable, "{unit} must split: {err}");
+                assert!(err.to_string().contains("i32::MAX days"), "{err}");
+            }
+        }
+        Ok(())
+    }
 }

encodings/datetime-parts/src/compute/compare.rs

@@ -52,7 +52,10 @@ impl CompareKernel for DateTimeParts {
         let Some(options) = ext_dtype.metadata_opt::<Timestamp>() else {
             return Ok(None);
         };
-        let ts_parts = timestamp::split(timestamp, options.unit)?;
+        // The constant has no i32 days part.
+        let Ok(ts_parts) = timestamp::split(timestamp, options.unit) else {
+            return Ok(None);
+        };
 
         let result = match operator {
             CompareOperator::Eq => compare_eq(lhs, &ts_parts, nullability, ctx)?,
@@ -443,8 +446,7 @@ mod test {
         )
         .unwrap();
 
-        // Timestamp with a value larger than i32::MAX.
-        let rhs = dtp_array_from_timestamp(i64::MAX, rhs_validity);
+        let rhs = dtp_array_from_timestamp(i64::from(i32::MAX) + 1, rhs_validity);
 
         let comparison = lhs
             .clone()

encodings/datetime-parts/src/compute/rules.rs

@@ -13,12 +13,12 @@ use vortex_array::arrays::scalar_fn::AnyScalarFn;
 use vortex_array::arrays::scalar_fn::ScalarFn;
 use vortex_array::arrays::scalar_fn::ScalarFnArrayExt;
 use vortex_array::arrays::slice::SliceReduceAdaptor;
-use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::DType;
 use vortex_array::extension::datetime::Timestamp;
 use vortex_array::optimizer::ArrayOptimizer;
 use vortex_array::optimizer::rules::ArrayParentReduceRule;
 use vortex_array::optimizer::rules::ParentRuleSet;
+use vortex_array::scalar::Scalar;
 use vortex_array::scalar_fn::fns::between::Between;
 use vortex_array::scalar_fn::fns::binary::Binary;
 use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
@@ -138,8 +138,11 @@ impl ArrayParentReduceRule<DateTimeParts> for DTPComparisonPushDownRule {
                 // Keep the nullability of the replaced constant, so the rewritten comparison has
                 // the same result dtype as the original.
                 let target_dtype = days.dtype().with_nullability(c.dtype().nullability());
-                let constant = ConstantArray::new(days_value, len).into_array();
-                new_children.push(constant.cast(target_dtype)?);
+                // Narrowed days, such as u16, cannot hold every literal's days.
+                let Ok(days_scalar) = Scalar::from(days_value).cast(&target_dtype) else {
+                    return Ok(None);
+                };
+                new_children.push(ConstantArray::new(days_scalar, len).into_array());
             }
         }
 
@@ -190,6 +193,7 @@ fn is_constant_zero(array: &ArrayRef) -> bool {
 mod tests {
     use std::sync::LazyLock;
 
+    use rstest::rstest;
     use vortex_array::ArrayRef;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::BoolArray;
@@ -278,6 +282,42 @@ mod tests {
         ConstantArray::new(scalar, len).into_array()
     }
 
+    #[test]
+    fn days_pushdown_declines_constant_outside_i32_days() {
+        let constant = midnight_constant(i64::from(i32::MAX) + 1, TimeUnit::Seconds, 1);
+        assert_eq!(try_extract_days_constant(&constant), None);
+    }
+
+    #[rstest]
+    #[case::lt(Operator::Lt, true)]
+    #[case::lte(Operator::Lte, true)]
+    #[case::eq(Operator::Eq, false)]
+    #[case::not_eq(Operator::NotEq, true)]
+    #[case::gt(Operator::Gt, false)]
+    fn test_midnight_comparison_beyond_narrowed_days(#[case] op: Operator, #[case] expected: bool) {
+        let dtype = dtp_at_midnight(&[0], TimeUnit::Microseconds)
+            .dtype()
+            .clone();
+        let narrowed = DateTimeParts::try_new(
+            dtype,
+            PrimitiveArray::from_iter([20_000u16, 20_001]).into_array(),
+            ConstantArray::new(0u8, 2).into_array(),
+            ConstantArray::new(0u8, 2).into_array(),
+        )
+        .unwrap();
+        let constant = midnight_constant(2_932_896, TimeUnit::Microseconds, 2);
+        let comparison = Binary::try_new(narrowed.into_array(), constant, op)
+            .unwrap()
+            .into_array();
+
+        let mut ctx = SESSION.create_execution_ctx();
+        assert_arrays_eq!(
+            comparison.optimize().unwrap(),
+            BoolArray::from_iter([expected, expected]),
+            &mut ctx
+        );
+    }
+
     #[test]
     fn test_midnight_binary_comparison() {
         // DTP with days [0, 1, 2] at midnight

encodings/datetime-parts/src/timestamp.rs

@@ -12,7 +12,7 @@ pub const SECONDS_PER_DAY: i64 = 86_400; // 24 * 60 * 60
 /// Parts of a Unix timestamp (time since 1970-01-01 00:00:00 UTC).
 ///
 /// Broken down into:
-///  * Days since epoch (max range is "jiff"'s range)
+///  * Days since epoch (must fit `i32`, see [`days_fit`])
 ///  * Seconds within day (0-86399)
 ///  * Subseconds (range depends on `TimeUnit`)
 pub struct TimestampParts {
@@ -30,18 +30,39 @@ pub struct TimestampParts {
 ///
 /// # Errors
 ///
-/// Returns an error if `time_unit` is days, which cannot be split.
+/// Returns an error if `time_unit` is days, or if [`days_fit`] is false.
 pub fn split(timestamp: i64, time_unit: TimeUnit) -> VortexResult<TimestampParts> {
-    Ok(match time_unit {
+    let parts = match time_unit {
         TimeUnit::Nanoseconds => split_with_divisor::<1_000_000_000>(timestamp),
         TimeUnit::Microseconds => split_with_divisor::<1_000_000>(timestamp),
         TimeUnit::Milliseconds => split_with_divisor::<1_000>(timestamp),
         TimeUnit::Seconds => split_with_divisor::<1>(timestamp),
         TimeUnit::Days => vortex_bail!("Cannot handle day-level data"),
-    })
+    };
+    if !days_fit(timestamp, time_unit) {
+        vortex_bail!(
+            "Timestamp {timestamp}{time_unit} lies more than i32::MAX days from the epoch and \
+             cannot be split into days"
+        );
+    }
+    Ok(parts)
+}
+
+/// Whether the days since the epoch of `timestamp` fit the `i32` days part of [`TimestampParts`].
+pub fn days_fit(timestamp: i64, time_unit: TimeUnit) -> bool {
+    let ticks_per_second = match time_unit {
+        TimeUnit::Nanoseconds => 1_000_000_000,
+        TimeUnit::Microseconds => 1_000_000,
+        TimeUnit::Milliseconds => 1_000,
+        TimeUnit::Seconds => 1,
+        TimeUnit::Days => return false,
+    };
+    i32::try_from(timestamp / (SECONDS_PER_DAY * ticks_per_second)).is_ok()
 }
 
 /// Split a Unix timestamp into parts using a constant DIVISOR.
+///
+/// Days that do not fit `i32` wrap, so callers check [`days_fit`] first.
 #[inline]
 pub(crate) fn split_with_divisor<const DIVISOR: i64>(timestamp: i64) -> TimestampParts {
     let days = timestamp / (SECONDS_PER_DAY * DIVISOR);
@@ -161,6 +182,21 @@ mod tests {
         assert!(split(1, TimeUnit::Days).is_err());
     }
 
+    #[rstest]
+    #[case::seconds(TimeUnit::Seconds, false)]
+    #[case::milliseconds(TimeUnit::Milliseconds, false)]
+    #[case::microseconds(TimeUnit::Microseconds, true)]
+    #[case::nanoseconds(TimeUnit::Nanoseconds, true)]
+    fn split_rejects_days_outside_i32(#[case] unit: TimeUnit, #[case] fits: bool) {
+        for ts in [i64::MAX, i64::MIN] {
+            assert_eq!(days_fit(ts, unit), fits, "{ts}{unit}");
+            match split(ts, unit) {
+                Ok(parts) => assert_eq!(combine(parts, unit), ts, "{ts}{unit}"),
+                Err(err) => assert!(!fits && err.to_string().contains("i32::MAX days"), "{err}"),
+            }
+        }
+    }
+
     #[test]
     fn test_combine_seconds() {
         // 1970-01-02 01:02:03 UTC

vortex-array/src/aggregate_fn/fns/min_max/extension.rs

@@ -11,6 +11,7 @@ use crate::aggregate_fn::AggregateArgs;
 use crate::aggregate_fn::NumericalAggregateOpts;
 use crate::arrays::ExtensionArray;
 use crate::arrays::extension::ExtensionArrayExt;
+use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::scalar::Scalar;
 
@@ -20,15 +21,17 @@ pub(super) fn accumulate_extension(
     array: &ExtensionArray,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<()> {
-    let non_nullable_ext_dtype = array.ext_dtype().with_nullability(Nullability::NonNullable);
+    let dtype = DType::Extension(array.ext_dtype().with_nullability(Nullability::NonNullable));
     // A NaN in float storage comes back as NaN extrema, which `merge` skips or poisons with.
-    let local =
-        min_max(array.storage_array(), ctx, *args.options)?.map(|MinMaxResult { min, max }| {
-            MinMaxResult {
-                min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min),
-                max: Scalar::extension_ref(non_nullable_ext_dtype, max),
-            }
-        });
+    let local = min_max(array.storage_array(), ctx, *args.options)?
+        .map(|MinMaxResult { min, max }| {
+            // Storage extremes the extension type rejects are an error, not an invariant violation.
+            VortexResult::Ok(MinMaxResult {
+                min: Scalar::try_new(dtype.clone(), min.into_value())?,
+                max: Scalar::try_new(dtype, max.into_value())?,
+            })
+        })
+        .transpose()?;
     partial.merge(args, local);
     Ok(())
 }

vortex-array/src/aggregate_fn/fns/min_max/mod.rs

@@ -567,6 +567,7 @@ mod tests {
     use crate::arrays::ListArray;
     use crate::arrays::NullArray;
     use crate::arrays::PrimitiveArray;
+    use crate::arrays::TemporalArray;
     use crate::arrays::VarBinArray;
     use crate::dtype::BigCast;
     use crate::dtype::DType;
@@ -580,6 +581,7 @@ mod tests {
     use crate::expr::stats::Precision;
     use crate::expr::stats::Stat;
     use crate::expr::stats::StatsProvider;
+    use crate::extension::datetime::TimeUnit;
     use crate::match_each_decimal_value_type;
     use crate::scalar::DecimalValue;
     use crate::scalar::Scalar;
@@ -602,6 +604,28 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_timestamp_min_max_beyond_jiff_range() -> VortexResult<()> {
+        // 9999-12-31T00:00:00Z
+        const END_OF_9999_US: i64 = 253_402_214_400_000_000;
+        let timestamps = TemporalArray::new_timestamp(
+            PrimitiveArray::from_iter([0i64, END_OF_9999_US]).into_array(),
+            TimeUnit::Microseconds,
+            None,
+        )
+        .into_array();
+        let mut ctx = SESSION.create_execution_ctx();
+        let MinMaxResult { min, max } =
+            min_max(&timestamps, &mut ctx, NumericalAggregateOpts::default())?
+                .vortex_expect("non-empty array");
+        assert_eq!(min.as_extension().to_storage_scalar(), 0i64.into());
+        assert_eq!(
+            max.as_extension().to_storage_scalar(),
+            END_OF_9999_US.into()
+        );
+        Ok(())
+    }
+
     #[test]
     fn test_prim_min_max_multiple_null_runs() -> VortexResult<()> {
         // Several disjoint valid runs separated by nulls exercise the per-run fold; the extrema