diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 7a20b34e675..f96c2b08cc9 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -63,10 +63,13 @@ 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. | +| — | `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. | ## Needs re-porting (real conflicts against `0.85.0` — deferred, not dropped) @@ -77,16 +80,22 @@ 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. | +| — | 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. | -| 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. | + +## 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** | +| 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. 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 @@ -94,22 +103,42 @@ 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. - -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. +`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. + ## Known remaining exposure 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/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/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index e080eb7f04c..1d98d0e011b 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -1,47 +1,166 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; 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::extension::datetime::AnyTemporal; +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) { - // 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>, + target_dtype: &DType, + ctx: &mut ExecutionCtx, +) -> 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()))? + .execute::(ctx)?; + + let converted = cast_date_values_to_timestamp(&source_i64, *source_unit, *target_unit, ctx)?; + + let new_storage = converted + .into_array() + .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 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 &validity { + Validity::NonNullable | Validity::AllValid => { + for &value in input { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(conversion.convert(value)?) }; + } + } + Validity::AllInvalid => { + for _ in 0..input.len() { + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(0i64) }; + } + } + 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 { + 0i64 + }; + // SAFETY: output has sufficient capacity for all pushed values. + unsafe { output.push_unchecked(converted) }; + } + } + } + + Ok(PrimitiveArray::new(output.freeze(), validity)) +} + #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -61,6 +180,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 +219,50 @@ mod tests { assert_eq!(output.dtype(), &new_dtype); } + #[test] + 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()))?; + + 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() -> 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(); + + 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()))?; + + 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] fn cast_different_ext_dtype() { let original_dtype = 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()); +} 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 { 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); 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" + ); +}