From c595fee35ca7c3ed4c9e85280ee8d2b0c742d626 Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 16:18:27 -0700 Subject: [PATCH 01/12] docs: confirm the Arrow Map alias patch is superseded, not just droppable Traced through vortex-arrow's current Map handling instead of trusting the merge-conflict shape: upstream 0.85.0 replaced the List> alias entirely with a native Map array/dtype (vortex-array/src/dtype/map.rs, vortex-array/src/arrays/map/), and both the import and export paths already recurse into map key/value fields for extension metadata the way Spice's 3e7fa40d3 did. Ran the vortex-arrow map test suite on this branch (17 passing, including the exact nested-extension-type scenario 3e7fa40d3 fixed) rather than asserting this from the diff alone. Moves row 14 from "needs re-porting" to "dropped" with that run as evidence. --- SPICE_PATCHES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 7a20b34e675..fb67b25b8fc 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -67,6 +67,7 @@ content in favor of this rewrite, not by a code-level resolution. The actual cod | — | `ea3890528` "duckdb 1.5.5 (#8895)" | — | Not a Spice patch: identical to upstream's own `0.80.0` release commit (`037b1a8778db`), already superseded by landing on `0.85.0`. | | — | `68badd689` "TableStrategy hides panic... (#8672)" | — | Not a Spice patch: `(cherry picked from commit 4abe3d04f...)`, authored by a Vortex/spiraldb maintainer, cherry-picked onto `spiceai-54` ahead of a release. Already upstream. | | — | `b5498e8e1` "get CI green" (#81), `8d2efe7a7` style import sort | — | CI/lint housekeeping for subsystems Spice doesn't build (java/, python/, ffi/) or trivial reformatting of files already changed upstream. Not carried; not fork state Spice depends on. | +| 14 | Arrow Map alias — `DType` alias + map-entry recursion in the session importer | `840e746a9`, `3e7fa40d3` (superseded the stale `1a6dc54f1` origin, then this) | **Superseded by a full upstream rewrite, not just a fix.** Upstream 0.85.0 replaced the `List>` alias both Spice's patch and the pre-split `vortex-array/src/arrow` code used with a **native `Map` array/dtype**: `vortex-array/src/dtype/map.rs` (`MapDType`) and `vortex-array/src/arrays/map/` (`Array`, backed by a `ListView>`). Both `vortex-arrow/src/convert.rs` (import) and `vortex-arrow/src/executor/mod.rs` (export) dispatch `DataType::Map` to this native path, and `vortex-arrow/src/session.rs` recurses into map key/value fields for extension-type metadata exactly as Spice's `3e7fa40d3` did. **Ran `cargo test -p vortex-arrow --lib map` on this branch: 17 tests pass**, including `executor::map::tests::map_roundtrip_preserves_nested_uuid_fields` (the exact extension-recursion scenario `3e7fa40d3` fixed) and `session::tests::schema_roundtrip_preserves_map_uuid_fields`. Upstream's own coverage here is broader than Spice's two patches were. Attempting to cherry-pick the old alias-based patch onto this would be a regression, not a port — the P0 defect (spiceai/spiceai#13524) is closed by a better mechanism than the one Spice shipped. **Remaining gap**: this proves the fix at the `vortex-arrow` crate level; the Spice main repo's own guard (`crates/vortex/src/persistent/mod.rs::map_column_roundtrips_through_a_vortex_file`, per `docs/dev/fork_patches.md`) still needs re-running once this branch is pinned, since that test exercises the full write path through the vendored `vortex-datafusion` sink, which this fork does not build. | ## Needs re-porting (real conflicts against `0.85.0` — deferred, not dropped) @@ -86,7 +87,6 @@ reading the new upstream code and re-implementing the fix, not resolving conflic | — | Restore `load_full` + `HashMap` import "dropped in rebase" | `ab4f0b177` | `vortex-array/src/arc_swap_map.rs` is deleted at `0.85.0` (modify/delete conflict) — check whether the underlying function this restored still has a reason to exist before re-porting; may be moot. Not previously in the ledger. | | — | Migrate extension date-cast + Arrow Map test to ctx execute APIs | `d41e094cf` | Conflicts in `cast.rs` (same `ArrayView` signature issue as row 7) and `vortex-arrow/src/convert.rs`. Bundle with the row 7/15/16 re-port and the Map-alias re-port below — touches both areas. Not previously in the ledger. | | — | Restore lint checks on forks | `bb80c537b` | Conflicts across `.github/workflows/ci.yml`, `cast.rs`, `vortex-datafusion/src/persistent/sink.rs` (the last is out of scope). Low priority — CI config for this fork, not behavior Spice depends on. | -| 14 (current fix) | **Arrow Map alias restoration** — `DType` alias + map-entry recursion in the session importer | `840e746a9`, `3e7fa40d3` (2026-08-26; supersedes the stale `1a6dc54f1` the old ledger pointed at) | Real conflicts in `vortex-arrow/src/convert.rs`, `vortex-arrow/src/executor/{map,mod}.rs`, `vortex-arrow/src/session.rs` (the `vortex-arrow` crate itself is new since `1a6dc54f1`'s time, from the `vortex-array/src/arrow` crate split). **Highest priority in this list** — this is the exact defect from spiceai/spiceai#13524 (every `Map`-typed write fails on flush), and it has already been silently lost from a re-cut once before. Do not re-cut past this branch without landing it. | ## Untracked in the previous ledger — needs its own audit From a98a78c1b090de909e1fe1c94e90107881d43e4d Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 16:26:57 -0700 Subject: [PATCH 02/12] cast: re-port vortex.date -> vortex.timestamp extension casts (#28) to 0.85.0 Ported forward from spiceai-54 (7e5b08151), which doesn't apply as a cherry-pick: CastReduce::cast moved from `fn cast(array: &ExtensionArray, ...)` to `fn cast(array: ArrayView<'_, Extension>, ...)`, and the "different extension type" branch now eagerly delegates to the storage array's generic cast instead of falling through. Confirmed the regression is real before porting, not assumed: casting a Date(Days) array to Timestamp(Nanoseconds) on this branch failed with "No CastKernel to cast canonical array vortex.primitive from i32 to vortex.timestamp[ns](i64)" (see the removed probe in this file's history). Without this patch a pushed-down `CAST(date_col AS TIMESTAMP)` fails the scan outright (spiceai/spiceai#13624's array half). Validity resolution needed its own adaptation: PrimitiveArray::validity_mask() no longer exists (validity may now be lazy and needs an ExecutionCtx to resolve), so this uses `legacy_session()` -- the crate's own sanctioned shim for "code paths that cannot yet thread an explicit session through" -- the same pattern Validity::is_valid/is_null already use. Carries forward both of the original patch's tests unchanged in behavior, confirmed passing against 0.85.0: cast_date_days_to_timestamp_nanoseconds and cast_date_days_to_timestamp_seconds_nullable, plus all 8 pre-existing extension-cast tests still pass. --- .../src/arrays/extension/compute/cast.rs | 212 ++++++++++++++++++ 1 file changed, 212 insertions(+) diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index e080eb7f04c..8db0493600c 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -1,21 +1,47 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::AllOr; use crate::ArrayRef; use crate::IntoArray; +use crate::ToCanonical; use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; +use crate::arrays::PrimitiveArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::PType; +use crate::executor::VortexSessionExecute; +use crate::extension::datetime::AnyTemporal; +use crate::legacy_session; +use crate::extension::datetime::TemporalMetadata; +use crate::extension::datetime::TimeUnit; use crate::scalar_fn::fns::cast::CastReduce; impl CastReduce for Extension { fn cast(array: ArrayView<'_, Extension>, dtype: &DType) -> VortexResult> { if !array.dtype().eq_ignore_nullability(dtype) { + // A date can cast to a timestamp by scaling its storage value, not just + // relabelling the dtype (spiceai/spiceai#13624): the general storage delegate + // below has no `CastKernel` from a date's primitive storage to a timestamp's, + // so without this the cast fails outright rather than converting. + if let Some(new_storage) = cast_temporal_date_to_timestamp(array, dtype)? { + let DType::Extension(ext_dtype) = dtype else { + unreachable!( + "cast_temporal_date_to_timestamp only returns Some for an Extension target" + ); + }; + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); + } + // Target is not the same extension type. // Delegate to the storage array's cast. return Ok(Some(array.storage_array().cast(dtype.clone())?)); @@ -42,6 +68,138 @@ impl CastReduce for Extension { } } +fn cast_temporal_date_to_timestamp( + array: ArrayView<'_, Extension>, + target_dtype: &DType, +) -> VortexResult> { + let DType::Extension(target_ext_dtype) = target_dtype else { + return Ok(None); + }; + + let Some(source_temporal) = array.ext_dtype().metadata_opt::() else { + return Ok(None); + }; + let Some(target_temporal) = target_ext_dtype.metadata_opt::() else { + return Ok(None); + }; + + let (TemporalMetadata::Date(source_unit), TemporalMetadata::Timestamp(target_unit, _)) = + (source_temporal, target_temporal) + else { + return Ok(None); + }; + + let source_i64 = array + .storage_array() + .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?; + let source_i64 = source_i64.to_primitive(); + + let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit)?; + + converted + .into_array() + .cast(target_ext_dtype.storage_dtype().clone()) + .map(Some) +} + +fn cast_date_values_to_timestamp( + values: &PrimitiveArray, + source_unit: TimeUnit, + target_unit: TimeUnit, +) -> VortexResult { + let (multiply, divide) = date_to_timestamp_scale(source_unit, target_unit)?; + + // `execute_mask` needs an `ExecutionCtx`, which this reduce-only cast path doesn't have + // threaded through it. `legacy_session` is this crate's own sanctioned shim for exactly + // that case (its own doc: "code paths that cannot yet thread an explicit session through"). + let validity = values.validity()?; + let mask = validity.execute_mask(values.len(), &mut legacy_session().create_execution_ctx())?; + + let input = values.as_slice::(); + let mut output = BufferMut::with_capacity(input.len()); + match mask.bit_buffer() { + AllOr::All => { + for &value in input { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(convert_temporal_value(value, multiply, divide)?) }; + } + } + AllOr::None => { + for _ in 0..input.len() { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + AllOr::Some(bits) => { + for (&value, valid) in input.iter().zip(bits.iter()) { + if valid { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { + output.push_unchecked(convert_temporal_value(value, multiply, divide)?) + }; + } else { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + } + } + + Ok(PrimitiveArray::new(output.freeze(), validity)) +} + +fn date_to_timestamp_scale( + source_unit: TimeUnit, + target_unit: TimeUnit, +) -> VortexResult<(i64, i64)> { + let source_ns = to_nanoseconds(source_unit)?; + let target_ns = to_nanoseconds(target_unit)?; + + if source_ns >= target_ns { + let multiply = source_ns / target_ns; + return Ok((multiply, 1)); + } + + let divide = target_ns / source_ns; + Ok((1, divide)) +} + +fn to_nanoseconds(unit: TimeUnit) -> VortexResult { + match unit { + TimeUnit::Nanoseconds => Ok(1), + TimeUnit::Microseconds => Ok(1_000), + TimeUnit::Milliseconds => Ok(1_000_000), + TimeUnit::Seconds => Ok(1_000_000_000), + TimeUnit::Days => Ok(86_400_000_000_000), + } +} + +fn convert_temporal_value(value: i64, multiply: i64, divide: i64) -> VortexResult { + let mut scaled = i128::from(value) + .checked_mul(i128::from(multiply)) + .ok_or_else(|| { + vortex_error::vortex_err!( + Compute: "Date value {value} overflows while scaling to timestamp" + ) + })?; + + if divide != 1 { + let divisor = i128::from(divide); + if scaled % divisor != 0 { + vortex_bail!( + Compute: "Date value {value} cannot be represented exactly in target timestamp unit" + ); + } + scaled /= divisor; + } + + if scaled < i128::from(i64::MIN) || scaled > i128::from(i64::MAX) { + vortex_bail!(Compute: "Date value {value} overflows target timestamp range"); + } + + Ok(scaled as i64) +} + #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -61,6 +219,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::executor::VortexSessionExecute; + use crate::extension::datetime::Date; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; @@ -99,6 +258,59 @@ mod tests { assert_eq!(output.dtype(), &new_dtype); } + #[test] + fn cast_date_days_to_timestamp_nanoseconds() { + let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); + let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased(); + + let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array()); + let output = arr + .into_array() + .cast(DType::Extension(target_dtype.clone())) + .unwrap() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + + assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + + let storage = output.storage_array().to_primitive(); + assert_eq!( + storage.as_slice::(), + &[0, 86_400_000_000_000, -86_400_000_000_000] + ); + } + + #[test] + fn cast_date_days_to_timestamp_seconds_nullable() { + let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); + let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased(); + + let arr = ExtensionArray::new( + source_dtype, + PrimitiveArray::from_option_iter([Some(0i32), None, Some(2)]).into_array(), + ); + + let output = arr + .into_array() + .cast(DType::Extension(target_dtype.clone())) + .unwrap() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + + assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + + let storage = output.storage_array().to_primitive(); + assert_eq!( + storage.scalar_at(0).unwrap().as_primitive().as_::(), + Some(0) + ); + assert!(storage.scalar_at(1).unwrap().is_null()); + assert_eq!( + storage.scalar_at(2).unwrap().as_primitive().as_::(), + Some(172_800) + ); + } + #[test] fn cast_different_ext_dtype() { let original_dtype = From dc9892190c9d86b0d48bf831229e020eb6677b5f Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 16:56:09 -0700 Subject: [PATCH 03/12] fix(array): resolve fixed-offset timezones, convert date scalars/arrays to timestamps instead of re-labelling (#75, #78, #93) Re-ports two related fork patch families from spiceai-54 to 0.85.0 as fresh implementations rather than cherry-picks, since neither applies: both touch functions upstream restructured, and the second depends on the first already being in place (resolve_timezone) to even compile against. Fixed-offset timezone resolution (#75, #78): Arrow allows a timestamp's timezone to be an IANA name or a fixed offset (+00:00), but jiff only resolves IANA names, so any offset-form timezone aborted the process via `.vortex_expect("unknown timezone")` -- in scalar validation, Display rendering, and TemporalMetadata::to_jiff. Ports the fork's resolve_timezone (new file, 38 passing tests carried over verbatim) and threads it through all three sites, plus Scalar::try_extension_ref (a fallible counterpart to the panicking extension_ref) through the three call sites that built extension scalars infallibly on a read path (min/max aggregation, scalar_at, the constant-folding reduce rule) and DateTimeParts::scalar_at. Date -> timestamp conversion (#93, spiceai/spiceai#13624): Scalar::cast sent any cast targeting an extension type through the target's storage type and re-labelled the result, discarding what the source meant -- vortex.date and vortex.timestamp both count from the epoch but in different units. A scan falsifies `cast(col as timestamp) > lit` into `cast(max(col) as timestamp) <= lit` and casts a file's statistic scalar to decide whether to read it at all, so this silently pruned files whose rows would have matched. Ports DateToTimestamp (new file, shared by both the array CastKernel and the scalar cast so the two cannot disagree about which files/rows to keep) and Timestamp::storage_range, replacing Jiff Span-based validation whose limits are not a timestamp's at either end (span floor is one nanosecond above i64::MIN, which a real timestamp[ns] column holds). This supersedes a98a78c1b's array-only port on this branch: that port used legacy_session() as a workaround for CastReduce's missing ExecutionCtx. Discovered from an untracked later Spice commit that this fork's own answer was to split into CastReduce (buffer-free restructuring) and a proper CastKernel (which does get a real ctx) instead -- the same split this commit uses, requiring an explicit `register_execute_parent_kernel(Cast.id(), Extension, ...)` that has no upstream equivalent (Extension had no Cast kernel at 0.85.0 at all, only the reduce path). Verified: cargo test -p vortex-array -p vortex-datetime-parts --lib: 3428 + 72 passed, 0 failed. cargo check across vortex-array, vortex-arrow, vortex-layout, vortex-file, vortex-io, vortex-utils, vortex-duckdb: clean, zero warnings. Not ported: the DuckDB TIMESTAMP_TZ dtype/scalar conversion's is_utc_timezone gate (vortex-duckdb/src/convert/{dtype,scalar}.rs) -- vortex-duckdb is not a crate Spice's Cargo.toml consumes, directly or transitively, so this is fork hygiene rather than something Spice's re-pin needs. Left for a follow-up. --- encodings/datetime-parts/src/ops.rs | 10 +- .../src/aggregate_fn/fns/min_max/extension.rs | 14 +- .../src/arrays/extension/compute/cast.rs | 242 ++++++-------- .../src/arrays/extension/compute/rules.rs | 2 +- .../src/arrays/extension/vtable/kernel.rs | 3 + .../src/arrays/extension/vtable/operations.rs | 4 +- vortex-array/src/extension/datetime/cast.rs | 206 ++++++++++++ .../src/extension/datetime/matcher.rs | 39 ++- vortex-array/src/extension/datetime/mod.rs | 4 + .../src/extension/datetime/timestamp.rs | 304 ++++++++++++++++-- .../src/extension/datetime/timezone.rs | 199 ++++++++++++ vortex-array/src/scalar/cast.rs | 12 +- vortex-array/src/scalar/constructor.rs | 16 +- .../src/scalar/typed_view/extension/mod.rs | 59 +++- .../src/scalar/typed_view/extension/tests.rs | 158 +++++++++ 15 files changed, 1071 insertions(+), 201 deletions(-) create mode 100644 vortex-array/src/extension/datetime/cast.rs create mode 100644 vortex-array/src/extension/datetime/timezone.rs diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index 57a1f60173f..ab8fa362cb7 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -65,9 +65,11 @@ impl OperationsVTable for DateTimeParts { options.unit, ); - Ok(Scalar::extension::( - options.clone(), - Scalar::primitive(ts, ext.storage_dtype().nullability()), - )) + // Build the extension scalar fallibly: an extension type whose metadata the storage value + // cannot satisfy (an unresolvable timezone, say) must surface as an error rather than abort + // the process, since `execute_scalar` dispatches here when reading a value out of a + // timestamp column stored under this encoding. + let storage = Scalar::primitive(ts, ext.storage_dtype().nullability()); + Scalar::try_extension_ref(ext, storage) } } 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 e981c7c6af1..73c5618599c 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/extension.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/extension.rs @@ -19,15 +19,21 @@ pub(super) fn accumulate_extension( ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let non_nullable_ext_dtype = array.ext_dtype().with_nullability(Nullability::NonNullable); + // Build the extension scalars fallibly: an extension type whose metadata the storage value + // cannot satisfy (an unresolvable timezone, say) must fail the aggregation rather than abort + // the process, since this runs on the write path when building zone maps. let local = min_max( array.storage_array(), ctx, NumericalAggregateOpts::default(), )? - .map(|MinMaxResult { min, max }| MinMaxResult { - min: Scalar::extension_ref(non_nullable_ext_dtype.clone(), min), - max: Scalar::extension_ref(non_nullable_ext_dtype, max), - }); + .map(|MinMaxResult { min, max }| -> VortexResult { + Ok(MinMaxResult { + min: Scalar::try_extension_ref(non_nullable_ext_dtype.clone(), min)?, + max: Scalar::try_extension_ref(non_nullable_ext_dtype, max)?, + }) + }) + .transpose()?; partial.merge(local); Ok(()) } diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 8db0493600c..1d98d0e011b 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -3,12 +3,10 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_mask::AllOr; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::ToCanonical; use crate::array::ArrayView; use crate::arrays::Extension; use crate::arrays::ExtensionArray; @@ -17,60 +15,73 @@ use crate::arrays::extension::ExtensionArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::PType; -use crate::executor::VortexSessionExecute; use crate::extension::datetime::AnyTemporal; -use crate::legacy_session; +use crate::extension::datetime::DateToTimestamp; use crate::extension::datetime::TemporalMetadata; use crate::extension::datetime::TimeUnit; +use crate::scalar_fn::fns::cast::CastKernel; use crate::scalar_fn::fns::cast::CastReduce; +use crate::validity::Validity; impl CastReduce for Extension { fn cast(array: ArrayView<'_, Extension>, dtype: &DType) -> VortexResult> { - if !array.dtype().eq_ignore_nullability(dtype) { - // A date can cast to a timestamp by scaling its storage value, not just - // relabelling the dtype (spiceai/spiceai#13624): the general storage delegate - // below has no `CastKernel` from a date's primitive storage to a timestamp's, - // so without this the cast fails outright rather than converting. - if let Some(new_storage) = cast_temporal_date_to_timestamp(array, dtype)? { - let DType::Extension(ext_dtype) = dtype else { - unreachable!( - "cast_temporal_date_to_timestamp only returns Some for an Extension target" - ); - }; - return Ok(Some( - ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), - )); - } - - // Target is not the same extension type. + let DType::Extension(ext_dtype) = dtype else { + // Target is not an extension type. // Delegate to the storage array's cast. return Ok(Some(array.storage_array().cast(dtype.clone())?)); + }; + + if array.ext_dtype().eq_ignore_nullability(ext_dtype) { + // Same extension type: restructure by casting the storage array. This is buffer-free + // (`cast` returns a lazy cast expression), so it stays in the reduce phase. + let new_storage = match array + .storage_array() + .cast(ext_dtype.storage_dtype().clone()) + { + Ok(arr) => arr, + Err(e) => { + tracing::warn!("Failed to cast storage array: {e}"); + return Ok(None); + } + }; + + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); } + // A different extension type (e.g. Date -> Timestamp) requires per-value conversion that + // reads buffers, so it is deferred to the execute phase in `CastKernel`, where an + // `ExecutionCtx` is available. + Ok(None) + } +} + +impl CastKernel for Extension { + fn cast( + array: ArrayView<'_, Extension>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { let DType::Extension(ext_dtype) = dtype else { - unreachable!("Already verified we have an extension dtype"); + // Non-extension targets are restructured buffer-free by `CastReduce`. + return Ok(None); }; - let new_storage = match array - .storage_array() - .cast(ext_dtype.storage_dtype().clone()) - { - Ok(arr) => arr, - Err(e) => { - tracing::warn!("Failed to cast storage array: {e}"); - return Ok(None); - } - }; + if let Some(new_storage) = cast_temporal_date_to_timestamp(&array, dtype, ctx)? { + return Ok(Some( + ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), + )); + } - Ok(Some( - ExtensionArray::new(ext_dtype.clone(), new_storage).into_array(), - )) + Ok(None) } } fn cast_temporal_date_to_timestamp( - array: ArrayView<'_, Extension>, + array: &ArrayView<'_, Extension>, target_dtype: &DType, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let DType::Extension(target_ext_dtype) = target_dtype else { return Ok(None); @@ -91,56 +102,58 @@ fn cast_temporal_date_to_timestamp( let source_i64 = array .storage_array() - .cast(DType::Primitive(PType::I64, array.dtype().nullability()))?; - let source_i64 = source_i64.to_primitive(); + .cast(DType::Primitive(PType::I64, array.dtype().nullability()))? + .execute::(ctx)?; - let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit)?; + let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit, ctx)?; - converted + let new_storage = converted .into_array() - .cast(target_ext_dtype.storage_dtype().clone()) - .map(Some) + .cast(target_ext_dtype.storage_dtype().clone())? + .execute::(ctx)? + .into_array(); + + Ok(Some(new_storage)) } fn cast_date_values_to_timestamp( values: &PrimitiveArray, source_unit: TimeUnit, target_unit: TimeUnit, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let (multiply, divide) = date_to_timestamp_scale(source_unit, target_unit)?; - - // `execute_mask` needs an `ExecutionCtx`, which this reduce-only cast path doesn't have - // threaded through it. `legacy_session` is this crate's own sanctioned shim for exactly - // that case (its own doc: "code paths that cannot yet thread an explicit session through"). - let validity = values.validity()?; - let mask = validity.execute_mask(values.len(), &mut legacy_session().create_execution_ctx())?; + let conversion = DateToTimestamp::new(source_unit, target_unit)?; let input = values.as_slice::(); + let validity = values.validity()?; let mut output = BufferMut::with_capacity(input.len()); - match mask.bit_buffer() { - AllOr::All => { + + match &validity { + Validity::NonNullable | Validity::AllValid => { for &value in input { // SAFETY: output has sufficient capacity for all pushed values. - unsafe { output.push_unchecked(convert_temporal_value(value, multiply, divide)?) }; + unsafe { output.push_unchecked(conversion.convert(value)?) }; } } - AllOr::None => { + Validity::AllInvalid => { for _ in 0..input.len() { // SAFETY: output has sufficient capacity for all pushed values. unsafe { output.push_unchecked(0i64) }; } } - AllOr::Some(bits) => { - for (&value, valid) in input.iter().zip(bits.iter()) { - if valid { - // SAFETY: output has sufficient capacity for all pushed values. - unsafe { - output.push_unchecked(convert_temporal_value(value, multiply, divide)?) - }; + Validity::Array(_) => { + // Resolve validity to a boolean mask once. Null slots keep a placeholder 0 so a garbage + // source value in a null slot cannot trip the overflow check in + // `DateToTimestamp::convert`; the output re-uses `validity`, so those slots stay null. + let mask = validity.execute_mask(input.len(), ctx)?; + for (i, &value) in input.iter().enumerate() { + let converted = if mask.value(i) { + conversion.convert(value)? } else { - // SAFETY: output has sufficient capacity for all pushed values. - unsafe { output.push_unchecked(0i64) }; - } + 0i64 + }; + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(converted) }; } } } @@ -148,58 +161,6 @@ fn cast_date_values_to_timestamp( Ok(PrimitiveArray::new(output.freeze(), validity)) } -fn date_to_timestamp_scale( - source_unit: TimeUnit, - target_unit: TimeUnit, -) -> VortexResult<(i64, i64)> { - let source_ns = to_nanoseconds(source_unit)?; - let target_ns = to_nanoseconds(target_unit)?; - - if source_ns >= target_ns { - let multiply = source_ns / target_ns; - return Ok((multiply, 1)); - } - - let divide = target_ns / source_ns; - Ok((1, divide)) -} - -fn to_nanoseconds(unit: TimeUnit) -> VortexResult { - match unit { - TimeUnit::Nanoseconds => Ok(1), - TimeUnit::Microseconds => Ok(1_000), - TimeUnit::Milliseconds => Ok(1_000_000), - TimeUnit::Seconds => Ok(1_000_000_000), - TimeUnit::Days => Ok(86_400_000_000_000), - } -} - -fn convert_temporal_value(value: i64, multiply: i64, divide: i64) -> VortexResult { - let mut scaled = i128::from(value) - .checked_mul(i128::from(multiply)) - .ok_or_else(|| { - vortex_error::vortex_err!( - Compute: "Date value {value} overflows while scaling to timestamp" - ) - })?; - - if divide != 1 { - let divisor = i128::from(divide); - if scaled % divisor != 0 { - vortex_bail!( - Compute: "Date value {value} cannot be represented exactly in target timestamp unit" - ); - } - scaled /= divisor; - } - - if scaled < i128::from(i64::MIN) || scaled > i128::from(i64::MAX) { - vortex_bail!(Compute: "Date value {value} overflows target timestamp range"); - } - - Ok(scaled as i64) -} - #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -259,29 +220,28 @@ mod tests { } #[test] - fn cast_date_days_to_timestamp_nanoseconds() { + fn cast_date_days_to_timestamp_nanoseconds() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let source_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); let target_dtype = Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased(); let arr = ExtensionArray::new(source_dtype, buffer![0i32, 1, -1].into_array()); let output = arr .into_array() - .cast(DType::Extension(target_dtype.clone())) - .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) - .unwrap(); - - assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + .cast(DType::Extension(target_dtype.clone()))?; - let storage = output.storage_array().to_primitive(); - assert_eq!( - storage.as_slice::(), - &[0, 86_400_000_000_000, -86_400_000_000_000] - ); + let expected = ExtensionArray::new( + target_dtype, + buffer![0i64, 86_400_000_000_000, -86_400_000_000_000].into_array(), + ) + .into_array(); + assert_arrays_eq!(output, expected, &mut ctx); + Ok(()) } #[test] - fn cast_date_days_to_timestamp_seconds_nullable() { + fn cast_date_days_to_timestamp_seconds_nullable() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let source_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let target_dtype = Timestamp::new(TimeUnit::Seconds, Nullability::Nullable).erased(); @@ -292,23 +252,15 @@ mod tests { let output = arr .into_array() - .cast(DType::Extension(target_dtype.clone())) - .unwrap() - .execute::(&mut SESSION.create_execution_ctx()) - .unwrap(); - - assert_eq!(output.dtype(), &DType::Extension(target_dtype)); + .cast(DType::Extension(target_dtype.clone()))?; - let storage = output.storage_array().to_primitive(); - assert_eq!( - storage.scalar_at(0).unwrap().as_primitive().as_::(), - Some(0) - ); - assert!(storage.scalar_at(1).unwrap().is_null()); - assert_eq!( - storage.scalar_at(2).unwrap().as_primitive().as_::(), - Some(172_800) - ); + let expected = ExtensionArray::new( + target_dtype, + PrimitiveArray::from_option_iter([Some(0i64), None, Some(172_800)]).into_array(), + ) + .into_array(); + assert_arrays_eq!(output, expected, &mut ctx); + Ok(()) } #[test] diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index d5a0e8d0060..8169fa527d7 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -36,7 +36,7 @@ impl ArrayReduceRule for ExtensionConstantRule { }; let storage_scalar = const_array.scalar().clone(); - let ext_scalar = Scalar::extension_ref(array.ext_dtype().clone(), storage_scalar); + let ext_scalar = Scalar::try_extension_ref(array.ext_dtype().clone(), storage_scalar)?; let constant_with_extension_scalar = ConstantArray::new(ext_scalar, array.len()).into_array(); diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index c90458f1d5f..a81863a3ed6 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -11,6 +11,8 @@ use crate::optimizer::kernels::ArrayKernelsExt; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::cast::CastExecuteAdaptor; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -19,5 +21,6 @@ pub(crate) fn initialize(session: &VortexSession) { Extension, CompareExecuteAdaptor(Extension), ); + kernels.register_execute_parent_kernel(Cast.id(), Extension, CastExecuteAdaptor(Extension)); kernels.register_execute_parent_kernel(Dict.id(), Extension, TakeExecuteAdaptor(Extension)); } diff --git a/vortex-array/src/arrays/extension/vtable/operations.rs b/vortex-array/src/arrays/extension/vtable/operations.rs index 66de94b596a..446e3684ef7 100644 --- a/vortex-array/src/arrays/extension/vtable/operations.rs +++ b/vortex-array/src/arrays/extension/vtable/operations.rs @@ -16,9 +16,9 @@ impl OperationsVTable for Extension { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(Scalar::extension_ref( + Scalar::try_extension_ref( array.ext_dtype().clone(), array.storage_array().execute_scalar(index, ctx)?, - )) + ) } } diff --git a/vortex-array/src/extension/datetime/cast.rs b/vortex-array/src/extension/datetime/cast.rs new file mode 100644 index 00000000000..cf0a8cd772e --- /dev/null +++ b/vortex-array/src/extension/datetime/cast.rs @@ -0,0 +1,206 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Value conversion between temporal extension types. +//! +//! A `vortex.date` value is a count of `source_unit` since the epoch and a `vortex.timestamp` +//! value is a count of `target_unit` since the same epoch, so converting between them is a +//! rescale of the storage value. Both the array kernel and the scalar cast need it, and they +//! need to agree: a scan falsifies `cast(col as timestamp) > lit` into +//! `cast(max(col) as timestamp) <= lit`, binds `max(col)` to a literal and casts *that*, so a +//! file is pruned on the scalar conversion and its surviving rows are filtered on the array +//! one. Two conversions that disagree drop rows that match. +//! +//! [`DateToTimestamp`] is the one implementation of it, so neither caller can accept a value +//! the other refuses. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; + +/// Converts `vortex.date` values into `vortex.timestamp` values of another unit. +/// +/// Built once per conversion, so the per-value work is arithmetic and two comparisons. +pub(crate) struct DateToTimestamp { + /// The scale to apply. Exactly one of these is ever other than 1, so a value is multiplied + /// and then divided without losing the low bits of a widening conversion. + multiply: i64, + divide: i64, + /// The instants the target `vortex.timestamp` can hold, in the target's own unit. + min: i64, + max: i64, +} + +impl DateToTimestamp { + /// Prepare the conversion from `source_unit` dates to `target_unit` timestamps. + pub(crate) fn new(source_unit: TimeUnit, target_unit: TimeUnit) -> VortexResult { + let source_ns = to_nanoseconds(source_unit)?; + let target_ns = to_nanoseconds(target_unit)?; + + let (multiply, divide) = if source_ns >= target_ns { + (source_ns / target_ns, 1) + } else { + (1, target_ns / source_ns) + }; + + let (min, max) = Timestamp::storage_range(target_unit)?; + Ok(Self { + multiply, + divide, + min, + max, + }) + } + + /// Rescale one date value into the target timestamp's unit. + /// + /// Refuses rather than rounds or wraps: a value the target unit cannot hold exactly, one + /// that leaves the `i64` range, and one that lands outside the instants a + /// `vortex.timestamp` can represent are all errors, because a silently moved instant is a + /// wrong comparison. + pub(crate) fn convert(&self, value: i64) -> VortexResult { + let mut scaled = i128::from(value) + .checked_mul(i128::from(self.multiply)) + .ok_or_else(|| { + vortex_error::vortex_err!( + Compute: "Date value {value} overflows while scaling to timestamp" + ) + })?; + + if self.divide != 1 { + let divisor = i128::from(self.divide); + if scaled % divisor != 0 { + vortex_bail!( + Compute: "Date value {value} cannot be represented exactly in target timestamp unit" + ); + } + scaled /= divisor; + } + + let scaled = i64::try_from(scaled).map_err(|_| { + vortex_error::vortex_err!(Compute: "Date value {value} overflows target timestamp range") + })?; + + // The `i64` range is wider than the instants a timestamp scalar accepts, and this is + // where the array and the scalar would otherwise part company: an array can hold the + // value, and a scalar built from it is rejected by `Timestamp`'s validation, so a + // file's statistics fail to cast while the rows they gate convert happily. The bounds + // are `Timestamp`'s own, so there is one range rather than two that must agree. + if scaled < self.min || scaled > self.max { + vortex_bail!( + Compute: "Date value {value} is outside the range a timestamp can represent" + ); + } + + Ok(scaled) + } +} + +/// The number of nanoseconds in one `unit`. +fn to_nanoseconds(unit: TimeUnit) -> VortexResult { + match unit { + TimeUnit::Nanoseconds => Ok(1), + TimeUnit::Microseconds => Ok(1_000), + TimeUnit::Milliseconds => Ok(1_000_000), + TimeUnit::Seconds => Ok(1_000_000_000), + TimeUnit::Days => Ok(86_400_000_000_000), + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[test] + fn days_scale_up_to_nanoseconds() -> VortexResult<()> { + let conversion = DateToTimestamp::new(TimeUnit::Days, TimeUnit::Nanoseconds)?; + assert_eq!(conversion.convert(19_783)?, 1_709_251_200_000_000_000); + Ok(()) + } + + #[test] + fn milliseconds_scale_up_to_nanoseconds() -> VortexResult<()> { + let conversion = DateToTimestamp::new(TimeUnit::Milliseconds, TimeUnit::Nanoseconds)?; + assert_eq!( + conversion.convert(1_709_251_200_000)?, + 1_709_251_200_000_000_000 + ); + Ok(()) + } + + #[test] + fn a_value_the_target_unit_cannot_hold_exactly_is_refused() -> VortexResult<()> { + // One millisecond, asked for in seconds: the conversion divides, and 1 is not a whole + // number of seconds. Rounding it would move the instant. + let conversion = DateToTimestamp::new(TimeUnit::Milliseconds, TimeUnit::Seconds)?; + assert!(conversion.convert(1).is_err()); + // A whole second still converts. + assert_eq!(conversion.convert(2_000)?, 2); + Ok(()) + } + + #[test] + fn a_value_that_leaves_the_i64_range_is_refused() -> VortexResult<()> { + let conversion = DateToTimestamp::new(TimeUnit::Days, TimeUnit::Nanoseconds)?; + assert!(conversion.convert(i64::MAX).is_err()); + Ok(()) + } + + /// A date past the last instant a timestamp can hold has to be refused. + /// + /// The `i64` range does not catch these: `i32::MAX` days is a valid `Date32` value and + /// scales into `i64` seconds with room to spare, but no timestamp scalar can be built from + /// the result. Converting it anyway is what puts the array and the scalar out of step — + /// the array holds it, and the scalar a file's statistics go through does not. + #[rstest] + #[case(TimeUnit::Seconds)] + #[case(TimeUnit::Milliseconds)] + #[case(TimeUnit::Microseconds)] + fn a_date_beyond_the_last_representable_instant_is_refused( + #[case] target_unit: TimeUnit, + ) -> VortexResult<()> { + let conversion = DateToTimestamp::new(TimeUnit::Days, target_unit)?; + + assert!( + conversion.convert(i64::from(i32::MAX)).is_err(), + "the largest Date32 value is far past the last instant" + ); + // Jiff stops at the end of year 9999; 2_932_896 days is 10000-01-01. + assert!( + conversion.convert(2_932_896).is_err(), + "the first day past the end of the range is refused" + ); + assert!( + conversion.convert(2_932_895).is_ok(), + "the last day inside the range still converts" + ); + Ok(()) + } + + /// The range is the target's, not the source's. + #[test] + fn the_representable_range_narrows_with_the_target_unit() -> VortexResult<()> { + // 9999-12-31 fits in seconds through microseconds; nanoseconds run out in 2262. + let last_day = 2_932_895; + assert!( + DateToTimestamp::new(TimeUnit::Days, TimeUnit::Seconds)? + .convert(last_day) + .is_ok() + ); + assert!( + DateToTimestamp::new(TimeUnit::Days, TimeUnit::Nanoseconds)? + .convert(last_day) + .is_err() + ); + Ok(()) + } + + #[test] + fn a_timestamp_in_days_has_no_conversion() { + assert!(DateToTimestamp::new(TimeUnit::Days, TimeUnit::Days).is_err()); + } +} diff --git a/vortex-array/src/extension/datetime/matcher.rs b/vortex-array/src/extension/datetime/matcher.rs index e64cb289824..d538f54c089 100644 --- a/vortex-array/src/extension/datetime/matcher.rs +++ b/vortex-array/src/extension/datetime/matcher.rs @@ -13,6 +13,7 @@ use crate::extension::datetime::Date; use crate::extension::datetime::Time; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; +use crate::extension::datetime::resolve_timezone; /// Matcher for temporal extension data types. pub struct AnyTemporal; @@ -74,7 +75,7 @@ impl TemporalMetadata<'_> { Some(tz) => Ok(TemporalJiff::Zoned( jiff::Timestamp::UNIX_EPOCH .checked_add(unit.to_jiff_span(v)?)? - .in_tz(tz.as_ref())?, + .to_zoned(resolve_timezone(tz.as_ref())?), )), }, } @@ -103,3 +104,39 @@ impl Display for TemporalJiff { } } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_error::VortexResult; + + use crate::extension::datetime::TemporalMetadata; + use crate::extension::datetime::TimeUnit; + + /// Arrow allows a fixed UTC offset in place of an IANA name, and Iceberg maps every + /// `timestamptz` column to `+00:00`. + #[rstest] + // A zero offset resolves to UTC, which `jiff` annotates by name. + #[case("+00:00", "1970-01-01T00:00:01+00:00[UTC]")] + #[case("-05:30", "1969-12-31T18:30:01-05:30[-05:30]")] + #[case("+09:00", "1970-01-01T09:00:01+09:00[+09:00]")] + #[case("America/New_York", "1969-12-31T19:00:01-05:00[America/New_York]")] + fn to_jiff_resolves_timezone( + #[case] timezone: &str, + #[case] expected: &str, + ) -> VortexResult<()> { + let tz = Some(Arc::from(timezone)); + let metadata = TemporalMetadata::Timestamp(&TimeUnit::Seconds, &tz); + assert_eq!(metadata.to_jiff(1)?.to_string(), expected); + Ok(()) + } + + #[test] + fn to_jiff_rejects_unknown_timezone() { + let tz = Some(Arc::from("Not/A/Timezone")); + let metadata = TemporalMetadata::Timestamp(&TimeUnit::Seconds, &tz); + assert!(metadata.to_jiff(1).is_err()); + } +} diff --git a/vortex-array/src/extension/datetime/mod.rs b/vortex-array/src/extension/datetime/mod.rs index 86c6df0737f..3dccec3178a 100644 --- a/vortex-array/src/extension/datetime/mod.rs +++ b/vortex-array/src/extension/datetime/mod.rs @@ -3,14 +3,18 @@ //! Datetime extension DTypes, compatible with Apache Arrow. +mod cast; mod date; mod matcher; mod time; mod timestamp; +mod timezone; mod unit; +pub(crate) use cast::*; pub use date::*; pub use matcher::*; pub use time::*; pub use timestamp::*; +pub use timezone::*; pub use unit::*; diff --git a/vortex-array/src/extension/datetime/timestamp.rs b/vortex-array/src/extension/datetime/timestamp.rs index d07f68aa684..a9bffd5fbea 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 vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -22,6 +21,7 @@ use crate::dtype::extension::ExtDType; use crate::dtype::extension::ExtId; use crate::dtype::extension::ExtVTable; use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::resolve_timezone; use crate::scalar::ScalarValue; /// Timestamp DType. @@ -55,6 +55,38 @@ impl Timestamp { ExtDType::try_new(options, DType::Primitive(PType::I64, nullability)) .vortex_expect("failed to create timestamp dtype") } + + /// The inclusive range of storage values a `vortex.timestamp` in `unit` can hold. + /// + /// Taken from Jiff's own limits, and the one definition of them: `unpack_native` refuses a + /// scalar outside this range and `DateToTimestamp` refuses to convert a date into one, so + /// a converted date can always be carried by a scalar. A scan prunes a file on the scalar + /// conversion and filters the rows it kept on the array one, so the two ranges parting + /// company is a file dropped by a comparison the row filter would never have made. + /// + /// Nanoseconds are the whole of `i64`: its floor is 1677-09-21 and its ceiling 2262-04-11, + /// both far inside Jiff's range, so the conversion below clamps rather than narrowing. + pub(crate) fn storage_range(unit: TimeUnit) -> VortexResult<(i64, i64)> { + Ok(match unit { + TimeUnit::Nanoseconds => ( + i64::try_from(jiff::Timestamp::MIN.as_nanosecond()).unwrap_or(i64::MIN), + i64::try_from(jiff::Timestamp::MAX.as_nanosecond()).unwrap_or(i64::MAX), + ), + TimeUnit::Microseconds => ( + jiff::Timestamp::MIN.as_microsecond(), + jiff::Timestamp::MAX.as_microsecond(), + ), + TimeUnit::Milliseconds => ( + jiff::Timestamp::MIN.as_millisecond(), + jiff::Timestamp::MAX.as_millisecond(), + ), + TimeUnit::Seconds => ( + jiff::Timestamp::MIN.as_second(), + jiff::Timestamp::MAX.as_second(), + ), + TimeUnit::Days => vortex_bail!("Timestamp does not support Days time unit"), + }) + } } /// Options for the Timestamp DType. @@ -89,22 +121,64 @@ pub enum TimestampValue<'a> { Nanoseconds(i64, Option<&'a Arc>), } +impl TimestampValue<'_> { + /// The storage value and the unit it counts in. + fn storage(&self) -> (i64, TimeUnit) { + match self { + TimestampValue::Seconds(v, _) => (*v, TimeUnit::Seconds), + TimestampValue::Milliseconds(v, _) => (*v, TimeUnit::Milliseconds), + TimestampValue::Microseconds(v, _) => (*v, TimeUnit::Microseconds), + TimestampValue::Nanoseconds(v, _) => (*v, TimeUnit::Nanoseconds), + } + } + + /// The timezone the value renders in, if it carries one. + fn timezone(&self) -> Option<&Arc> { + match self { + TimestampValue::Seconds(_, tz) + | TimestampValue::Milliseconds(_, tz) + | TimestampValue::Microseconds(_, tz) + | TimestampValue::Nanoseconds(_, tz) => *tz, + } + } + + /// The instant this value denotes, or `None` for a count no timestamp can hold. + /// + /// Every step reports rather than aborts. The unchecked `Span` constructors panic outside + /// the span range and `Timestamp + Span` panics outside the timestamp range, so a + /// `Display` impl — which has no way to report a failure — cannot be built on them. + fn to_jiff(&self) -> Option { + // Every `i64` is a nanosecond instant and this constructor takes all of them, whereas + // a span's nanosecond floor is one above `i64::MIN`. `unpack_native` admits that whole + // range, so rendering has to as well. + if let TimestampValue::Nanoseconds(v, _) = self { + return jiff::Timestamp::from_nanosecond(i128::from(*v)).ok(); + } + + let (value, unit) = self.storage(); + let span = unit.to_jiff_span(value).ok()?; + jiff::Timestamp::UNIX_EPOCH.checked_add(span).ok() + } +} + 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), + // A count that denotes no instant must not abort a `Display` impl, which has no way to + // report the failure. `unpack_native` refuses to build a scalar from one, but the + // variant is public and constructible on its own, so render the raw count and its unit. + let Some(ts) = self.to_jiff() else { + let (value, unit) = self.storage(); + return write!(f, "{value}{unit}"); }; - let ts = jiff::Timestamp::UNIX_EPOCH + span; - match tz { + match self.timezone() { None => write!(f, "{ts}"), - Some(tz) => { - let adjusted_ts = ts.in_tz(tz.as_ref()).vortex_expect("unknown timezone"); - write!(f, "{adjusted_ts}",) - } + // A timezone that does not resolve must not abort it either. Render the underlying + // UTC timestamp instead. + Some(tz) => match resolve_timezone(tz.as_ref()) { + Ok(zone) => write!(f, "{}", ts.to_zoned(zone)), + Err(_) => write!(f, "{ts}"), + }, } } } @@ -197,33 +271,30 @@ 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), - ), + let value = 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"), }; - // Validate the storage value is within the valid range for Timestamp. - let ts = jiff::Timestamp::UNIX_EPOCH - .checked_add(span) - .map_err(|e| vortex_err!("Invalid timestamp scalar: {}", e))?; + // Compare against the timestamp's own range rather than routing the value through a + // Jiff `Span`. A span's limits are not a timestamp's at either end: they stop one + // short of `i64::MIN` nanoseconds, which is a valid 1677 instant, and they run past + // the last instant everywhere else, where the unchecked constructors abort rather + // than report. This range is also what `DateToTimestamp` converts into, so a scalar + // accepts exactly the values the array kernel produces. + let (min, max) = Self::storage_range(metadata.unit)?; + vortex_ensure!( + (min..=max).contains(&ts_value), + "Invalid timestamp scalar: {ts_value} {} is outside the instants a timestamp can represent ({min}..={max})", + metadata.unit + ); + // Validate the timezone resolves, accepting both IANA names and fixed UTC offsets. if let Some(tz) = tz { - ts.in_tz(tz.as_ref()) + resolve_timezone(tz.as_ref()) .map_err(|e| vortex_err!("Invalid timezone for timestamp scalar: {}", e))?; } @@ -241,6 +312,7 @@ mod tests { use crate::dtype::Nullability::Nullable; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; + use crate::extension::datetime::TimestampValue; use crate::scalar::PValue; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -319,4 +391,168 @@ mod tests { let data = [0x00u8, 0x0A, 0x00, b'U', b'T', b'C']; assert!(vtable.deserialize_metadata(&data).is_err()); } + + /// A fixed UTC offset renders at that offset. Iceberg emits `+00:00` for every `timestamptz`. + #[test] + fn display_renders_fixed_offset() { + let utc = Arc::from("+00:00"); + assert_eq!( + TimestampValue::Seconds(1, Some(&utc)).to_string(), + "1970-01-01T00:00:01+00:00[UTC]" + ); + + let tokyo = Arc::from("+09:00"); + assert_eq!( + TimestampValue::Seconds(1, Some(&tokyo)).to_string(), + "1970-01-01T09:00:01+09:00[+09:00]" + ); + } + + /// `Display` has no way to report a failure, so a timezone that does not resolve falls back to + /// rendering the underlying UTC timestamp rather than aborting the process. + #[test] + fn display_falls_back_on_unresolvable_timezone() { + let bad = Arc::from("Not/A/Timezone"); + let unzoned = TimestampValue::Seconds(1, None).to_string(); + assert_eq!(TimestampValue::Seconds(1, Some(&bad)).to_string(), unzoned); + } + + /// A storage value outside the timestamp range has to be reported, not panicked on. + /// + /// The hazard is Jiff's `Span`, whose limits are wider than a timestamp's and whose + /// unchecked constructors abort outside them: validating a storage value by building a + /// span from it takes the process down on the integer extrema below. `storage_range` is + /// total over `i64`, so every value here is reported. Each is one a `vortex.timestamp` + /// array can hold, so any read of one reaches this path. + /// + /// Nanoseconds are absent because every `i64` is a valid nanosecond timestamp — see + /// `every_i64_is_a_valid_nanosecond_timestamp`. + #[rstest::rstest] + #[case(TimeUnit::Seconds, i64::MAX)] + #[case(TimeUnit::Seconds, i64::MIN)] + #[case(TimeUnit::Milliseconds, i64::MAX)] + #[case(TimeUnit::Microseconds, i64::MAX)] + // Just past the last instant, and well short of the span limit — the half of the range + // that a span reports on, pinned alongside the half it aborts on. + #[case(TimeUnit::Seconds, 253_402_300_800)] + fn an_out_of_range_storage_value_is_an_error_not_a_panic( + #[case] unit: TimeUnit, + #[case] storage_value: i64, + ) { + let dtype = DType::Extension(Timestamp::new(unit, Nullable).erased()); + let err = Scalar::try_new(dtype, Some(storage_value.into())) + .expect_err("a value past the last instant is not a timestamp"); + assert!( + err.to_string().contains("Invalid timestamp scalar"), + "the error has to name the problem, got: {err}" + ); + } + + /// Every `i64` is a valid nanosecond timestamp, its floor included. + /// + /// `i64::MIN` nanoseconds is 1677-09-21T00:12:43.145224192Z and `i64::MAX` is + /// 2262-04-11T23:47:16.854775807Z, both far inside Jiff's range. A Jiff `Span` stops one + /// short of `i64::MIN`, so validating through one refuses a value a `vortex.timestamp[ns]` + /// array carries — the scalar and the array disagreeing about the same storage value, + /// which is the shape of defect this module exists to rule out. No `vortex.date` reaches + /// this value (neither of its units divides it), so the range is what has to be right + /// here, not the conversion. + #[rstest::rstest] + #[case(i64::MIN)] + #[case(i64::MIN + 1)] + #[case(0)] + #[case(i64::MAX)] + fn every_i64_is_a_valid_nanosecond_timestamp(#[case] storage_value: i64) { + let dtype = DType::Extension(Timestamp::new(TimeUnit::Nanoseconds, Nullable).erased()); + Scalar::try_new(dtype, Some(storage_value.into())) + .expect("every i64 nanosecond count is an instant a timestamp can represent"); + } + + /// The range a date converts into is exactly the range a scalar accepts. + /// + /// `DateToTimestamp` refuses to produce a value outside `Timestamp::storage_range` so that + /// a converted date can always be carried by a scalar. That only holds while this is the + /// same range the scalar itself enforces, so pin the two together at both ends — a scan + /// prunes a file on the scalar conversion and filters the rows it kept on the array one. + #[rstest::rstest] + #[case(TimeUnit::Seconds)] + #[case(TimeUnit::Milliseconds)] + #[case(TimeUnit::Microseconds)] + #[case(TimeUnit::Nanoseconds)] + fn the_storage_range_is_exactly_what_a_scalar_accepts(#[case] unit: TimeUnit) { + let (min, max) = Timestamp::storage_range(unit).expect("a range for every stored unit"); + let scalar = |v: i64| { + Scalar::try_new( + DType::Extension(Timestamp::new(unit, Nullable).erased()), + Some(v.into()), + ) + }; + + assert!( + scalar(min).is_ok(), + "{unit}: the floor of the range is an instant" + ); + assert!( + scalar(max).is_ok(), + "{unit}: the ceiling of the range is an instant" + ); + // Guarded, because the nanosecond range is the whole of `i64` and has no outside. + if min > i64::MIN { + assert!( + scalar(min - 1).is_err(), + "{unit}: one below the floor is not" + ); + } + if max < i64::MAX { + assert!( + scalar(max + 1).is_err(), + "{unit}: one above the ceiling is not" + ); + } + } + + /// A timestamp is never stored in days, so it has no range of storage values. + #[test] + fn a_timestamp_in_days_has_no_storage_range() { + assert!(Timestamp::storage_range(TimeUnit::Days).is_err()); + } + + /// Rendering has to reach every instant a scalar can carry, the nanosecond floor included. + /// + /// A `Display` impl has no way to report a failure, so building one on `Span`'s unchecked + /// constructors turns an unrenderable value into an abort. `i64::MIN` nanoseconds is the + /// case that matters: `unpack_native` admits it, so it reaches this path as a scalar. + #[rstest::rstest] + #[case( + TimestampValue::Nanoseconds(i64::MIN, None), + "1677-09-21T00:12:43.145224192Z" + )] + #[case( + TimestampValue::Nanoseconds(i64::MAX, None), + "2262-04-11T23:47:16.854775807Z" + )] + #[case(TimestampValue::Nanoseconds(0, None), "1970-01-01T00:00:00Z")] + #[case(TimestampValue::Seconds(1_709_251_200, None), "2024-03-01T00:00:00Z")] + fn display_renders_every_instant_a_scalar_carries( + #[case] value: TimestampValue<'_>, + #[case] expected: &str, + ) { + assert_eq!(value.to_string(), expected); + } + + /// A count that denotes no instant renders as itself rather than aborting. + /// + /// `unpack_native` refuses to build a scalar from one, but `TimestampValue` is public and + /// constructible on its own, so `Display` still has to survive it. + #[rstest::rstest] + #[case(TimestampValue::Seconds(i64::MAX, None), "9223372036854775807s")] + #[case(TimestampValue::Seconds(i64::MIN, None), "-9223372036854775808s")] + #[case(TimestampValue::Milliseconds(i64::MAX, None), "9223372036854775807ms")] + #[case(TimestampValue::Seconds(253_402_300_800, None), "253402300800s")] + fn display_falls_back_on_a_count_that_is_not_an_instant( + #[case] value: TimestampValue<'_>, + #[case] expected: &str, + ) { + assert_eq!(value.to_string(), expected); + } } diff --git a/vortex-array/src/extension/datetime/timezone.rs b/vortex-array/src/extension/datetime/timezone.rs new file mode 100644 index 00000000000..358b841ce49 --- /dev/null +++ b/vortex-array/src/extension/datetime/timezone.rs @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Resolution of Arrow timezone strings into `jiff` time zones. + +use jiff::tz::Offset; +use jiff::tz::TimeZone; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +/// Resolves an Arrow timezone string into a [`TimeZone`]. +/// +/// Arrow permits the timezone of a timestamp to be either an IANA time zone name +/// (`America/New_York`) or a fixed offset from UTC (`+00:00`). `jiff` only resolves the former, +/// so fixed offsets are parsed directly. Producers that emit the offset form are common: Iceberg +/// maps every `timestamptz` column to `+00:00`. +/// +/// # Examples +/// +/// ``` +/// use vortex_array::extension::datetime::resolve_timezone; +/// +/// assert_eq!(resolve_timezone("+00:00")?, jiff::tz::TimeZone::UTC); +/// assert_eq!(resolve_timezone("-05:30")?.to_fixed_offset()?.seconds(), -19_800); +/// assert!(resolve_timezone("Not/A/Timezone").is_err()); +/// # Ok::<(), Box>(()) +/// ``` +pub fn resolve_timezone(timezone: &str) -> VortexResult { + if let Ok(tz) = TimeZone::get(timezone) { + return Ok(tz); + } + + let Some(seconds) = parse_utc_offset_seconds(timezone) else { + vortex_bail!("unknown timezone: {timezone}"); + }; + let offset = Offset::from_seconds(seconds) + .map_err(|e| vortex_err!("timezone offset out of range: {timezone}: {e}"))?; + + Ok(TimeZone::fixed(offset)) +} + +/// Whether an Arrow timezone string denotes UTC. +/// +/// UTC has more than one Arrow spelling: any zero fixed offset such as `+00:00` (which is what +/// Iceberg and `arrow-rs` emit), and the IANA names that are zero-offset for all time, such as +/// `UTC` and `Etc/UTC`. +/// +/// # Examples +/// +/// ``` +/// use vortex_array::extension::datetime::is_utc_timezone; +/// +/// assert!(is_utc_timezone("UTC")); +/// assert!(is_utc_timezone("Etc/UTC")); +/// assert!(is_utc_timezone("+00:00")); +/// assert!(is_utc_timezone("-00:00")); +/// assert!(!is_utc_timezone("+09:00")); +/// assert!(!is_utc_timezone("America/New_York")); +/// ``` +pub fn is_utc_timezone(timezone: &str) -> bool { + resolve_timezone(timezone).is_ok_and(|tz| match tz.to_fixed_offset() { + Ok(offset) => offset == Offset::UTC, + // A named IANA zone has no single offset, so it denotes UTC only if it is zero-offset with + // no transitions in either direction. + Err(_) => { + let epoch = jiff::Timestamp::UNIX_EPOCH; + tz.to_offset(epoch) == Offset::UTC + && tz.following(epoch).next().is_none() + && tz.preceding(epoch).next().is_none() + } + }) +} + +/// Parses the fixed UTC offset forms Arrow allows — `±HH:MM`, `±HHMM` and `±HH` — into a signed +/// number of seconds. Returns `None` for anything else, including offsets whose minutes component +/// is not a valid minute. +fn parse_utc_offset_seconds(timezone: &str) -> Option { + let bytes = timezone.as_bytes(); + + let digits = match bytes.len() { + // ±HH:MM + 6 if bytes[3] == b':' => [bytes[1], bytes[2], bytes[4], bytes[5]], + // ±HHMM + 5 => [bytes[1], bytes[2], bytes[3], bytes[4]], + // ±HH + 3 => [bytes[1], bytes[2], b'0', b'0'], + _ => return None, + }; + + let mut values = [0i32; 4]; + for (value, digit) in values.iter_mut().zip(digits) { + *value = i32::from(digit.checked_sub(b'0').filter(|d| *d < 10)?); + } + + let minutes = values[2] * 10 + values[3]; + if minutes > 59 { + return None; + } + // Hours are not bounded here; `Offset::from_seconds` rejects anything beyond ±25:59:59. + let seconds = (values[0] * 10 + values[1]) * 60 * 60 + minutes * 60; + + match bytes[0] { + b'+' => Some(seconds), + b'-' => Some(-seconds), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use jiff::tz::TimeZone; + use rstest::rstest; + use vortex_error::VortexResult; + + use crate::extension::datetime::is_utc_timezone; + use crate::extension::datetime::resolve_timezone; + + #[rstest] + // Iceberg emits this form for every `timestamptz` column. + #[case("+00:00", 0)] + #[case("-00:00", 0)] + #[case("+09:00", 9 * 60 * 60)] + #[case("-05:30", -(5 * 60 * 60 + 30 * 60))] + #[case("+14:00", 14 * 60 * 60)] + // The abbreviated forms Arrow also permits. + #[case("+0930", 9 * 60 * 60 + 30 * 60)] + #[case("-08", -(8 * 60 * 60))] + fn resolves_fixed_offset( + #[case] timezone: &str, + #[case] expected_seconds: i32, + ) -> VortexResult<()> { + let tz = resolve_timezone(timezone)?; + assert_eq!(tz.to_fixed_offset()?.seconds(), expected_seconds); + Ok(()) + } + + #[test] + fn resolves_iana_name() -> VortexResult<()> { + // IANA lookup still takes precedence over offset parsing. + let tz = resolve_timezone("America/New_York")?; + assert_eq!(tz.iana_name(), Some("America/New_York")); + assert_eq!(resolve_timezone("UTC")?, TimeZone::UTC); + Ok(()) + } + + #[rstest] + #[case("Not/A/Timezone")] + #[case("")] + // Missing sign. + #[case("00:00")] + // Wrong separator, and a bare separator in the abbreviated position. + #[case("+00-00")] + #[case("+0:000")] + // Non-digits where digits are required. + #[case("+aa:bb")] + #[case("+ab")] + // Minutes out of range. + #[case("+00:60")] + #[case("+0099")] + // Hours beyond the ±25:59:59 offset range. + #[case("+26:00")] + #[case("-99")] + fn rejects_invalid_timezone(#[case] timezone: &str) { + assert!( + resolve_timezone(timezone).is_err(), + "expected {timezone} to be rejected" + ); + } + + #[rstest] + // Every spelling of a zero offset denotes UTC. + #[case("+00:00", true)] + #[case("-00:00", true)] + #[case("+0000", true)] + #[case("-0000", true)] + #[case("+00", true)] + #[case("-00", true)] + // So do the IANA names that are zero-offset for all time. + #[case("UTC", true)] + #[case("Etc/UTC", true)] + #[case("Etc/GMT", true)] + #[case("Zulu", true)] + #[case("Universal", true)] + // Non-zero offsets and zones that are not UTC do not. + #[case("+09:00", false)] + #[case("-05:30", false)] + #[case("+00:01", false)] + #[case("Etc/GMT+5", false)] + // A zone that is zero-offset only for part of its history is not UTC. + #[case("Europe/London", false)] + #[case("America/New_York", false)] + // Neither is anything unresolvable. + #[case("Not/A/Timezone", false)] + #[case("", false)] + fn detects_utc_timezone(#[case] timezone: &str, #[case] expected: bool) { + assert_eq!(is_utc_timezone(timezone), expected, "for {timezone}"); + } +} diff --git a/vortex-array/src/scalar/cast.rs b/vortex-array/src/scalar/cast.rs index 0ce8c331f1d..103af620cff 100644 --- a/vortex-array/src/scalar/cast.rs +++ b/vortex-array/src/scalar/cast.rs @@ -53,9 +53,15 @@ impl Scalar { return Scalar::try_new(target_dtype.clone(), self.value().cloned()); } - // TODO(connor): This isn't really correct for extension types. - // If the target is an extension type, then we want to cast to its storage type. - if let Some(ext_dtype) = target_dtype.as_extension_opt() { + // If the target is an extension type, cast to its storage type and re-label. This is + // only sound when the source carries no interpretation of its own: an extension source + // knows what its values mean and converts itself in the match below. Re-labelling one + // would hand back the source's number under the target's meaning — a `vortex.date[ms]` + // and a `vortex.timestamp[ns]` share `i64` storage, so the date would come back as an + // instant a million times too small. + if !matches!(self.dtype(), DType::Extension(..)) + && let Some(ext_dtype) = target_dtype.as_extension_opt() + { let cast_storage_scalar_value = self.cast(ext_dtype.storage_dtype())?.into_value(); return Scalar::try_new(target_dtype.clone(), cast_storage_scalar_value); } diff --git a/vortex-array/src/scalar/constructor.rs b/vortex-array/src/scalar/constructor.rs index 3ea25be816e..098e8aad76f 100644 --- a/vortex-array/src/scalar/constructor.rs +++ b/vortex-array/src/scalar/constructor.rs @@ -240,12 +240,24 @@ impl Scalar { /// /// # Panics /// - /// Panics if the storage dtype of `ext_dtype` does not match `value`'s dtype. + /// Panics if the storage dtype of `ext_dtype` does not match `storage_scalar`'s dtype, or if + /// `storage_scalar` is not a valid value for the extension type. Use + /// [`Scalar::try_extension_ref`] where an invalid storage value is recoverable. pub fn extension_ref(ext_dtype: ExtDTypeRef, storage_scalar: Scalar) -> Self { + Self::try_extension_ref(ext_dtype, storage_scalar) + .vortex_expect("unable to construct an extension `Scalar`") + } + + /// Creates a new extension scalar wrapping the given storage value, returning an error if + /// `storage_scalar` is not a valid value for the extension type. + /// + /// # Panics + /// + /// Panics if the storage dtype of `ext_dtype` does not match `storage_scalar`'s dtype. + pub fn try_extension_ref(ext_dtype: ExtDTypeRef, storage_scalar: Scalar) -> VortexResult { assert_eq!(ext_dtype.storage_dtype(), storage_scalar.dtype()); Self::try_new(DType::Extension(ext_dtype), storage_scalar.into_value()) - .vortex_expect("unable to construct an extension `Scalar`") } /// Creates a union scalar from a type ID and child scalar. diff --git a/vortex-array/src/scalar/typed_view/extension/mod.rs b/vortex-array/src/scalar/typed_view/extension/mod.rs index b9f524c7077..b9e73b75cc9 100644 --- a/vortex-array/src/scalar/typed_view/extension/mod.rs +++ b/vortex-array/src/scalar/typed_view/extension/mod.rs @@ -13,6 +13,9 @@ use vortex_error::vortex_panic; use crate::dtype::DType; use crate::dtype::extension::ExtDTypeRef; +use crate::extension::datetime::AnyTemporal; +use crate::extension::datetime::DateToTimestamp; +use crate::extension::datetime::TemporalMetadata; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -99,11 +102,17 @@ impl<'a> ExtScalar<'a> { return Scalar::try_new(target_dtype.clone(), self.value.cloned()); } - // We only allow casting to the same extension dtype for now. - if let DType::Extension(ext_dtype) = target_dtype - && self.ext_dtype.eq_ignore_nullability(ext_dtype) - { - return Scalar::try_new(target_dtype.clone(), self.value.cloned()); + if let DType::Extension(ext_dtype) = target_dtype { + // The same extension dtype: the value already means what the target says it does. + if self.ext_dtype.eq_ignore_nullability(ext_dtype) { + return Scalar::try_new(target_dtype.clone(), self.value.cloned()); + } + + // A different one needs the value converted, which is defined only for the pairs + // `Extension`'s `CastKernel` converts at the array level. + if let Some(scalar) = self.cast_date_to_timestamp(ext_dtype, target_dtype)? { + return Ok(scalar); + } } vortex_bail!( @@ -113,6 +122,46 @@ impl<'a> ExtScalar<'a> { target_dtype ); } + + /// Convert a `vortex.date` scalar to a `vortex.timestamp` one, rescaling the value. + /// + /// Returns `Ok(None)` for any other pair of extension types, leaving the caller to refuse + /// the cast. The conversion mirrors `Extension`'s `CastKernel`, which is what a scan + /// applies to the rows of a file whose statistics were compared through here. + fn cast_date_to_timestamp( + &self, + target_ext_dtype: &ExtDTypeRef, + target_dtype: &DType, + ) -> VortexResult> { + let (Some(source_temporal), Some(target_temporal)) = ( + self.ext_dtype.metadata_opt::(), + target_ext_dtype.metadata_opt::(), + ) else { + return Ok(None); + }; + + let (TemporalMetadata::Date(source_unit), TemporalMetadata::Timestamp(target_unit, _)) = + (source_temporal, target_temporal) + else { + return Ok(None); + }; + + // Null is handled by `Scalar::cast` before it reaches an extension scalar, and by the + // nullability check above; a null that gets this far still converts to a null. + let Some(value) = self.to_storage_scalar().as_primitive().as_opt::() else { + return Ok(None); + }; + let Some(value) = value else { + return Ok(Some(Scalar::try_new(target_dtype.clone(), None)?)); + }; + + let converted = DateToTimestamp::new(*source_unit, *target_unit)?.convert(value)?; + + let storage_value = Scalar::primitive(converted, target_dtype.nullability()) + .cast(target_ext_dtype.storage_dtype())? + .into_value(); + Ok(Some(Scalar::try_new(target_dtype.clone(), storage_value)?)) + } } // TODO(connor): In the future we may want to allow implementors to customize this behavior. diff --git a/vortex-array/src/scalar/typed_view/extension/tests.rs b/vortex-array/src/scalar/typed_view/extension/tests.rs index d9411806db6..9e1cf3f7484 100644 --- a/vortex-array/src/scalar/typed_view/extension/tests.rs +++ b/vortex-array/src/scalar/typed_view/extension/tests.rs @@ -9,12 +9,20 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::EmptyMetadata; +use crate::IntoArray; +use crate::arrays::ExtensionArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::extension::ExtDType; use crate::dtype::extension::ExtId; use crate::dtype::extension::ExtVTable; +use crate::executor::VortexSessionExecute; +use crate::extension::datetime::Date; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -346,3 +354,153 @@ fn test_ext_scalar_with_metadata() { let ext = scalar.as_extension(); assert_eq!(ext.ext_dtype().metadata::(), &1234); } + +/// A `vortex.date` scalar has to convert into a `vortex.timestamp` one, not be re-labelled. +/// +/// A scan falsifies `cast(col as timestamp) > lit` into `cast(max(col) as timestamp) <= lit` +/// and binds `max(col)` to a literal, so this cast decides whether a file is read at all. +/// Casting through the target's storage type instead — which is what happens when an +/// extension source is not given the chance to convert itself — hands back the date's own +/// number under the timestamp's meaning, and the two share `i64` storage for `Date64`, so it +/// does not even fail: the file is pruned and its matching rows never load. +#[test] +fn test_ext_scalar_cast_date_to_timestamp() { + // 2024-03-01, as days and as milliseconds since the epoch. + const DAYS: i32 = 19_783; + const MILLIS: i64 = 1_709_251_200_000; + const NANOS: i64 = 1_709_251_200_000_000_000; + + let nanos_dtype = + DType::Extension(Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased()); + + let from_days = Scalar::try_new( + DType::Extension(Date::new(TimeUnit::Days, Nullability::NonNullable).erased()), + Some(DAYS.into()), + ) + .unwrap(); + let casted = from_days.cast(&nanos_dtype).unwrap(); + assert_eq!(casted.dtype(), &nanos_dtype); + assert_eq!( + casted.as_extension().to_storage_scalar(), + Scalar::primitive(NANOS, Nullability::NonNullable), + "a date in days has to scale into the target's unit" + ); + + let from_millis = Scalar::try_new( + DType::Extension(Date::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased()), + Some(MILLIS.into()), + ) + .unwrap(); + let casted = from_millis.cast(&nanos_dtype).unwrap(); + assert_eq!(casted.dtype(), &nanos_dtype); + assert_eq!( + casted.as_extension().to_storage_scalar(), + Scalar::primitive(NANOS, Nullability::NonNullable), + "sharing i64 storage with the target is not a reason to skip the conversion" + ); +} + +/// The scalar cast has to land on the same instant as the array kernel. +/// +/// One of the two is applied to a file's statistics and the other to its rows, so a +/// disagreement between them is not a type error — it is a file pruned by a comparison the +/// row filter would never have made, or a scan that fails on a statistic whose rows convert +/// fine. +/// +/// Each value is cast on its own rather than as a batch: an array cast fails as a whole, so a +/// batch containing one unconvertible value would say nothing about the ones beside it. +#[test] +fn test_ext_scalar_cast_date_to_timestamp_matches_the_array_kernel() { + let session = crate::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Ordinary dates, then the boundaries: the largest `Date32` value, the first day no + // timestamp can represent (10000-01-01) and the last one that can. The boundaries are + // where a range check applied to one path and not the other shows up — the array holds + // the converted value and the scalar's validation rejects it. + for (source_unit, values) in [ + ( + TimeUnit::Days, + vec![0i64, 19_783, -1, i64::from(i32::MAX), 2_932_896, 2_932_895], + ), + ( + TimeUnit::Milliseconds, + vec![ + 0i64, + 1_709_251_200_000, + -86_400_000, + i64::from(i32::MAX) * 86_400_000, + 253_402_300_800_000, + 253_402_300_799_000, + ], + ), + ] { + for target_unit in [ + TimeUnit::Seconds, + TimeUnit::Milliseconds, + TimeUnit::Microseconds, + TimeUnit::Nanoseconds, + ] { + let source_dtype = Date::new(source_unit, Nullability::NonNullable).erased(); + let target_dtype = + DType::Extension(Timestamp::new(target_unit, Nullability::NonNullable).erased()); + + for value in &values { + let (storage, scalar_value) = if source_unit == TimeUnit::Days { + let Ok(days) = i32::try_from(*value) else { + continue; + }; + (PrimitiveArray::from_iter([days]).into_array(), days.into()) + } else { + ( + PrimitiveArray::from_iter([*value]).into_array(), + (*value).into(), + ) + }; + + let array_cast = ExtensionArray::new(source_dtype.clone(), storage) + .into_array() + .cast(target_dtype.clone()) + .and_then(|a| a.execute::(&mut ctx)); + + let scalar = + Scalar::try_new(DType::Extension(source_dtype.clone()), Some(scalar_value)) + .unwrap(); + let scalar_cast = scalar.cast(&target_dtype); + + match array_cast { + Ok(array) => { + let expected = array.into_array().execute_scalar(0, &mut ctx).unwrap(); + assert_eq!( + scalar_cast.unwrap(), + expected, + "{source_unit} -> {target_unit} disagrees on {value}" + ); + } + Err(_) => assert!( + scalar_cast.is_err(), + "{source_unit} -> {target_unit} converts {value}, which the array kernel refuses" + ), + } + } + } + } +} + +/// A pair of extension types with no defined conversion has to be refused, not re-labelled. +/// +/// `vortex.timestamp[ms]` and `vortex.timestamp[ns]` both store `i64`, so re-labelling one as +/// the other returns an instant a million times too small without any error. The array kernel +/// refuses this pair, and so must the scalar cast. +#[test] +fn test_ext_scalar_cast_between_timestamp_units_is_refused() { + let millis = Scalar::try_new( + DType::Extension(Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased()), + Some(1_709_251_200_000i64.into()), + ) + .unwrap(); + + let nanos = + DType::Extension(Timestamp::new(TimeUnit::Nanoseconds, Nullability::NonNullable).erased()); + assert!(millis.cast(&nanos).is_err()); +} From 587c05fe157dc57a5f11f2cb18ae9a2400d840b9 Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 16:56:52 -0700 Subject: [PATCH 04/12] docs: record the timezone and date/timestamp cast re-ports in the ledger Moves rows 4, 7, 15, 16 from "needs re-porting" to a new applied section, each with the commit and the test command that verified it, and notes that the earlier a98a78c1b array-cast port is superseded rather than left as a separate row. --- SPICE_PATCHES.md | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index fb67b25b8fc..45090b300ac 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -79,15 +79,23 @@ reading the new upstream code and re-implementing the fix, not resolving conflic | # | Patch | Old origin | What changed upstream | |---|---|---|---| | 6 | Avoid session lock re-entry in writer init | `c536c9aed` (#29) | `vortex-file/src/writer.rs` is +369/-44 lines different at `0.85.0`; nothing to cherry-pick onto. Still a **GAP** — no guard test exists in the Spice main repo either. | -| 7, 15, 16 | `vortex.date`→`vortex.timestamp` array + scalar casts, timestamp validation via `storage_range` | `7e5b08151` (#28), this branch's `3c5246867` (#93) | `CastReduce::cast` signature changed from `fn cast(array: &ExtensionArray, ...)` to `fn cast(array: ArrayView<'_, Extension>, ...)`. The exact function Spice added (`cast_temporal_date_to_timestamp`) doesn't exist at `0.85.0`. Real conflicts in `vortex-array/src/arrays/extension/compute/cast.rs` and `vortex-array/src/extension/datetime/{timestamp,mod}.rs`. High priority: row 15/16 guard spiceai/spiceai#13624 (wrongly pruned files / silently wrong data). | -| 4 | Fixed-offset timezone resolution | `6cdea73d6` (#75), `5b4bee108` (#78) | `vortex-array/src/extension/datetime/timezone.rs` is **deleted** in the merge-conflict sense at `0.85.0` (moved/restructured) — real conflicts also in `vortex-array/src/scalar/constructor.rs` and `vortex-duckdb/src/convert/{dtype,scalar}.rs`. The previous ledger's Verify (`test -f timezone.rs`) never actually proved this row; re-verify the real behavior once re-ported. | | 10 (real fix) | `UncompressedSizeInBytes` in `vortex-array/src/arrays/dict/take.rs` + `stats/mod.rs` | `a9ef29dea` (previously **not in the ledger at all** — the old ledger cited only the vendored, out-of-scope `6712e9ffa`) | `0.85.0` moved this into a new `vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs` implementing `AggregateFnVTable`. Plausible the redesign already subsumes this fix; re-verify against the new module rather than assuming either way. | | 11 | Intra-file decode parallelism | `26b274c72` (#62) | Real conflicts in `vortex-file/Cargo.toml`, `vortex-file/src/tests.rs`; `vortex-layout/src/scan/split_by.rs` also changed upstream (small diff, not yet compared line-by-line). Perf-only; still a ledger **GAP** (no guard). | | — | Cache `ArrayKernels` as a per-`ExecutionCtx` snapshot | `ccaa55627` | Real conflict in `vortex-array/src/executor.rs`. Perf-only. Not previously in the ledger. | | — | Restore `load_full` + `HashMap` import "dropped in rebase" | `ab4f0b177` | `vortex-array/src/arc_swap_map.rs` is deleted at `0.85.0` (modify/delete conflict) — check whether the underlying function this restored still has a reason to exist before re-porting; may be moot. Not previously in the ledger. | -| — | Migrate extension date-cast + Arrow Map test to ctx execute APIs | `d41e094cf` | Conflicts in `cast.rs` (same `ArrayView` signature issue as row 7) and `vortex-arrow/src/convert.rs`. Bundle with the row 7/15/16 re-port and the Map-alias re-port below — touches both areas. Not previously in the ledger. | | — | Restore lint checks on forks | `bb80c537b` | Conflicts across `.github/workflows/ci.yml`, `cast.rs`, `vortex-datafusion/src/persistent/sink.rs` (the last is out of scope). Low priority — CI config for this fork, not behavior Spice depends on. | +## Applied — re-ported as fresh implementations (not clean cherry-picks) + +These landed as new commits against `0.85.0` rather than cherry-picks, since the code they touch was restructured. Each is behaviorally verified, not just compiled. + +| # | Patch | New commit | Old origin | Verify | +|---|---|---|---|---| +| 4 | Fixed-offset timezone resolution (`resolve_timezone`, threaded through scalar validation, `Display`, and `TemporalMetadata::to_jiff`; `Scalar::try_extension_ref` through the 4 sites that built extension scalars infallibly on a read path) | `dc9892190` | `6cdea73d6` (#75), `5b4bee108` (#78) | `cargo test -p vortex-array --lib extension::datetime::timezone extension::datetime::matcher` — **ran, 42 pass** | +| 7, 15, 16 | `vortex.date`→`vortex.timestamp` array + scalar casts via a shared `DateToTimestamp`, `Timestamp::storage_range` replacing Jiff-`Span`-based validation | `dc9892190` (supersedes this branch's earlier `a98a78c1b`, which used a `legacy_session()` workaround instead of the proper `CastKernel` split found in this row's own later Spice history) | `7e5b08151` (#28), this branch's `3c5246867` (#93) | `cargo test -p vortex-array --lib arrays::extension::compute::cast scalar::typed_view::extension` — **ran, all pass**; full crate: `cargo test -p vortex-array -p vortex-datetime-parts --lib` — **3428 + 72 pass, 0 failed** | + +Not ported: the DuckDB `TIMESTAMP_TZ` dtype/scalar `is_utc_timezone` gate (`vortex-duckdb/src/convert/{dtype,scalar}.rs`) — `vortex-duckdb` is not a crate Spice's `Cargo.toml` consumes, directly or transitively. Fork hygiene, not something Spice's re-pin needs; left open. + ## Untracked in the previous ledger — needs its own audit Six commits landed on `spiceai-54` on 2026-09-11 (author Ben Chambers) that were never added From 8c14139400af7ad1518abf56c53a93a3b2f538aa Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 16:59:32 -0700 Subject: [PATCH 05/12] docs: confirm the writer-init session-lock re-entry hazard has no live call site vortex-file/src/writer.rs's array-context construction was rewritten to call the lock-free, Vec-returning session.enabled_component_ids(...) once instead of the SessionGuard-returning session.arrays() twice in one statement -- the exact pattern the original patch existed to prevent. Swept every remaining session.arrays() call site across the crates Spice consumes; none re-invoke it within a live guard's scope. No code change needed; moves row 6 from "needs re-porting" to "dropped" with that sweep as evidence. --- SPICE_PATCHES.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 45090b300ac..d2cc8ed142a 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -63,10 +63,10 @@ content in favor of this rewrite, not by a code-level resolution. The actual cod | 9 | Pushdown TRUE-bubbling / empty `IN` list | `8044a8470` (#8) | Both touched files are in `vortex-datafusion`, which Spice vendors directly into `crates/vortex` in the main repo (`vortex-datafusion = { path = "crates/vortex" }`) and does not take from this fork at all. Not fork state. | | 10 (partial) | `UncompressedSizeInBytes` in `vortex-datafusion/src/persistent/format.rs` | `6712e9ffa` (#3) | Same as row 9 — vendored, out of scope for this fork. | | 5 | Balanced `list_contains` OR tree | `d694abda6` (#37) | The function this commit added (`or_arrays_balanced`/`array_depth`) no longer exists anywhere in `spiceai-54`'s own later history — superseded internally by the untracked Sept-11 series (see below), which itself converges with `0.85.0`'s own `vortex_utils::iter::ReduceBalancedIterExt`. Dead code; nothing to carry. | -| 14 (old origin only) | Original Arrow Map alias | `1a6dc54f1` | Superseded by the Aug-26 restoration pair (see **Needs re-porting** below); its files (`vortex-array/src/arrow/*`) predate the `vortex-arrow` crate split and no longer exist. | | — | `ea3890528` "duckdb 1.5.5 (#8895)" | — | Not a Spice patch: identical to upstream's own `0.80.0` release commit (`037b1a8778db`), already superseded by landing on `0.85.0`. | | — | `68badd689` "TableStrategy hides panic... (#8672)" | — | Not a Spice patch: `(cherry picked from commit 4abe3d04f...)`, authored by a Vortex/spiraldb maintainer, cherry-picked onto `spiceai-54` ahead of a release. Already upstream. | | — | `b5498e8e1` "get CI green" (#81), `8d2efe7a7` style import sort | — | CI/lint housekeeping for subsystems Spice doesn't build (java/, python/, ffi/) or trivial reformatting of files already changed upstream. Not carried; not fork state Spice depends on. | +| 6 | Avoid session lock re-entry in writer init | `c536c9aed` (#29) | The hazard was calling the `SessionGuard`-returning `session.arrays()` twice within one statement (`ArrayContext::new(self.session.arrays()...).with_registry(self.session.arrays()...)`), so the two guards overlapped for the whole statement instead of the first being dropped before the second was taken. `vortex-file/src/writer.rs`'s array-context construction was rewritten into `new_array_context`, which calls the lock-free, `Vec`-returning `session.enabled_component_ids(...)` exactly once and reuses the result — the guard-returning call this row exists to guard against isn't made here at all anymore. Swept every `.arrays()` call site across `vortex-file`, `vortex-array`, `vortex-layout`, `vortex-io`, `vortex-scan`, `vortex-btrblocks`, `vortex-utils`, `vortex-session` (`grep -rn "session\.arrays()"`, 10 hits): every one either calls it once or binds the guard to a variable and reuses that variable, never re-invoking the accessor within a live guard's scope. No code change; the hazard's precondition doesn't exist in the crates Spice consumes. Still a **GAP** in the sense that no regression test pins this — if a future edit reintroduces a double call, nothing here would catch it. | | 14 | Arrow Map alias — `DType` alias + map-entry recursion in the session importer | `840e746a9`, `3e7fa40d3` (superseded the stale `1a6dc54f1` origin, then this) | **Superseded by a full upstream rewrite, not just a fix.** Upstream 0.85.0 replaced the `List>` alias both Spice's patch and the pre-split `vortex-array/src/arrow` code used with a **native `Map` array/dtype**: `vortex-array/src/dtype/map.rs` (`MapDType`) and `vortex-array/src/arrays/map/` (`Array`, backed by a `ListView>`). Both `vortex-arrow/src/convert.rs` (import) and `vortex-arrow/src/executor/mod.rs` (export) dispatch `DataType::Map` to this native path, and `vortex-arrow/src/session.rs` recurses into map key/value fields for extension-type metadata exactly as Spice's `3e7fa40d3` did. **Ran `cargo test -p vortex-arrow --lib map` on this branch: 17 tests pass**, including `executor::map::tests::map_roundtrip_preserves_nested_uuid_fields` (the exact extension-recursion scenario `3e7fa40d3` fixed) and `session::tests::schema_roundtrip_preserves_map_uuid_fields`. Upstream's own coverage here is broader than Spice's two patches were. Attempting to cherry-pick the old alias-based patch onto this would be a regression, not a port — the P0 defect (spiceai/spiceai#13524) is closed by a better mechanism than the one Spice shipped. **Remaining gap**: this proves the fix at the `vortex-arrow` crate level; the Spice main repo's own guard (`crates/vortex/src/persistent/mod.rs::map_column_roundtrips_through_a_vortex_file`, per `docs/dev/fork_patches.md`) still needs re-running once this branch is pinned, since that test exercises the full write path through the vendored `vortex-datafusion` sink, which this fork does not build. | ## Needs re-porting (real conflicts against `0.85.0` — deferred, not dropped) @@ -78,7 +78,6 @@ reading the new upstream code and re-implementing the fix, not resolving conflic | # | Patch | Old origin | What changed upstream | |---|---|---|---| -| 6 | Avoid session lock re-entry in writer init | `c536c9aed` (#29) | `vortex-file/src/writer.rs` is +369/-44 lines different at `0.85.0`; nothing to cherry-pick onto. Still a **GAP** — no guard test exists in the Spice main repo either. | | 10 (real fix) | `UncompressedSizeInBytes` in `vortex-array/src/arrays/dict/take.rs` + `stats/mod.rs` | `a9ef29dea` (previously **not in the ledger at all** — the old ledger cited only the vendored, out-of-scope `6712e9ffa`) | `0.85.0` moved this into a new `vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs` implementing `AggregateFnVTable`. Plausible the redesign already subsumes this fix; re-verify against the new module rather than assuming either way. | | 11 | Intra-file decode parallelism | `26b274c72` (#62) | Real conflicts in `vortex-file/Cargo.toml`, `vortex-file/src/tests.rs`; `vortex-layout/src/scan/split_by.rs` also changed upstream (small diff, not yet compared line-by-line). Perf-only; still a ledger **GAP** (no guard). | | — | Cache `ArrayKernels` as a per-`ExecutionCtx` snapshot | `ccaa55627` | Real conflict in `vortex-array/src/executor.rs`. Perf-only. Not previously in the ledger. | From 1672b7328807922fe267fed1de431cb02268bc76 Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:03:50 -0700 Subject: [PATCH 06/12] fix(array): propagate UncompressedSizeInBytes through dict take, include it in PRUNING_STATS Re-ports a9ef29dea from spiceai-54: this fork's own ledger only tracked the vendored, out-of-scope half of the original patch (vortex-datafusion), never noticing this core-crate half existed. Confirmed it doesn't exist at 0.85.0 independently of the aggregate-fn redesign that moved computation of the stat itself into a new AggregateFnVTable module -- that redesign changes how the value is computed per array type, not whether it survives PRUNING_STATS or a dict take, so this is still needed on top of it. Without PRUNING_STATS including it, the stat isn't guaranteed to be computed at write time. Without take.rs propagating it (alongside Min/Max/IsConstant, which already were), a Dict-encoded column's byte-size statistic silently disappears after any take/filter, and the optimizer mis-sizes joins built over what remains -- a planning regression, not a correctness one, but a silent one either way. Added two tests (none existed for propagate_take_stats before): one confirming the new stat carries forward inexactly, one confirming IsConstant's existing inexact-propagation path (added by the same original patch) still fires correctly when an index is null and the exact path above it doesn't apply. Verified: cargo test -p vortex-array --lib: 3430 passed (up from 3428), 0 failed. cargo check across vortex-array, vortex-arrow, vortex-layout, vortex-file, vortex-io, vortex-utils, vortex-datetime-parts: clean. --- vortex-array/src/arrays/dict/take.rs | 98 +++++++++++++++++++++++++--- vortex-array/src/stats/mod.rs | 1 + 2 files changed, 90 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/arrays/dict/take.rs b/vortex-array/src/arrays/dict/take.rs index d557d34b72f..f3201e13dd6 100644 --- a/vortex-array/src/arrays/dict/take.rs +++ b/vortex-array/src/arrays/dict/take.rs @@ -156,17 +156,97 @@ pub(crate) fn propagate_take_stats( st.set(Stat::IsConstant, Precision::exact(true)); } } - let inexact_min_max = [Stat::Min, Stat::Max] - .into_iter() - .filter_map(|stat| match source.statistics().get(stat).into_inexact() { - Precision::Exact(scalar) | Precision::Inexact(scalar) => { - scalar.into_value().map(|sv| (stat, Precision::Inexact(sv))) - } - Precision::Absent => None, - }) - .collect::>(); + let inexact_min_max = [ + Stat::Min, + Stat::Max, + Stat::UncompressedSizeInBytes, + Stat::IsConstant, + ] + .into_iter() + .filter_map(|stat| match source.statistics().get(stat).into_inexact() { + Precision::Exact(scalar) | Precision::Inexact(scalar) => { + scalar.into_value().map(|sv| (stat, Precision::Inexact(sv))) + } + Precision::Absent => None, + }) + .collect::>(); st.combine_sets( &(unsafe { StatsSet::new_unchecked(inexact_min_max) }).as_typed_ref(source.dtype()), ) }) } + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + + use vortex_error::VortexExpect; + + use super::*; + use crate::dtype::Nullability; + use crate::scalar::Scalar; + + /// `UncompressedSizeInBytes` has to survive a dict take, inexactly, the same way `Min`/`Max` + /// already do: the optimizer sizes joins from it, and a value's absence after a take is read + /// as "unknown" rather than "zero", so losing it silently degrades plans over dictionary- + /// encoded columns rather than failing loudly. + #[test] + fn uncompressed_size_in_bytes_propagates_through_take_as_inexact() -> VortexResult<()> { + let source = buffer![1i32, 2, 3, 4].into_array(); + source.statistics().set( + Stat::UncompressedSizeInBytes, + Precision::exact( + Scalar::primitive(16u64, Nullability::NonNullable) + .into_value() + .vortex_expect("non-null size"), + ), + ); + + let target = buffer![1i32, 3].into_array(); + let indices = buffer![0u64, 2].into_array(); + propagate_take_stats(&source, &target, &indices)?; + + assert_eq!( + target + .statistics() + .get_as::(Stat::UncompressedSizeInBytes), + Precision::Inexact(16), + "the source's size should carry forward, downgraded to inexact" + ); + Ok(()) + } + + /// `IsConstant` is handled twice: exactly when every index is valid (the block above the + /// one this row extends), and inexactly here otherwise. A null index breaks the "every + /// output value equals the source's constant" guarantee, so the exact path must not fire, + /// but the source's own `IsConstant` reading (whatever it is) is still a reasonable inexact + /// carry-forward rather than dropping the stat entirely. + #[test] + fn is_constant_propagates_inexactly_through_take_with_a_null_index() -> VortexResult<()> { + let source = ConstantArray::new(Scalar::primitive(7i32, Nullability::NonNullable), 4) + .into_array(); + source.statistics().set( + Stat::IsConstant, + Precision::exact( + Scalar::from(true) + .into_value() + .vortex_expect("non-null bool"), + ), + ); + + let target = ConstantArray::new(Scalar::primitive(7i32, Nullability::NonNullable), 2) + .into_array(); + // A nullable index array with one null: `indices_all_valid` is false, so the exact + // IsConstant path above must not run. + let indices = crate::arrays::PrimitiveArray::from_option_iter([Some(0u64), None]) + .into_array(); + propagate_take_stats(&source, &target, &indices)?; + + assert_eq!( + target.statistics().get_as::(Stat::IsConstant), + Precision::Inexact(true), + "IsConstant should still carry forward inexactly when it can't be proven exact" + ); + Ok(()) + } +} diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs index f7610fc71e4..244cb1e1de5 100644 --- a/vortex-array/src/stats/mod.rs +++ b/vortex-array/src/stats/mod.rs @@ -41,6 +41,7 @@ pub const PRUNING_STATS: &[Stat] = &[ Stat::Sum, Stat::NullCount, Stat::NaNCount, + Stat::UncompressedSizeInBytes, ]; pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec { From 6e37f0c3ac73745ab750dddcc12d3322a40ef15c Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:04:17 -0700 Subject: [PATCH 07/12] docs: record the UncompressedSizeInBytes re-port in the ledger Moves row 10's real fix from "needs re-porting" to the applied section with its commit and the new tests that verify it, and notes explicitly that the AggregateFnVTable redesign doesn't subsume it -- it changes computation, not propagation or inclusion in PRUNING_STATS. --- SPICE_PATCHES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index d2cc8ed142a..80141c45ba1 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -78,7 +78,6 @@ reading the new upstream code and re-implementing the fix, not resolving conflic | # | Patch | Old origin | What changed upstream | |---|---|---|---| -| 10 (real fix) | `UncompressedSizeInBytes` in `vortex-array/src/arrays/dict/take.rs` + `stats/mod.rs` | `a9ef29dea` (previously **not in the ledger at all** — the old ledger cited only the vendored, out-of-scope `6712e9ffa`) | `0.85.0` moved this into a new `vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs` implementing `AggregateFnVTable`. Plausible the redesign already subsumes this fix; re-verify against the new module rather than assuming either way. | | 11 | Intra-file decode parallelism | `26b274c72` (#62) | Real conflicts in `vortex-file/Cargo.toml`, `vortex-file/src/tests.rs`; `vortex-layout/src/scan/split_by.rs` also changed upstream (small diff, not yet compared line-by-line). Perf-only; still a ledger **GAP** (no guard). | | — | Cache `ArrayKernels` as a per-`ExecutionCtx` snapshot | `ccaa55627` | Real conflict in `vortex-array/src/executor.rs`. Perf-only. Not previously in the ledger. | | — | Restore `load_full` + `HashMap` import "dropped in rebase" | `ab4f0b177` | `vortex-array/src/arc_swap_map.rs` is deleted at `0.85.0` (modify/delete conflict) — check whether the underlying function this restored still has a reason to exist before re-porting; may be moot. Not previously in the ledger. | @@ -92,6 +91,7 @@ These landed as new commits against `0.85.0` rather than cherry-picks, since the |---|---|---|---|---| | 4 | Fixed-offset timezone resolution (`resolve_timezone`, threaded through scalar validation, `Display`, and `TemporalMetadata::to_jiff`; `Scalar::try_extension_ref` through the 4 sites that built extension scalars infallibly on a read path) | `dc9892190` | `6cdea73d6` (#75), `5b4bee108` (#78) | `cargo test -p vortex-array --lib extension::datetime::timezone extension::datetime::matcher` — **ran, 42 pass** | | 7, 15, 16 | `vortex.date`→`vortex.timestamp` array + scalar casts via a shared `DateToTimestamp`, `Timestamp::storage_range` replacing Jiff-`Span`-based validation | `dc9892190` (supersedes this branch's earlier `a98a78c1b`, which used a `legacy_session()` workaround instead of the proper `CastKernel` split found in this row's own later Spice history) | `7e5b08151` (#28), this branch's `3c5246867` (#93) | `cargo test -p vortex-array --lib arrays::extension::compute::cast scalar::typed_view::extension` — **ran, all pass**; full crate: `cargo test -p vortex-array -p vortex-datetime-parts --lib` — **3428 + 72 pass, 0 failed** | +| 10 (real fix) | `UncompressedSizeInBytes` propagation through `vortex-array/src/arrays/dict/take.rs`, and inclusion in `stats::PRUNING_STATS` | `1672b7328` | `a9ef29dea` (previously **not in the ledger at all** — the old ledger cited only the vendored, out-of-scope `6712e9ffa`) | New tests added (none existed before): `cargo test -p vortex-array --lib arrays::dict::take` — **ran, 2 pass** (`uncompressed_size_in_bytes_propagates_through_take_as_inexact`, `is_constant_propagates_inexactly_through_take_with_a_null_index`); full crate `cargo test -p vortex-array --lib` — **3430 pass, 0 failed**. Not subsumed by the `AggregateFnVTable` redesign — that changes how the stat is computed per array type, not whether `PRUNING_STATS` includes it or `take` propagates it. | Not ported: the DuckDB `TIMESTAMP_TZ` dtype/scalar `is_utc_timezone` gate (`vortex-duckdb/src/convert/{dtype,scalar}.rs`) — `vortex-duckdb` is not a crate Spice's `Cargo.toml` consumes, directly or transitively. Fork hygiene, not something Spice's re-pin needs; left open. From 184648a648c78fbab5180e843c7e4b8413395ca8 Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:07:15 -0700 Subject: [PATCH 08/12] docs: re-scope the Sept-11 IN-list series, confirm ab4f0b177 is moot Read the full commit sequence rather than titles alone: the one non-perf commit in the Sept-11 series (a panic fix in an extension list_contains kernel) is itself deleted three commits later the same day, when the extension-specific kernel it patched is replaced by generic-probe handling -- matching 0.85.0, which never had that kernel. Not independently reproduced against 0.85.0, so this is a reading, not a run, but the series is not a live-bug row regardless: reframed from "needs its own audit" to "a real but sizeable performance feature port against a since-rewritten 944-line file," deprioritized below the smaller items for that reason. Also confirmed ab4f0b177 is moot (the file it touches doesn't exist at 0.85.0, the function it restores isn't referenced anywhere) and flagged ccaa55627's actual scope: it touches the central ExecutionCtx hot path, not a leaf file, so it carries more blast radius than the other open items here. --- SPICE_PATCHES.md | 49 ++++++++++++++++++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 12 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 80141c45ba1..d1ab4fb93ed 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -66,6 +66,7 @@ content in favor of this rewrite, not by a code-level resolution. The actual cod | — | `ea3890528` "duckdb 1.5.5 (#8895)" | — | Not a Spice patch: identical to upstream's own `0.80.0` release commit (`037b1a8778db`), already superseded by landing on `0.85.0`. | | — | `68badd689` "TableStrategy hides panic... (#8672)" | — | Not a Spice patch: `(cherry picked from commit 4abe3d04f...)`, authored by a Vortex/spiraldb maintainer, cherry-picked onto `spiceai-54` ahead of a release. Already upstream. | | — | `b5498e8e1` "get CI green" (#81), `8d2efe7a7` style import sort | — | CI/lint housekeeping for subsystems Spice doesn't build (java/, python/, ffi/) or trivial reformatting of files already changed upstream. Not carried; not fork state Spice depends on. | +| — | `ab4f0b177` "restore load_full + HashMap import dropped in rebase" | — | `vortex-array/src/arc_swap_map.rs` doesn't exist at `0.85.0` at all (not moved, gone), and `grep -rn "load_full"` across `vortex-array/src` finds zero matches. Spice's own fix for a Spice-only prior rebase mistake in a file structure that no longer exists. Nothing to port. | | 6 | Avoid session lock re-entry in writer init | `c536c9aed` (#29) | The hazard was calling the `SessionGuard`-returning `session.arrays()` twice within one statement (`ArrayContext::new(self.session.arrays()...).with_registry(self.session.arrays()...)`), so the two guards overlapped for the whole statement instead of the first being dropped before the second was taken. `vortex-file/src/writer.rs`'s array-context construction was rewritten into `new_array_context`, which calls the lock-free, `Vec`-returning `session.enabled_component_ids(...)` exactly once and reuses the result — the guard-returning call this row exists to guard against isn't made here at all anymore. Swept every `.arrays()` call site across `vortex-file`, `vortex-array`, `vortex-layout`, `vortex-io`, `vortex-scan`, `vortex-btrblocks`, `vortex-utils`, `vortex-session` (`grep -rn "session\.arrays()"`, 10 hits): every one either calls it once or binds the guard to a variable and reuses that variable, never re-invoking the accessor within a live guard's scope. No code change; the hazard's precondition doesn't exist in the crates Spice consumes. Still a **GAP** in the sense that no regression test pins this — if a future edit reintroduces a double call, nothing here would catch it. | | 14 | Arrow Map alias — `DType` alias + map-entry recursion in the session importer | `840e746a9`, `3e7fa40d3` (superseded the stale `1a6dc54f1` origin, then this) | **Superseded by a full upstream rewrite, not just a fix.** Upstream 0.85.0 replaced the `List>` alias both Spice's patch and the pre-split `vortex-array/src/arrow` code used with a **native `Map` array/dtype**: `vortex-array/src/dtype/map.rs` (`MapDType`) and `vortex-array/src/arrays/map/` (`Array`, backed by a `ListView>`). Both `vortex-arrow/src/convert.rs` (import) and `vortex-arrow/src/executor/mod.rs` (export) dispatch `DataType::Map` to this native path, and `vortex-arrow/src/session.rs` recurses into map key/value fields for extension-type metadata exactly as Spice's `3e7fa40d3` did. **Ran `cargo test -p vortex-arrow --lib map` on this branch: 17 tests pass**, including `executor::map::tests::map_roundtrip_preserves_nested_uuid_fields` (the exact extension-recursion scenario `3e7fa40d3` fixed) and `session::tests::schema_roundtrip_preserves_map_uuid_fields`. Upstream's own coverage here is broader than Spice's two patches were. Attempting to cherry-pick the old alias-based patch onto this would be a regression, not a port — the P0 defect (spiceai/spiceai#13524) is closed by a better mechanism than the one Spice shipped. **Remaining gap**: this proves the fix at the `vortex-arrow` crate level; the Spice main repo's own guard (`crates/vortex/src/persistent/mod.rs::map_column_roundtrips_through_a_vortex_file`, per `docs/dev/fork_patches.md`) still needs re-running once this branch is pinned, since that test exercises the full write path through the vendored `vortex-datafusion` sink, which this fork does not build. | @@ -79,8 +80,7 @@ reading the new upstream code and re-implementing the fix, not resolving conflic | # | Patch | Old origin | What changed upstream | |---|---|---|---| | 11 | Intra-file decode parallelism | `26b274c72` (#62) | Real conflicts in `vortex-file/Cargo.toml`, `vortex-file/src/tests.rs`; `vortex-layout/src/scan/split_by.rs` also changed upstream (small diff, not yet compared line-by-line). Perf-only; still a ledger **GAP** (no guard). | -| — | Cache `ArrayKernels` as a per-`ExecutionCtx` snapshot | `ccaa55627` | Real conflict in `vortex-array/src/executor.rs`. Perf-only. Not previously in the ledger. | -| — | Restore `load_full` + `HashMap` import "dropped in rebase" | `ab4f0b177` | `vortex-array/src/arc_swap_map.rs` is deleted at `0.85.0` (modify/delete conflict) — check whether the underlying function this restored still has a reason to exist before re-porting; may be moot. Not previously in the ledger. | +| — | Cache `ArrayKernels` as a per-`ExecutionCtx` snapshot | `ccaa55627` | Real conflict in `vortex-array/src/executor.rs`. Adds a `KernelSnapshot` type and `ArrayKernels::snapshot()` (confirmed absent at `0.85.0`) into the hot `execute_until` path on the central `ExecutionCtx` struct — higher blast radius than the other items here if ported carelessly. Perf-only. Not previously in the ledger. | | — | Restore lint checks on forks | `bb80c537b` | Conflicts across `.github/workflows/ci.yml`, `cast.rs`, `vortex-datafusion/src/persistent/sink.rs` (the last is out of scope). Low priority — CI config for this fork, not behavior Spice depends on. | ## Applied — re-ported as fresh implementations (not clean cherry-picks) @@ -101,16 +101,41 @@ Six commits landed on `spiceai-54` on 2026-09-11 (author Ben Chambers) that were to this file, so this re-cut is also the first time they've been evaluated: `6c9ffc507`, `2f1a22ada`, `af1c5b301`, `f73241661`, `95d40c8bb`, `aff66352b` — a rewrite of -`IN`-list / `list_contains` handling (null-bearing lists, extension-type lists, constant-list -set-probing), overlapping row 5's file (`vortex-array/src/scalar_fn/fns/list_contains/mod.rs`) -and, per the audit, written against a `list_contains` dispatch shape `0.85.0` has since -restructured around a generic `process_matches::` with offset reinterpret-casting. -`0.85.0` independently converged on the same `vortex_utils::iter::ReduceBalancedIterExt` -utility this series uses, which is *why* row 5's original patch is dead code — but the -null-handling and extension-list logic itself is likely still needed and was not attempted -as a cherry-pick in this pass (expected to conflict, not confirmed dropped). This is the -single largest remaining item: budget a dedicated session, comparing upstream's new dispatch -shape file-by-file before re-applying. +`IN`-list / `list_contains` handling, overlapping row 5's file +(`vortex-array/src/scalar_fn/fns/list_contains/mod.rs`) and, per the audit, written against a +`list_contains` dispatch shape `0.85.0` has since restructured around a generic +`process_matches::` with offset reinterpret-casting. `0.85.0` independently converged on +the same `vortex_utils::iter::ReduceBalancedIterExt` utility this series uses, which is *why* +row 5's original patch is dead code. + +**Re-read in full and re-scoped**: this series is Spice's own *performance* work, not a fix +for a pre-existing upstream bug — every commit in it is `perf(...)` except one, and that one +(`af1c5b301`, "decline a null-bearing extension list instead of panicking") fixes a panic in +an extension-type `list_contains` kernel that **`2f1a22ada`, three commits later the same day, +deletes outright** ("removes the `Scalar::new_unchecked` that retagged each element... Moving +[extension handling] there drops the per-batch cost to nothing"). By the end of the series, +Spice's own code no longer has an extension-specific `list_contains` kernel at all — matching +`0.85.0`, which never had one either (`grep` finds no `extension` reference anywhere in +`vortex-array/src/scalar_fn/fns/list_contains/mod.rs`). The panic's precondition — a +Spice-only kernel calling `Scalar::new_unchecked` on a null extension element — exists on +neither `spiceai-54`'s current tip nor `0.85.0`. **Not independently reproduced against +0.85.0** (no repro attempted; this is a reading of the commit sequence, not a run) — call this +unverified rather than confirmed, but it is not a live-bug row regardless: at worst it is +moot, not open. + +What the series actually adds beyond that: a hash-set membership probe for constant `IN` +lists instead of one `Eq` array + OR per element, an interval-based statistics-pruning +rewrite (a single top-level `OR` of gaps, replacing an `AND`-of-equalities that the optimizer +compared pairwise — quadratic in list length), an FSST-level `ListContainsElementKernel` that +answers `IN` over compressed strings without decompressing, and a word-at-a-time null-mask +probe. All four are `perf(...)`-labeled, real, and substantial (the largest single commit +here, `f73241661`, is also the one with the most test evidence: 39 recorded answers in +`vortex-array/tests/in_list_differential.rs`, described as failing under 5 of 6 mutations to +the pruning logic in `vortex-layout/tests/list_contains_pruning.rs` — neither test file exists +at `0.85.0`). Re-implementing this is a **feature port against a since-rewritten 944-line +file**, not a patch cherry-pick — deprioritized below the smaller items in this ledger for +that reason, not because it lacks value. Budget a dedicated session that reads +`process_matches::` first. Also unresolved from the previous audit and not re-checked here: `4e2d62654` (#87, "Absolute split concurrency, stop deriving concurrency defaults from host parallelism") — thematically From b18002d7e7f112985bceb3ebb1c023cec59bdfeb Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:15:34 -0700 Subject: [PATCH 09/12] fix(vortex-layout): adapt flat reader's cherry-picked tests to BoundExpression dafa0960a (cherry-picked earlier on this branch as "Fix flat reader subrange decode reuse") added a #[cfg(test)] module using APIs from spiceai-54's 2026-06-25 state: an undefined SESSION static (should be new_session()), an ArrayContext passed without the .into() 0.85.0's write_stream signature needs, and six call sites passing a raw Expression where LayoutReader's filter_evaluation/projection_evaluation now require &BoundExpression (upstream's "propagate bound expressions through layouts", #9126). This went undetected because the earlier verification pass ran `cargo check` without --tests, so the test module's compile errors never surfaced -- surfaced only now while adding cargo check --tests coverage for a different patch. Fixed by binding each expression once per branch with the same `.bind(reader.dtype())` pattern the file's own flat_identity test already used. Verified: cargo test -p vortex-layout --lib layouts::flat::reader -- 9 passed, 0 failed. This is the first real test run of dafa0960a's actual fix (segment-request coalescing across regressed query patterns), not just a compile check. --- vortex-layout/src/layouts/flat/reader.rs | 47 ++++++++++++++++-------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index f655dcf27f4..87155b55e8e 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -312,7 +312,7 @@ mod test { let (ptr, eof) = SequenceId::root().split(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::new(segments), array.to_array_stream().sequenced(ptr), eof, @@ -381,7 +381,7 @@ mod test { #[case] pattern: RegressedQueryPattern, ) -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = new_session().with_handle(handle); let mut ctx = session.create_execution_ctx(); let requests = Arc::new(AtomicUsize::new(0)); @@ -393,10 +393,17 @@ mod test { Arc::clone(&requests), ) .await?; - let first = - reader.projection_evaluation(&(0..3), &root(), MaskFuture::new_true(3))?; - let second = - reader.projection_evaluation(&(3..6), &root(), MaskFuture::new_true(3))?; + let identity = root().bind(reader.dtype())?; + let first = reader.projection_evaluation( + &(0..3), + &identity, + MaskFuture::new_true(3), + )?; + let second = reader.projection_evaluation( + &(3..6), + &identity, + MaskFuture::new_true(3), + )?; let (first, second): (VortexResult, VortexResult) = futures::join!(first, second); @@ -410,7 +417,7 @@ mod test { Arc::clone(&requests), ) .await?; - let filter = not_eq(root(), lit(0_i32)); + let filter = not_eq(root(), lit(0_i32)).bind(reader.dtype())?; let first = reader.filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3))?; let second = @@ -427,13 +434,14 @@ mod test { Arc::clone(&requests), ) .await?; - let filter = gt(root(), lit(1_i32)); + let filter = gt(root(), lit(1_i32)).bind(reader.dtype())?; + let identity = root().bind(reader.dtype())?; let first_mask = reader.filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3))?; let second_mask = reader.filter_evaluation(&(3..6), &filter, MaskFuture::new_true(3))?; - let first = reader.projection_evaluation(&(0..3), &root(), first_mask)?; - let second = reader.projection_evaluation(&(3..6), &root(), second_mask)?; + let first = reader.projection_evaluation(&(0..3), &identity, first_mask)?; + let second = reader.projection_evaluation(&(3..6), &identity, second_mask)?; let (first, second): (VortexResult, VortexResult) = futures::join!(first, second); @@ -446,14 +454,15 @@ mod test { let expected = source.clone().apply(&projection)?; let reader = counted_flat_reader(&session, source, Arc::clone(&requests)).await?; + let bound_projection = projection.bind(reader.dtype())?; let first = reader.projection_evaluation( &(0..3), - &projection, + &bound_projection, MaskFuture::new_true(3), )?; let second = reader.projection_evaluation( &(3..6), - &projection, + &bound_projection, MaskFuture::new_true(3), )?; @@ -484,12 +493,14 @@ mod test { .filter(Mask::from_iter([false, true, true]))?; let reader = counted_flat_reader(&session, source, Arc::clone(&requests)).await?; + let filter = filter.bind(reader.dtype())?; + let identity = root().bind(reader.dtype())?; let first_mask = reader.filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3))?; let second_mask = reader.filter_evaluation(&(3..6), &filter, MaskFuture::new_true(3))?; - let first = reader.projection_evaluation(&(0..3), &root(), first_mask)?; - let second = reader.projection_evaluation(&(3..6), &root(), second_mask)?; + let first = reader.projection_evaluation(&(0..3), &identity, first_mask)?; + let second = reader.projection_evaluation(&(3..6), &identity, second_mask)?; let (first, second): (VortexResult, VortexResult) = futures::join!(first, second); @@ -521,12 +532,16 @@ mod test { .apply(&projection)?; let reader = counted_flat_reader(&session, source, Arc::clone(&requests)).await?; + let filter = filter.bind(reader.dtype())?; + let bound_projection = projection.bind(reader.dtype())?; let first_mask = reader.filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3))?; let second_mask = reader.filter_evaluation(&(3..6), &filter, MaskFuture::new_true(3))?; - let first = reader.projection_evaluation(&(0..3), &projection, first_mask)?; - let second = reader.projection_evaluation(&(3..6), &projection, second_mask)?; + let first = + reader.projection_evaluation(&(0..3), &bound_projection, first_mask)?; + let second = + reader.projection_evaluation(&(3..6), &bound_projection, second_mask)?; let (first, second): (VortexResult, VortexResult) = futures::join!(first, second); From 121d5a724836c7f8451827f5be714a550a037846 Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:22:17 -0700 Subject: [PATCH 10/12] feat(vortex-layout): absolute split concurrency, stop deriving defaults from host parallelism (#87) Re-ports 4e2d62654 from spiceai-54. Confirmed live at 0.85.0 before porting (not assumed): ScanBuilder::with_concurrency's setting was multiplied by get_available_parallelism() at two sites (scan_builder.rs, repeated_scan.rs), so with_concurrency(1) still ran one split per worker thread rather than serially -- a caller sizing a memory reservation against the configured number under-reserves by the host's core count. MultiLayoutDataSource's deferred-file-open concurrency had the same defect one level worse: it used get_available_parallelism() as the value directly, so file-open width scaled with the node under a cgroup CPU quota narrower than the machine, the same class of bug crates/cpu-budget exists to prevent elsewhere. Introduces SplitConcurrency (PerWorker/Absolute) with one effective() that resolves either into a width, so a caller reading the setting and the scan honoring it cannot disagree about what it meant. with_absolute_concurrency is new; with_concurrency keeps its existing per-worker meaning. MultiLayoutDataSource's deferred-open default is now the fixed DEFAULT_CONCURRENCY (8), not the host's core count. The bound only holds if split construction is lazy under it: constructing a split task registers its reads with the I/O system, so building every task up-front (the previous LazyScanState::Preparing / RepeatedScan::execute path) let with_absolute_concurrency(1) still register every split's reads before the first one was polled. RepeatedScan::execute now delegates to a new split_tasks() lazy iterator that RepeatedScan::execute_stream buffers by the configured concurrency, and ScanBuilder::into_stream only runs `prepare` on the CPU pool before streaming through the same lazy path rather than collecting every task first. execute()/build() are unchanged in meaning (unbounded by the setting) and now say so in their docs. Ported the two regression tests from the original patch (previously none existed for this): both count registered splits via a new SplitProbe/ ProbeLayoutReader test harness in scan/test.rs and assert the high-water mark equals the configured concurrency, covering absolute concurrency 1 and 3, ordered and unordered, at both the RepeatedScan and ScanBuilder layers. Verified: cargo test -p vortex-layout --lib: 255 passed (up from 249), 0 failed -- including all 6 new/changed test cases. cargo test -p vortex-file --lib: 139 passed, 0 failed (downstream consumer of the scan path). cargo check --tests across vortex-array, vortex-arrow, vortex-layout, vortex-file, vortex-io, vortex-utils, vortex-datetime-parts: clean, zero warnings. Not ported: the vortex-file import-reordering hunks in this original commit (driver.rs, request.rs, segments/source.rs) are pure rustfmt churn already reflected in 0.85.0's own ordering; nothing to apply. --- vortex-layout/src/scan/multi.rs | 15 +- vortex-layout/src/scan/repeated_scan.rs | 160 ++++++++++++++++++--- vortex-layout/src/scan/scan_builder.rs | 168 +++++++++++++++------- vortex-layout/src/scan/test.rs | 180 ++++++++++++++++++++++++ 4 files changed, 449 insertions(+), 74 deletions(-) diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index f0b7a506e08..fa1bb375176 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -58,12 +58,18 @@ use vortex_scan::PartitionStream; use vortex_scan::ScanRequest; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::scan_builder::ScanBuilder; /// Default concurrency for opening deferred readers. +/// +/// A fixed width rather than the host's parallelism: a caller running under a CPU +/// entitlement narrower than the machine (a cgroup quota) cannot see that count +/// reflected in `available_parallelism`, so deriving the default from it opens +/// files far wider than the share the process was granted. Callers that want the +/// machine's parallelism can still ask for it through +/// [`MultiLayoutDataSource::with_concurrency`]. const DEFAULT_CONCURRENCY: usize = 8; /// An async factory that produces a [`LayoutReaderRef`]. @@ -129,7 +135,7 @@ impl MultiLayoutDataSource { session: &VortexSession, ) -> Self { let dtype = first.dtype().clone(); - let concurrency = get_available_parallelism().unwrap_or(DEFAULT_CONCURRENCY); + let concurrency = DEFAULT_CONCURRENCY; let total = 1 + remaining.len(); let mut sizes = byte_sizes; @@ -177,7 +183,7 @@ impl MultiLayoutDataSource { byte_sizes: Vec>, session: &VortexSession, ) -> Self { - let concurrency = get_available_parallelism().unwrap_or(DEFAULT_CONCURRENCY); + let concurrency = DEFAULT_CONCURRENCY; let mut sizes = byte_sizes; if sizes.is_empty() { @@ -208,7 +214,8 @@ impl MultiLayoutDataSource { /// Sets the concurrency for opening deferred readers. /// /// Controls how many file opens run in parallel via `buffer_unordered`. - /// Defaults to the number of available CPU cores. + /// Defaults to 8 — a fixed width, not the host's parallelism, so a process + /// under a CPU quota does not open files against the whole machine's core count. pub fn with_concurrency(mut self, concurrency: usize) -> Self { self.concurrency = concurrency; self diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 413761b8103..06506c1b711 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use futures::Stream; use futures::future::BoxFuture; +use futures::future::Either as EitherFuture; use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; @@ -23,12 +24,13 @@ use vortex_io::runtime::BlockingRuntime; use vortex_io::session::RuntimeSessionExt; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; +use crate::scan::scan_builder::SplitConcurrency; use crate::scan::splits::Splits; use crate::scan::tasks::TaskContext; +use crate::scan::tasks::TaskFuture; use crate::scan::tasks::split_exec; /// A projected subset (by indices, range, and filter) of rows from a Vortex data source. @@ -47,8 +49,8 @@ pub struct RepeatedScan { selection: Selection, /// The natural splits of the file. splits: Splits, - /// The number of splits to make progress on concurrently **per-thread**. - concurrency: usize, + /// How many splits to make progress on concurrently. + concurrency: SplitConcurrency, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. map_fn: Arc VortexResult + Send + Sync>, /// Maximal number of rows to read (after filtering) @@ -98,7 +100,7 @@ impl RepeatedScan { row_range: Option>, selection: Selection, splits: Splits, - concurrency: usize, + concurrency: SplitConcurrency, map_fn: Arc VortexResult + Send + Sync>, limit: Option, dtype: DType, @@ -119,10 +121,23 @@ impl RepeatedScan { } } + /// Constructs a task per row split of the scan, returned as a vector of futures. + /// + /// Note that this registers the reads of *every* split up-front. Prefer + /// [`Self::execute_stream`], which builds the same tasks lazily so that the configured + /// concurrency bounds how many splits have reads in flight. pub fn execute( &self, row_range: Option>, ) -> VortexResult>>>> { + self.split_tasks(row_range).collect() + } + + /// The row range of each split this scan will read, in file order. + /// + /// Computing every range up-front is cheap — it only intersects ranges — unlike + /// constructing the split tasks, which registers I/O. See [`Self::split_tasks`]. + fn split_ranges(&self, row_range: Option>) -> Vec> { let selection_range: Option> = match &self.selection { Selection::IncludeByIndex(buf) if !buf.is_empty() => { Some(buf[0]..buf[buf.len() - 1] + 1) @@ -142,7 +157,7 @@ impl RepeatedScan { None => Either::Left(vec.iter().copied()), Some(range) => { if range.is_empty() { - return Ok(Vec::new()); + return Vec::new(); } let lo = vec.partition_point(|&x| x <= range.start); let hi = vec.partition_point(|&x| x < range.end); @@ -160,7 +175,7 @@ impl RepeatedScan { None => Either::Left(ranges.iter().cloned()), Some(range) => { if range.is_empty() { - return Ok(Vec::new()); + return Vec::new(); } Either::Right(ranges.iter().filter_map(move |r| { let start = cmp::max(r.start, range.start); @@ -171,8 +186,24 @@ impl RepeatedScan { }), }; - let mut limit = self.limit; - let mut tasks = Vec::new(); + ranges.collect() + } + + /// Builds one split task at a time, in file order. + /// + /// Constructing a task registers its reads with the I/O system — that eager registration is + /// deliberate prefetch, see [`split_exec`] — so this iterator must stay lazy. Whatever + /// buffers it is what bounds the splits with reads in flight: [`Self::execute_stream`] + /// buffers by the scan's [`SplitConcurrency`], while [`Self::execute`] collects it and so + /// registers every split at once. + /// + /// A split whose task fails to construct yields the error and ends the iterator. + fn split_tasks( + &self, + row_range: Option>, + ) -> impl Iterator>>> + Send + 'static + use { + let mut ranges = self.split_ranges(row_range).into_iter(); + let selection = self.selection.clone(); let ctx = Arc::new(TaskContext { filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), reader: Arc::clone(&self.layout_reader), @@ -180,32 +211,47 @@ impl RepeatedScan { mapper: Arc::clone(&self.map_fn), }); - for range in ranges { - let row_mask = self.selection.row_mask(&range); - if row_mask.mask().all_false() { - continue; - } + let mut limit = self.limit; + let mut finished = false; - tasks.push(split_exec(Arc::clone(&ctx), row_mask, limit.as_mut())?); - if limit.is_some_and(|l| l == 0) { - break; - } - } + iter::from_fn(move || { + loop { + if finished { + return None; + } - Ok(tasks) + let range = ranges.next()?; + let row_mask = selection.row_mask(&range); + if row_mask.mask().all_false() { + continue; + } + + let task = split_exec(Arc::clone(&ctx), row_mask, limit.as_mut()); + if task.is_err() || limit.is_some_and(|l| l == 0) { + finished = true; + } + return Some(task); + } + }) } + /// Streams the scan, holding at most [`SplitConcurrency::effective`] splits in flight. + /// + /// Splits are constructed lazily under that bound, so a split's reads are not registered + /// until it is admitted. pub fn execute_stream( &self, row_range: Option>, ) -> VortexResult> + Send + 'static + use> { use futures::StreamExt; - let num_workers = get_available_parallelism().unwrap_or(1); - let concurrency = self.concurrency * num_workers; + let concurrency = self.concurrency.effective(); let handle = self.session.handle(); let stream = - futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task)); + futures::stream::iter(self.split_tasks(row_range)).map(move |task| match task { + Ok(task) => EitherFuture::Left(handle.spawn(task)), + Err(err) => EitherFuture::Right(futures::future::ready(Err(err))), + }); let stream = if self.ordered { stream.buffered(concurrency).boxed() @@ -225,3 +271,73 @@ fn intersect_ranges(left: Option<&Range>, right: Option>) -> Opt (Some(l), Some(r)) => Some(cmp::max(l.start, r.start)..cmp::min(l.end, r.end)), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + + use crate::scan::scan_builder::ScanBuilder; + use crate::scan::test::ProbeLayoutReader; + use crate::scan::test::SplitProbe; + use crate::scan::test::assert_bounded_scan; + use crate::scan::test::session_with_handle; + + /// The configured concurrency must bound the splits whose reads have been registered, so + /// split tasks have to be constructed lazily under the buffer rather than all up-front. + /// `1` is genuinely serial. + #[rstest] + #[case::serial(1, true)] + #[case::serial_unordered(1, false)] + #[case::bounded(3, true)] + #[case::bounded_unordered(3, false)] + fn absolute_concurrency_bounds_registered_splits( + #[case] concurrency: usize, + #[case] ordered: bool, + ) -> VortexResult<()> { + const ROW_COUNT: usize = 8; + + let mut ctx = array_session().create_execution_ctx(); + let probe = Arc::new(SplitProbe::default()); + let reader = Arc::new(ProbeLayoutReader::new(ROW_COUNT, Arc::clone(&probe))); + + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + + let scan = ScanBuilder::new(session, reader) + .with_absolute_concurrency(concurrency) + .with_ordered(ordered) + .prepare()?; + + let mut values = Vec::new(); + for (yielded, chunk) in runtime + .block_on_stream(scan.execute_stream(None)?) + .enumerate() + { + let prim = chunk?.execute::(&mut ctx)?; + values.extend(prim.into_buffer::().iter().copied()); + + let admitted = yielded + 1 + concurrency; + assert!( + probe.registered() <= admitted, + "{} splits registered after {} yielded, expected at most {admitted}", + probe.registered(), + yielded + 1, + ); + } + + if !ordered { + values.sort_unstable(); + } + assert_bounded_scan(&values, ROW_COUNT, &probe, concurrency); + + Ok(()) + } +} diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index f7abbb21fb2..db2b416cb22 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -27,7 +27,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_io::runtime::BlockingRuntime; -use vortex_io::runtime::Handle; use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_metrics::MetricsRegistry; @@ -45,6 +44,35 @@ use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; +/// How a scan's configured split concurrency maps onto splits actually in flight. +/// +/// The distinction is load-bearing for any caller that reasons about the number: +/// the per-worker form is a lookahead multiplier, so the width it produces depends +/// on the machine, while the absolute form is the width itself. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SplitConcurrency { + /// `n` splits per worker thread. The effective width is `n * available + /// parallelism`, which makes it a good default for throughput and a poor one + /// for a bound — see [`ScanBuilder::with_absolute_concurrency`]. + PerWorker(usize), + /// Exactly `n` splits in flight, whatever parallelism the host reports. + Absolute(usize), +} + +impl SplitConcurrency { + /// The number of splits to hold in flight. + /// + /// The single place either form becomes a width, so a caller that reads the + /// setting and the scan that honors it cannot disagree about what it meant. + pub fn effective(self) -> usize { + match self { + Self::PerWorker(n) => n.saturating_mul(get_available_parallelism().unwrap_or(1)), + Self::Absolute(n) => n, + } + .max(1) + } +} + /// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs. /// /// A scan has three independent row restriction mechanisms: @@ -72,8 +100,8 @@ pub struct ScanBuilder { /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare) /// uses them instead of walking the layout. natural_splits: Option>, - /// The number of splits to make progress on concurrently **per-thread**. - concurrency: usize, + /// How many splits to make progress on concurrently. + concurrency: SplitConcurrency, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, @@ -102,7 +130,7 @@ impl ScanBuilder { natural_splits: None, // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. - concurrency: 4, + concurrency: SplitConcurrency::PerWorker(4), map_fn: Arc::new(Ok), metrics_registry: None, file_stats: None, @@ -225,16 +253,46 @@ impl ScanBuilder { ) } - /// Returns the per-worker row-split concurrency. + /// Returns the effective concurrency. pub fn concurrency(&self) -> usize { + self.concurrency.effective() + } + + /// Returns the configured row-split concurrency. + pub fn split_concurrency(&self) -> SplitConcurrency { self.concurrency } /// The number of row splits to make progress on concurrently per-thread, must /// be greater than 0. + /// + /// The effective width is this value times the available parallelism. Callers + /// that must bound the splits actually in flight — because they size a memory + /// budget against it, or run under a CPU entitlement narrower than the host — + /// want [`Self::with_absolute_concurrency`] instead. pub fn with_concurrency(mut self, concurrency: usize) -> Self { assert!(concurrency > 0); - self.concurrency = concurrency; + self.concurrency = SplitConcurrency::PerWorker(concurrency); + self + } + + /// The exact number of row splits to hold in flight, independent of how many + /// worker threads the host offers. Must be greater than 0. + /// + /// `1` is genuinely serial. Use this when the number has to mean something to + /// the caller — a memory reservation sized per in-flight split, or a scan + /// confined to a CPU entitlement that [`get_available_parallelism`] cannot see + /// (it reports the machine's cores, so under a cgroup quota the per-worker form + /// multiplies by the node's core count rather than the share this process was + /// granted). + /// + /// The bound covers a split's reads, not just its execution: the streaming paths + /// ([`Self::into_stream`] and [`RepeatedScan::execute_stream`]) construct split tasks + /// lazily, so a split registers its reads only once admitted. [`Self::build`] returns + /// every task at once and so is not bounded by this setting. + pub fn with_absolute_concurrency(mut self, concurrency: usize) -> Self { + assert!(concurrency > 0); + self.concurrency = SplitConcurrency::Absolute(concurrency); self } @@ -364,6 +422,9 @@ impl ScanBuilder { } /// Constructs a task per row split of the scan, returned as a vector of futures. + /// + /// Every split registers its reads up-front, so the configured concurrency does not bound + /// this. Use [`Self::into_stream`] when the caller needs that bound. pub fn build(self) -> VortexResult>>>> { // The ultimate short circuit if self.limit.is_some_and(|l| l == 0) { @@ -392,20 +453,11 @@ impl ScanBuilder { enum LazyScanState { Builder(Option>>), - Preparing(PreparingScan), + Preparing(Task>>), Stream(BoxStream<'static, VortexResult>), Error(Option), } -type PreparedScanTasks = Vec>>>; - -struct PreparingScan { - ordered: bool, - concurrency: usize, - handle: Handle, - task: Task>>, -} - struct LazyScanStream { state: LazyScanState, } @@ -428,40 +480,19 @@ impl Stream for LazyScanStream { match &mut self.state { LazyScanState::Builder(builder) => { let builder = builder.take().vortex_expect("polled after completion"); - let ordered = builder.ordered; - let num_workers = get_available_parallelism().unwrap_or(1); - let concurrency = builder.concurrency * num_workers; - let handle = builder.session.handle(); - let task = handle - .spawn_cpu(move || builder.prepare().and_then(|scan| scan.execute(None))); - self.state = LazyScanState::Preparing(PreparingScan { - ordered, - concurrency, - handle, - task, - }); + // Only `prepare` runs on the CPU pool. Split tasks are built lazily by the + // stream below, so that the configured concurrency — not the split count — + // bounds how many splits have registered reads. + let task = builder.session.handle().spawn_cpu(move || builder.prepare()); + self.state = LazyScanState::Preparing(task); } - LazyScanState::Preparing(preparing) => { - match ready!(Pin::new(&mut preparing.task).poll(cx)) { - Ok(tasks) => { - let ordered = preparing.ordered; - let concurrency = preparing.concurrency; - let handle = preparing.handle.clone(); - let stream = - futures::stream::iter(tasks).map(move |task| handle.spawn(task)); - let stream = if ordered { - stream.buffered(concurrency).boxed() - } else { - stream.buffer_unordered(concurrency).boxed() - }; - let stream = stream - .filter_map(|chunk| async move { chunk.transpose() }) - .boxed(); - self.state = LazyScanState::Stream(stream); - } + LazyScanState::Preparing(task) => match ready!(Pin::new(task).poll(cx)) { + Ok(scan) => match scan.execute_stream(None) { + Ok(stream) => self.state = LazyScanState::Stream(stream.boxed()), Err(err) => self.state = LazyScanState::Error(Some(err)), - } - } + }, + Err(err) => self.state = LazyScanState::Error(Some(err)), + }, LazyScanState::Stream(stream) => return stream.as_mut().poll_next(cx), LazyScanState::Error(err) => return Poll::Ready(err.take().map(Err)), } @@ -985,6 +1016,47 @@ mod test { drop(runtime); } + /// [`ScanBuilder::into_stream`] must admit splits under the configured concurrency too: it + /// prepares the scan off-thread, then builds each split task lazily as the buffer admits it. + #[rstest::rstest] + #[case::serial(1)] + #[case::bounded(3)] + fn into_stream_bounds_registered_splits(#[case] concurrency: usize) -> VortexResult<()> { + const ROW_COUNT: usize = 8; + + let mut ctx = array_session().create_execution_ctx(); + let probe = Arc::new(crate::scan::test::SplitProbe::default()); + let reader = Arc::new(crate::scan::test::ProbeLayoutReader::new( + ROW_COUNT, + Arc::clone(&probe), + )); + + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + + let stream = ScanBuilder::new(session, reader) + .with_absolute_concurrency(concurrency) + .into_stream()?; + + let mut values = Vec::new(); + for (yielded, chunk) in runtime.block_on_stream(stream).enumerate() { + let prim = chunk?.execute::(&mut ctx)?; + values.extend(prim.into_buffer::().iter().copied()); + + let admitted = yielded + 1 + concurrency; + assert!( + probe.registered() <= admitted, + "{} splits registered after {} yielded, expected at most {admitted}", + probe.registered(), + yielded + 1, + ); + } + + crate::scan::test::assert_bounded_scan(&values, ROW_COUNT, &probe, concurrency); + + Ok(()) + } + #[test] fn into_stream_with_row_range() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-layout/src/scan/test.rs b/vortex-layout/src/scan/test.rs index 809639d169d..9b7039a14fe 100644 --- a/vortex-layout/src/scan/test.rs +++ b/vortex-layout/src/scan/test.rs @@ -1,13 +1,34 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::cmp; +use std::ops::Range; +use std::sync::Arc; use std::sync::LazyLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use vortex_array::IntoArray; +use vortex_array::MaskFuture; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::BoundExpression; +use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_io::runtime::Handle; use vortex_io::session::RuntimeSession; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_session::VortexSession; +use crate::ArrayFuture; +use crate::LayoutReader; +use crate::RowSplits; +use crate::SplitRange; use crate::session::LayoutSession; pub fn new_session() -> VortexSession { @@ -21,3 +42,162 @@ pub fn session_with_handle(handle: Handle) -> VortexSession { } pub static SCAN_SESSION: LazyLock = LazyLock::new(new_session); + +/// Counts how many splits a scan has registered reads for, and how many of those are still +/// outstanding. +/// +/// A split registers its reads when its task is *constructed* (see `scan::tasks::split_exec`), +/// which is what lets a scan's concurrency mean something: only admitted splits should have +/// registered. +#[derive(Debug, Default)] +pub struct SplitProbe { + registered: AtomicUsize, + outstanding: AtomicUsize, + max_outstanding: AtomicUsize, +} + +impl SplitProbe { + /// The number of splits whose reads have been registered. + pub fn registered(&self) -> usize { + self.registered.load(Ordering::SeqCst) + } + + /// The high-water mark of splits registered but not yet completed. + pub fn max_outstanding(&self) -> usize { + self.max_outstanding.load(Ordering::SeqCst) + } + + fn register(&self) { + self.registered.fetch_add(1, Ordering::SeqCst); + let outstanding = self.outstanding.fetch_add(1, Ordering::SeqCst) + 1; + self.max_outstanding + .fetch_max(outstanding, Ordering::SeqCst); + } + + fn complete(&self) { + self.outstanding.fetch_sub(1, Ordering::SeqCst); + } +} + +/// A reader that splits one row per split and reports registration through a [`SplitProbe`]. +/// +/// Each split projects to a single-element `i32` array holding its own row index, so a scan over +/// it yields `0..row_count`. +#[derive(Debug)] +pub struct ProbeLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + probe: Arc, +} + +impl ProbeLayoutReader { + pub fn new(row_count: usize, probe: Arc) -> Self { + Self { + name: Arc::from("probe"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count: row_count as u64, + probe, + } + } +} + +impl LayoutReader for ProbeLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + for split in (split_range.row_range().start + 1)..=split_range.row_range().end { + splits.push(split_range.row_offset() + split); + } + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &BoundExpression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &BoundExpression, + _mask: MaskFuture, + ) -> VortexResult { + let start = usize::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in usize"))?; + let end = usize::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in usize"))?; + + let values: VortexResult> = (start..end) + .map(|v| i32::try_from(v).map_err(|_| vortex_err!("split value must fit in i32"))) + .collect(); + let array = PrimitiveArray::from_iter(values?).into_array(); + + // Registering the projection is what submits this split's reads, so count it here + // rather than when the returned future is polled. + self.probe.register(); + let probe = Arc::clone(&self.probe); + + Ok(Box::pin(async move { + probe.complete(); + Ok(array) + })) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +/// Asserts that a scan of `row_count` splits yielded `0..row_count` while never holding more +/// than `concurrency` splits registered. +#[expect( + clippy::cast_possible_truncation, + reason = "test row counts are far below i32::MAX" +)] +pub fn assert_bounded_scan( + values: &[i32], + row_count: usize, + probe: &SplitProbe, + concurrency: usize, +) { + let expected: Vec = (0..row_count).map(|v| v as i32).collect(); + assert_eq!(values, expected.as_slice()); + assert_eq!(probe.registered(), row_count); + + let expected_max = cmp::min(concurrency, row_count); + assert_eq!( + probe.max_outstanding(), + expected_max, + "expected at most {expected_max} splits registered at once, \ + which requires split construction to be lazy under the concurrency buffer" + ); +} From f521d646941875956daa436f4e3ec287dfab87fd Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:23:06 -0700 Subject: [PATCH 11/12] docs: record the #87 split-concurrency port and the flat-reader test fix Moves #87 from an unresolved audit note to the applied section, and adds a row for the dafa0960a test-compatibility fix discovered while working on it. --- SPICE_PATCHES.md | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index d1ab4fb93ed..e3efd2e94e9 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -92,8 +92,10 @@ These landed as new commits against `0.85.0` rather than cherry-picks, since the | 4 | Fixed-offset timezone resolution (`resolve_timezone`, threaded through scalar validation, `Display`, and `TemporalMetadata::to_jiff`; `Scalar::try_extension_ref` through the 4 sites that built extension scalars infallibly on a read path) | `dc9892190` | `6cdea73d6` (#75), `5b4bee108` (#78) | `cargo test -p vortex-array --lib extension::datetime::timezone extension::datetime::matcher` — **ran, 42 pass** | | 7, 15, 16 | `vortex.date`→`vortex.timestamp` array + scalar casts via a shared `DateToTimestamp`, `Timestamp::storage_range` replacing Jiff-`Span`-based validation | `dc9892190` (supersedes this branch's earlier `a98a78c1b`, which used a `legacy_session()` workaround instead of the proper `CastKernel` split found in this row's own later Spice history) | `7e5b08151` (#28), this branch's `3c5246867` (#93) | `cargo test -p vortex-array --lib arrays::extension::compute::cast scalar::typed_view::extension` — **ran, all pass**; full crate: `cargo test -p vortex-array -p vortex-datetime-parts --lib` — **3428 + 72 pass, 0 failed** | | 10 (real fix) | `UncompressedSizeInBytes` propagation through `vortex-array/src/arrays/dict/take.rs`, and inclusion in `stats::PRUNING_STATS` | `1672b7328` | `a9ef29dea` (previously **not in the ledger at all** — the old ledger cited only the vendored, out-of-scope `6712e9ffa`) | New tests added (none existed before): `cargo test -p vortex-array --lib arrays::dict::take` — **ran, 2 pass** (`uncompressed_size_in_bytes_propagates_through_take_as_inexact`, `is_constant_propagates_inexactly_through_take_with_a_null_index`); full crate `cargo test -p vortex-array --lib` — **3430 pass, 0 failed**. Not subsumed by the `AggregateFnVTable` redesign — that changes how the stat is computed per array type, not whether `PRUNING_STATS` includes it or `take` propagates it. | +| — | `dafa0960a`'s own cherry-picked tests, adapted to `BoundExpression` | `b18002d7e` | (fixes a gap in this branch's own earlier work, not a fork patch) | The original cherry-pick (`dafa0960a`, "Fix flat reader subrange decode reuse") was only ever `cargo check`'d, not `cargo check --tests`'d, so its test module's incompatibility with `0.85.0`'s `LayoutReader` trait (raw `Expression` where `&BoundExpression` is required, an undefined `SESSION` static, a missing `.into()`) went undetected. `cargo test -p vortex-layout --lib layouts::flat::reader` — **ran, 9 pass** — the first real test run of that fix's actual behavior. | +| #87 | Absolute split concurrency; `MultiLayoutDataSource`'s deferred-open default stops scaling with host cores | `121d5a724` | `4e2d62654` | Confirmed live before porting: `get_available_parallelism()` was still used as a per-worker multiplier in `scan_builder.rs`/`repeated_scan.rs` and as `MultiLayoutDataSource`'s concurrency value directly. Adds `SplitConcurrency` (`PerWorker`/`Absolute`) and makes split-task construction lazy under the configured bound (`RepeatedScan::split_tasks`), so `with_absolute_concurrency(1)` is genuinely serial rather than still registering every split's reads up front. Ported both original regression tests (none existed before) via a new `SplitProbe`/`ProbeLayoutReader` test harness: `cargo test -p vortex-layout --lib` — **255 pass (up from 249), 0 failed**; downstream `cargo test -p vortex-file --lib` — **139 pass, 0 failed**. | -Not ported: the DuckDB `TIMESTAMP_TZ` dtype/scalar `is_utc_timezone` gate (`vortex-duckdb/src/convert/{dtype,scalar}.rs`) — `vortex-duckdb` is not a crate Spice's `Cargo.toml` consumes, directly or transitively. Fork hygiene, not something Spice's re-pin needs; left open. +Not ported: the DuckDB `TIMESTAMP_TZ` dtype/scalar `is_utc_timezone` gate (`vortex-duckdb/src/convert/{dtype,scalar}.rs`) — `vortex-duckdb` is not a crate Spice's `Cargo.toml` consumes, directly or transitively. Fork hygiene, not something Spice's re-pin needs; left open. Also not ported: `4e2d62654`'s vortex-file import-reordering hunks (pure rustfmt churn already reflected in `0.85.0`'s own ordering). ## Untracked in the previous ledger — needs its own audit @@ -137,11 +139,6 @@ file**, not a patch cherry-pick — deprioritized below the smaller items in thi that reason, not because it lacks value. Budget a dedicated session that reads `process_matches::` first. -Also unresolved from the previous audit and not re-checked here: `4e2d62654` (#87, "Absolute -split concurrency, stop deriving concurrency defaults from host parallelism") — thematically -related to row 13 but touches a different file set (`vortex-file/src/read/{driver,request}.rs`, -`vortex-file/src/segments/source.rs`, `vortex-layout/src/scan/*`) and conflicts for real at -`0.85.0`. Needs its own ledger row once re-ported. ## Known remaining exposure From ddc72b818a77220dfdd9e3a0a244a6b13d74da1e Mon Sep 17 00:00:00 2001 From: Viktor Yershov Date: Tue, 22 Sep 2026 17:24:39 -0700 Subject: [PATCH 12/12] docs: confirm row 11 (intra-file decode parallelism) is fully upstreamed vortex-layout/src/scan/split_by.rs at 0.85.0 already carries 26b274c72's subdivide_large_spans function and all 5 of its unit tests, byte-for-byte identical apart from unrelated test-helper API drift -- landed via upstream commit 9d3aafb06, which the row's own "Possibly upstream" guess had already flagged as a candidate origin. vortex-file's end-to-end test is present too. Ran both test suites rather than trusting the diff match alone. Moves the row from "needs re-porting" to "dropped". --- SPICE_PATCHES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index e3efd2e94e9..f96c2b08cc9 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -67,6 +67,7 @@ content in favor of this rewrite, not by a code-level resolution. The actual cod | — | `68badd689` "TableStrategy hides panic... (#8672)" | — | Not a Spice patch: `(cherry picked from commit 4abe3d04f...)`, authored by a Vortex/spiraldb maintainer, cherry-picked onto `spiceai-54` ahead of a release. Already upstream. | | — | `b5498e8e1` "get CI green" (#81), `8d2efe7a7` style import sort | — | CI/lint housekeeping for subsystems Spice doesn't build (java/, python/, ffi/) or trivial reformatting of files already changed upstream. Not carried; not fork state Spice depends on. | | — | `ab4f0b177` "restore load_full + HashMap import dropped in rebase" | — | `vortex-array/src/arc_swap_map.rs` doesn't exist at `0.85.0` at all (not moved, gone), and `grep -rn "load_full"` across `vortex-array/src` finds zero matches. Spice's own fix for a Spice-only prior rebase mistake in a file structure that no longer exists. Nothing to port. | +| 11 | Intra-file decode parallelism | `26b274c72` (#62) | The ledger's own "Possibly upstream" guess (row 11 also named `9d3aafb06`) was right. `vortex-layout/src/scan/split_by.rs` at `0.85.0` already has `subdivide_large_spans`, `MAX_SPLIT_ROWS` and all 5 of the original patch's unit tests, byte-for-byte identical except for unrelated pre-existing test-helper API drift (`Expression`→`BoundExpression`, `SCAN_SESSION`→`session_with_handle`) — landed via upstream commit `9d3aafb06` ("perf(scan): intra-file decode parallelism — sub-split large chunk spans (#8400)"), which `git log -- vortex-layout/src/scan/split_by.rs` shows as already in `0.85.0`'s own history. `vortex-file/src/tests.rs`'s end-to-end test (`test_large_flat_chunk_scan_subdivides_splits`) is present too. **Ran both**: `cargo test -p vortex-layout --lib scan::split_by` (8 pass) and `cargo test -p vortex-file --lib test_large_flat_chunk_scan_subdivides_splits` (1 pass). Nothing to port; the earlier "real conflicts" finding was from attempting to cherry-pick `26b274c72` itself onto content that already matches, not from the feature being missing. | | 6 | Avoid session lock re-entry in writer init | `c536c9aed` (#29) | The hazard was calling the `SessionGuard`-returning `session.arrays()` twice within one statement (`ArrayContext::new(self.session.arrays()...).with_registry(self.session.arrays()...)`), so the two guards overlapped for the whole statement instead of the first being dropped before the second was taken. `vortex-file/src/writer.rs`'s array-context construction was rewritten into `new_array_context`, which calls the lock-free, `Vec`-returning `session.enabled_component_ids(...)` exactly once and reuses the result — the guard-returning call this row exists to guard against isn't made here at all anymore. Swept every `.arrays()` call site across `vortex-file`, `vortex-array`, `vortex-layout`, `vortex-io`, `vortex-scan`, `vortex-btrblocks`, `vortex-utils`, `vortex-session` (`grep -rn "session\.arrays()"`, 10 hits): every one either calls it once or binds the guard to a variable and reuses that variable, never re-invoking the accessor within a live guard's scope. No code change; the hazard's precondition doesn't exist in the crates Spice consumes. Still a **GAP** in the sense that no regression test pins this — if a future edit reintroduces a double call, nothing here would catch it. | | 14 | Arrow Map alias — `DType` alias + map-entry recursion in the session importer | `840e746a9`, `3e7fa40d3` (superseded the stale `1a6dc54f1` origin, then this) | **Superseded by a full upstream rewrite, not just a fix.** Upstream 0.85.0 replaced the `List>` alias both Spice's patch and the pre-split `vortex-array/src/arrow` code used with a **native `Map` array/dtype**: `vortex-array/src/dtype/map.rs` (`MapDType`) and `vortex-array/src/arrays/map/` (`Array`, backed by a `ListView>`). Both `vortex-arrow/src/convert.rs` (import) and `vortex-arrow/src/executor/mod.rs` (export) dispatch `DataType::Map` to this native path, and `vortex-arrow/src/session.rs` recurses into map key/value fields for extension-type metadata exactly as Spice's `3e7fa40d3` did. **Ran `cargo test -p vortex-arrow --lib map` on this branch: 17 tests pass**, including `executor::map::tests::map_roundtrip_preserves_nested_uuid_fields` (the exact extension-recursion scenario `3e7fa40d3` fixed) and `session::tests::schema_roundtrip_preserves_map_uuid_fields`. Upstream's own coverage here is broader than Spice's two patches were. Attempting to cherry-pick the old alias-based patch onto this would be a regression, not a port — the P0 defect (spiceai/spiceai#13524) is closed by a better mechanism than the one Spice shipped. **Remaining gap**: this proves the fix at the `vortex-arrow` crate level; the Spice main repo's own guard (`crates/vortex/src/persistent/mod.rs::map_column_roundtrips_through_a_vortex_file`, per `docs/dev/fork_patches.md`) still needs re-running once this branch is pinned, since that test exercises the full write path through the vendored `vortex-datafusion` sink, which this fork does not build. | @@ -79,7 +80,6 @@ reading the new upstream code and re-implementing the fix, not resolving conflic | # | Patch | Old origin | What changed upstream | |---|---|---|---| -| 11 | Intra-file decode parallelism | `26b274c72` (#62) | Real conflicts in `vortex-file/Cargo.toml`, `vortex-file/src/tests.rs`; `vortex-layout/src/scan/split_by.rs` also changed upstream (small diff, not yet compared line-by-line). Perf-only; still a ledger **GAP** (no guard). | | — | Cache `ArrayKernels` as a per-`ExecutionCtx` snapshot | `ccaa55627` | Real conflict in `vortex-array/src/executor.rs`. Adds a `KernelSnapshot` type and `ArrayKernels::snapshot()` (confirmed absent at `0.85.0`) into the hot `execute_until` path on the central `ExecutionCtx` struct — higher blast radius than the other items here if ported carelessly. Perf-only. Not previously in the ledger. | | — | Restore lint checks on forks | `bb80c537b` | Conflicts across `.github/workflows/ci.yml`, `cast.rs`, `vortex-datafusion/src/persistent/sink.rs` (the last is out of scope). Low priority — CI config for this fork, not behavior Spice depends on. |