diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index d84f4b6be91..44406e00226 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/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 { 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 { + 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 { + 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_::() + .is_some_and(|ts| timestamp::days_fit(ts, time_unit)) + })) +} + #[inline] fn split_slice( days: &mut [MaybeUninit], @@ -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 = 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(()) + } } diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 96885800424..9e997db8c10 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -52,7 +52,10 @@ impl CompareKernel for DateTimeParts { let Some(options) = ext_dtype.metadata_opt::() 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() diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index 7ebe8477b00..e45926ea105 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/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 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 diff --git a/encodings/datetime-parts/src/timestamp.rs b/encodings/datetime-parts/src/timestamp.rs index a9e9bcc3322..6076bd49f0e 100644 --- a/encodings/datetime-parts/src/timestamp.rs +++ b/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 { - 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(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 diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index c740a4b70f5..033f01e870a 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java @@ -165,6 +165,10 @@ public VortexWriter build() throws IOException { } return new VortexWriter(pointer); } finally { + // Native code only borrows the schema, so the children the export allocated are ours to release. + if (ffi.snapshot().release != 0) { + ffi.release(); + } ffi.close(); } } diff --git a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java index 255201de6bb..d39b9f03904 100644 --- a/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java +++ b/java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java @@ -30,6 +30,7 @@ import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.VarBinaryVector; @@ -169,6 +170,18 @@ public void testCreateWriter() throws IOException { source.arrowSchema(allocator)); } + @Test + public void testCreateWriterReleasesSchemaExport() throws IOException { + String writePath = tempDir.resolve("test_schema_release.vortex").toUri().toString(); + Session session = Session.create(); + try (BufferAllocator allocator = new RootAllocator()) { + try (VortexWriter writer = VortexWriter.builder(session, writePath, personSchema(), allocator) + .build()) { + assertEquals(0, allocator.getAllocatedMemory()); + } + } + } + @Test public void testCreateWriterPlainLocalPath() throws IOException { Path outputPath = tempDir.resolve("test_create_plain_path.vortex"); diff --git a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs index 4c18fad0e6d..658d374604d 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/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(()) } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index ce3e938a8fd..707c1aa024b 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/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(×tamps, &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 diff --git a/vortex-array/src/extension/datetime/timestamp.rs b/vortex-array/src/extension/datetime/timestamp.rs index 42f00776216..5908d3f23d7 100644 --- a/vortex-array/src/extension/datetime/timestamp.rs +++ b/vortex-array/src/extension/datetime/timestamp.rs @@ -6,7 +6,6 @@ use std::fmt; use std::sync::Arc; -use jiff::Span; use jiff::fmt::temporal::DateTimeParser; use jiff::tz::TimeZone; use vortex_error::VortexExpect; @@ -104,13 +103,19 @@ pub enum TimestampValue<'a> { impl fmt::Display for TimestampValue<'_> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let (span, tz) = match self { - TimestampValue::Seconds(v, tz) => (Span::new().seconds(*v), *tz), - TimestampValue::Milliseconds(v, tz) => (Span::new().milliseconds(*v), *tz), - TimestampValue::Microseconds(v, tz) => (Span::new().microseconds(*v), *tz), - TimestampValue::Nanoseconds(v, tz) => (Span::new().nanoseconds(*v), *tz), + let (unit, value, tz) = match *self { + TimestampValue::Seconds(v, tz) => (TimeUnit::Seconds, v, tz), + TimestampValue::Milliseconds(v, tz) => (TimeUnit::Milliseconds, v, tz), + TimestampValue::Microseconds(v, tz) => (TimeUnit::Microseconds, v, tz), + TimestampValue::Nanoseconds(v, tz) => (TimeUnit::Nanoseconds, v, tz), + }; + // jiff cannot render the last 26 hours of year 9999, so show the raw count. + let Ok(ts) = unit + .to_jiff_span(value) + .and_then(|span| Ok(jiff::Timestamp::UNIX_EPOCH.checked_add(span)?)) + else { + return write!(f, "{value}{unit}"); }; - let ts = jiff::Timestamp::UNIX_EPOCH + span; match tz { None => write!(f, "{ts}"), @@ -211,39 +216,19 @@ impl ExtVTable for Timestamp { let ts_value = storage_value.as_primitive().cast::()?; let tz = metadata.tz.as_ref(); - let (span, value) = match metadata.unit { - TimeUnit::Nanoseconds => ( - Span::new().nanoseconds(ts_value), - TimestampValue::Nanoseconds(ts_value, tz), - ), - TimeUnit::Microseconds => ( - Span::new().microseconds(ts_value), - TimestampValue::Microseconds(ts_value, tz), - ), - TimeUnit::Milliseconds => ( - Span::new().milliseconds(ts_value), - TimestampValue::Milliseconds(ts_value, tz), - ), - TimeUnit::Seconds => ( - Span::new().seconds(ts_value), - TimestampValue::Seconds(ts_value, tz), - ), - TimeUnit::Days => { - vortex_bail!("Timestamp does not support Days time unit") - } - }; - - // Validate the storage value is within the valid range for Timestamp. - jiff::Timestamp::UNIX_EPOCH - .checked_add(span) - .map_err(|e| vortex_err!("Invalid timestamp scalar: {}", e))?; - if let Some(tz) = tz { resolve_time_zone(Some(tz)) .map_err(|e| e.with_context("Invalid timezone for timestamp scalar"))?; } - Ok(value) + // Any i64 count is valid, as in Arrow; jiff's range only limits Display. + Ok(match metadata.unit { + TimeUnit::Nanoseconds => TimestampValue::Nanoseconds(ts_value, tz), + TimeUnit::Microseconds => TimestampValue::Microseconds(ts_value, tz), + TimeUnit::Milliseconds => TimestampValue::Milliseconds(ts_value, tz), + TimeUnit::Seconds => TimestampValue::Seconds(ts_value, tz), + TimeUnit::Days => vortex_bail!("Timestamp does not support Days time unit"), + }) } } @@ -270,6 +255,28 @@ mod tests { Ok(()) } + #[cfg_attr(miri, ignore)] + #[rstest] + #[case::end_of_9999( + TimeUnit::Microseconds, + 253_402_214_400_000_000, + "253402214400000000µs" + )] + #[case::start_of_neg_9999(TimeUnit::Seconds, -377_705_116_800, "-377705116800s")] + #[case::max(TimeUnit::Microseconds, i64::MAX, "9223372036854775807µs")] + fn accept_timestamp_beyond_jiff_range( + #[case] unit: TimeUnit, + #[case] value: i64, + #[case] expected: &str, + ) -> VortexResult<()> { + for tz in [None, Some(Arc::from("UTC"))] { + let dtype = DType::Extension(Timestamp::new_with_tz(unit, tz, Nullable).erased()); + let scalar = Scalar::try_new(dtype, Some(ScalarValue::Primitive(PValue::I64(value))))?; + assert_eq!(format!("{}", scalar.as_extension()), expected); + } + Ok(()) + } + #[cfg_attr(miri, ignore)] #[rstest] #[case("+01:00", "1970-01-01T01:00:00+01:00[+01:00]")] diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 79748b69450..010f21e0c1d 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -21,6 +21,7 @@ use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_datetime_parts::DateTimeParts; use vortex_datetime_parts::TemporalParts; +use vortex_datetime_parts::can_split_temporal; use vortex_datetime_parts::split_temporal; use vortex_error::VortexResult; @@ -66,12 +67,17 @@ impl Scheme for TemporalScheme { fn expected_compression_ratio( &self, - _data: &ArrayAndStats, + data: &ArrayAndStats, _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, + exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - // Temporal compression (splitting into parts) is almost always beneficial. - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + // Second and millisecond timestamps past i32::MAX days cannot be split. + let splittable = TemporalArray::try_from(data.array().clone()) + .and_then(|temporal| can_split_temporal(&temporal, exec_ctx)); + match splittable { + Ok(true) => CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse), + _ => CompressionEstimate::Verdict(EstimateVerdict::Skip), + } } fn compress( diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index 03250fc4491..384e104008a 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -32,6 +32,7 @@ use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; +use vortex_datetime_parts::DateTimeParts; #[cfg(feature = "zstd")] use vortex_edition::EDITION_DECLARATIONS; #[cfg(feature = "zstd")] @@ -111,6 +112,27 @@ fn test_default_compressor_roundtrip() -> VortexResult<()> { Ok(()) } +#[test] +fn test_extreme_second_timestamps_roundtrip() -> VortexResult<()> { + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + let input = TemporalArray::new_timestamp( + PrimitiveArray::from_iter((0..2048i64).map(|i| match i % 3 { + 0 => i64::MAX - i, + 1 => i64::MIN + i, + _ => 1_700_000_000 + i, + })) + .into_array(), + TimeUnit::Seconds, + None, + ) + .into_array(); + let compressed = assert_roundtrip(&compressor, &input)?; + assert!(!compressed.is::()); + Ok(()) +} + #[rstest] #[case::unaligned({ let mut rng = StdRng::seed_from_u64(7); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 30547d298eb..e0d2594ec41 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -320,6 +320,17 @@ fn expand_struct_root( .first() .is_some_and(BoundExpression::is_root) { + // Masking `$.x` inside `$.x.q` would hide the path from the partitioner, which + // would then read all of `x`. Mask the whole path instead. + if validity.is_some() + && let Some(path) = field_path(&node, packed, fields)? + { + return Ok(Transformed { + value: with_validity(path), + changed: true, + order: TraversalOrder::Skip, + }); + } return Ok(Transformed::no(node)); } @@ -380,6 +391,35 @@ fn expand_struct_root( .into_inner()) } +/// Rebuilds a chain of `get_item` and `select` that starts at `$.field` on the field's expression +/// in `packed`, or returns `None` if `node` is not such a chain. +fn field_path( + node: &BoundExpression, + packed: &BoundExpression, + fields: &StructFields, +) -> VortexResult> { + let Some(scalar_fn) = node.as_scalar() else { + return Ok(None); + }; + if !scalar_fn.is::() && !scalar_fn.is::