diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index 22c2c693fd3..bf54369a166 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -98,6 +98,11 @@ jobs: RUST_BACKTRACE: full VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + VORTEX_IO_URING: "1" + VORTEX_IO_URING_RINGS: "4" + VORTEX_IO_URING_QUEUE_DEPTH: "128" + VORTEX_IO_URING_MIN_READ_SIZE: "0" + VORTEX_IO_URING_MAX_IN_FLIGHT: "512" run: | python3 scripts/random-access-split.py diff --git a/Cargo.lock b/Cargo.lock index b0f98336b0b..e78d6b6d559 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10240,6 +10240,7 @@ dependencies = [ "custom-labels", "futures", "glob", + "io-uring", "itertools 0.14.0", "kanal", "object_store", @@ -10247,6 +10248,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "rstest", + "rustix", "smol", "tempfile", "tokio", @@ -10345,11 +10347,13 @@ dependencies = [ "termtree", "tokio", "tracing", + "vortex-alp", "vortex-array", "vortex-arrow", "vortex-btrblocks", "vortex-buffer", "vortex-error", + "vortex-fastlanes", "vortex-flatbuffers", "vortex-io", "vortex-mask", diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..7448402c4a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -169,6 +169,7 @@ geoarrow = "0.8.0" geoarrow-cast = "0.8.0" get_dir = "0.5.0" glob = "0.3.2" +io-uring = "0.7.13" goldenfile = "1" half = { version = "2.7.1", features = ["std", "num-traits"] } hashbrown = "0.17.1" diff --git a/benchmark-results/pr-9416-random-access/HANDOVER.md b/benchmark-results/pr-9416-random-access/HANDOVER.md new file mode 100644 index 00000000000..d52f964120a --- /dev/null +++ b/benchmark-results/pr-9416-random-access/HANDOVER.md @@ -0,0 +1,88 @@ +# PR #9416 random-access handover + +## Scope + +- PR: https://github.com/vortex-data/vortex/pull/9416 +- Branch: `ji/partial-flat-random-access-batched-ring` +- Worktree: `/mnt/vortex-ssd/worktrees/random-access-pr9416-ring` +- Workload: feature-vectors / uniform, pinned to cores 0-7. +- The decoded ALPRD patch cache was removed. Do not reintroduce a data cache for the I/O comparison. +- The `action/bench-random-access` label was applied after removing the cache. + +## Retained change + +The ALPRD partial-read path no longer calls `clear_stats` on the newly constructed partial array. +Those `BitPacked`, `Patches`, `ALPRD`, and `FixedSizeList` arrays begin with empty statistics. The +fixed-width path still clears statistics inherited from its serialized array tree. + +Three five-second cached runs on 8 cores: + +| Variant | Vortex runs | Median | Lance median | +|---|---|---:|---:| +| No-cache baseline | 1.612 / 1.621 / 1.680 ms | 1.621 ms | 1.059 ms | +| Skip redundant ALPRD stats clear | 1.594 / 1.602 / 1.607 ms | 1.601 ms | 1.058 ms | + +This is a 1.2% Vortex improvement and changes neither I/O nor serialization. A final cold rerun is +still desirable, although the change occurs entirely after I/O. + +## Segment and patch measurements + +The complete 105-row measurement is in `feature-vectors-uniform-segments.csv`. + +- Uniform selects 105 physical Flat segments. +- Every segment has 256 vectors x 1,024 values = 262,144 values. +- Left buffer: 98,304 bytes per segment; right buffer: 753,664 bytes per segment. +- Patches per segment: min 953, mean 1,060.70, max 1,793. +- Patch density: 0.404627%. +- Patch bytes per segment: min 5,718, mean 6,364.23, max 10,758. +- Totals: 27,525,120 resident values, 111,374 patches, 668,244 patch bytes. +- The query selects 107,520 values and should intersect only about 435 patches, but the reader + reconstructs all 111,374 patches before slicing. +- Main data is already row-sliced: 384 left bytes + 2,944 right bytes per selected vector, or + 349,440 bytes total. Patch buffers remain unsliced and add 668,244 bytes. + +Conclusion: 256-vector Flat segments are not the main problem because left/right buffers support +partial row reads. Patch read/reconstruction granularity is the leak. + +## Rejected experiments + +1. `OnceLock` cache: hot 1.491 ms, but skipped patch reads and decode on repeated takes. + Removed as an invalid apples-to-apples I/O comparison. +2. Search bitpacked patch indices without bulk decode: median 1.640 ms versus 1.621 ms baseline. + The scalar probes were slower; removed. +3. Merge patch-index/value reads: median 1.741 ms. `FileSegmentSource` already coalesces nearby + reads, so this added slicing work; removed, including the temporary buffer API. +4. Latency-based blocking/io_uring routing: about 1.28 ms hot but about 26.7 ms cold because a small + cold metadata read falsely classified the file as hot; removed. + +## Profile and next work + +A 10-second no-cache Samply profile kept all eight Tokio workers busy. Dominant resolved self +frames were AArch64 atomics and mimalloc allocation/free, pointing to task/array/future construction +overhead rather than worker serialization. The local profile is +`/tmp/no-cache-feature-uniform.profile.json.gz` and is not portable with this branch. + +Next, count allocations and short-lived objects inside `resolve_alprd_pages` and final +canonicalization. Avoid two-stage patch I/O unless separately proven hot and cold: it can reduce +patch-value bytes but adds an I/O round. If a future format change is allowed, serialize ALPRD patch +chunk offsets. The 1,024-value patch chunk exactly matches one feature vector, enabling direct +lookup of the roughly four relevant patches. + +## Benchmark environment + +```text +taskset -c 0-7 +TOKIO_WORKER_THREADS=8 +RAYON_NUM_THREADS=8 +LANCE_IO_THREADS=8 +VORTEX_EXPERIMENTAL_PATCHED_ARRAY=1 +FLAT_LAYOUT_INLINE_ARRAY_NODE=1 +VORTEX_IO_URING=1 +VORTEX_IO_URING_RINGS=4 +VORTEX_IO_URING_QUEUE_DEPTH=128 +VORTEX_IO_URING_MIN_READ_SIZE=0 +VORTEX_IO_URING_MAX_IN_FLIGHT=512 +``` + +Validation completed: nightly formatting, `git diff --check`, `cargo check -p vortex-layout`, +release benchmark builds, feature-vectors/uniform smoke runs, and repeated hot comparisons. diff --git a/benchmark-results/pr-9416-random-access/feature-vectors-uniform-segments.csv b/benchmark-results/pr-9416-random-access/feature-vectors-uniform-segments.csv new file mode 100644 index 00000000000..6aedd5dbb73 --- /dev/null +++ b/benchmark-results/pr-9416-random-access/feature-vectors-uniform-segments.csv @@ -0,0 +1,106 @@ +segment_id,rows,values,patches,patch_bytes,left_bytes,right_bytes +26,256,262144,1031,6186,98304,753664 +62,256,262144,1018,6108,98304,753664 +81,256,262144,994,5964,98304,753664 +98,256,262144,993,5958,98304,753664 +99,256,262144,1011,6066,98304,753664 +119,256,262144,1007,6042,98304,753664 +124,256,262144,1017,6102,98304,753664 +220,256,262144,1022,6132,98304,753664 +231,256,262144,1018,6108,98304,753664 +232,256,262144,1049,6294,98304,753664 +337,256,262144,1041,6246,98304,753664 +439,256,262144,1067,6402,98304,753664 +475,256,262144,985,5910,98304,753664 +479,256,262144,1000,6000,98304,753664 +482,256,262144,1038,6228,98304,753664 +495,256,262144,1034,6204,98304,753664 +502,256,262144,1027,6162,98304,753664 +503,256,262144,1039,6234,98304,753664 +595,256,262144,1022,6132,98304,753664 +627,256,262144,1541,9246,98304,753664 +686,256,262144,987,5922,98304,753664 +704,256,262144,1016,6096,98304,753664 +712,256,262144,988,5928,98304,753664 +744,256,262144,1051,6306,98304,753664 +785,256,262144,980,5880,98304,753664 +808,256,262144,1054,6324,98304,753664 +815,256,262144,1008,6048,98304,753664 +821,256,262144,1014,6084,98304,753664 +840,256,262144,1089,6534,98304,753664 +849,256,262144,1015,6090,98304,753664 +865,256,262144,1030,6180,98304,753664 +908,256,262144,995,5970,98304,753664 +1063,256,262144,999,5994,98304,753664 +1081,256,262144,1775,10650,98304,753664 +1107,256,262144,1067,6402,98304,753664 +1131,256,262144,1027,6162,98304,753664 +1143,256,262144,1573,9438,98304,753664 +1148,256,262144,1051,6306,98304,753664 +1187,256,262144,1068,6408,98304,753664 +1227,256,262144,1041,6246,98304,753664 +1329,256,262144,993,5958,98304,753664 +1354,256,262144,973,5838,98304,753664 +1356,256,262144,1003,6018,98304,753664 +1374,256,262144,1045,6270,98304,753664 +1406,256,262144,963,5778,98304,753664 +1436,256,262144,1056,6336,98304,753664 +1460,256,262144,1000,6000,98304,753664 +1573,256,262144,953,5718,98304,753664 +1599,256,262144,996,5976,98304,753664 +1632,256,262144,1008,6048,98304,753664 +1667,256,262144,1033,6198,98304,753664 +1763,256,262144,1067,6402,98304,753664 +1778,256,262144,1014,6084,98304,753664 +1786,256,262144,1002,6012,98304,753664 +1816,256,262144,1017,6102,98304,753664 +1834,256,262144,1048,6288,98304,753664 +1869,256,262144,1042,6252,98304,753664 +1910,256,262144,1039,6234,98304,753664 +1953,256,262144,1011,6066,98304,753664 +1993,256,262144,1061,6366,98304,753664 +2010,256,262144,1078,6468,98304,753664 +2071,256,262144,1011,6066,98304,753664 +2091,256,262144,1047,6282,98304,753664 +2096,256,262144,1022,6132,98304,753664 +2146,256,262144,993,5958,98304,753664 +2176,256,262144,1030,6180,98304,753664 +2190,256,262144,1058,6348,98304,753664 +2236,256,262144,980,5880,98304,753664 +2238,256,262144,1064,6384,98304,753664 +2270,256,262144,1053,6318,98304,753664 +2282,256,262144,997,5982,98304,753664 +2389,256,262144,1064,6384,98304,753664 +2447,256,262144,1000,6000,98304,753664 +2484,256,262144,1046,6276,98304,753664 +2501,256,262144,1050,6300,98304,753664 +2567,256,262144,999,5994,98304,753664 +2569,256,262144,1051,6306,98304,753664 +2585,256,262144,1018,6108,98304,753664 +2639,256,262144,1507,9042,98304,753664 +2659,256,262144,1013,6078,98304,753664 +2694,256,262144,1533,9198,98304,753664 +2763,256,262144,1078,6468,98304,753664 +2842,256,262144,1067,6402,98304,753664 +2981,256,262144,1037,6222,98304,753664 +3001,256,262144,982,5892,98304,753664 +3007,256,262144,1018,6108,98304,753664 +3036,256,262144,990,5940,98304,753664 +3057,256,262144,973,5838,98304,753664 +3067,256,262144,962,5772,98304,753664 +3068,256,262144,1016,6096,98304,753664 +3219,256,262144,994,5964,98304,753664 +3298,256,262144,1034,6204,98304,753664 +3321,256,262144,983,5898,98304,753664 +3384,256,262144,1029,6174,98304,753664 +3410,256,262144,1012,6072,98304,753664 +3422,256,262144,1028,6168,98304,753664 +3578,256,262144,1032,6192,98304,753664 +3622,256,262144,1101,6606,98304,753664 +3637,256,262144,1036,6216,98304,753664 +3639,256,262144,975,5850,98304,753664 +3682,256,262144,1793,10758,98304,753664 +3760,256,262144,996,5976,98304,753664 +3849,256,262144,1480,8880,98304,753664 +3851,256,262144,999,5994,98304,753664 +3865,256,262144,1039,6234,98304,753664 diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 49a9bf92f58..fd1dc33f4a3 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -289,10 +289,14 @@ async fn register_v2_tables( .runtime_env() .object_store(table_url.object_store())?; - let fs: FileSystemRef = Arc::new(ObjectStoreFileSystem::new( - Arc::clone(&store), - SESSION.handle(), - )); + let fs: FileSystemRef = if benchmark_base.scheme() == "file" { + Arc::new(ObjectStoreFileSystem::local(SESSION.handle())) + } else { + Arc::new(ObjectStoreFileSystem::new( + Arc::clone(&store), + SESSION.handle(), + )) + }; let base_prefix = benchmark_base.path().trim_start_matches('/').to_string(); let fs = fs.with_prefix(base_prefix); diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index b6ee50d7b1f..f7c954d2222 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -71,6 +71,33 @@ pub struct ALPRDMetadata { patches: Option, } +impl ALPRDMetadata { + pub fn right_bit_width(&self) -> VortexResult { + u8::try_from(self.right_bit_width).map_err(|_| { + vortex_err!( + "right bit width {} does not fit in u8", + self.right_bit_width + ) + }) + } + + pub fn left_parts_dictionary(&self) -> VortexResult> { + self.dict + .get(..usize::try_from(self.dict_len)?) + .ok_or_else(|| vortex_err!("ALPRD dictionary length is out of bounds"))? + .iter() + .map(|&value| { + u16::try_from(value) + .map_err(|_| vortex_err!("ALPRD dictionary value {value} does not fit in u16")) + }) + .collect() + } + + pub fn patches(&self) -> Option<&PatchesMetadata> { + self.patches.as_ref() + } +} + impl ArrayHash for ALPRDData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.left_parts_dictionary.array_hash(state, accuracy); diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..115be03d89f 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -19,6 +19,7 @@ mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; pub use vtable::BitPacked; pub use vtable::BitPackedArray; +pub use vtable::BitPackedMetadata; pub(crate) fn initialize(session: &vortex_session::VortexSession) { vtable::initialize(session); diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..8be75c8e553 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -72,6 +72,22 @@ pub struct BitPackedMetadata { pub(crate) patches: Option, } +impl BitPackedMetadata { + pub fn bit_width(&self) -> VortexResult { + u8::try_from(self.bit_width) + .map_err(|_| vortex_err!("bit width {} does not fit in u8", self.bit_width)) + } + + pub fn offset(&self) -> VortexResult { + u16::try_from(self.offset) + .map_err(|_| vortex_err!("bit-packed offset {} does not fit in u16", self.offset)) + } + + pub fn patches(&self) -> Option<&PatchesMetadata> { + self.patches.as_ref() + } +} + impl ArrayHash for BitPackedData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.offset.hash(state); diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index bfd43cdb401..6ee096e747e 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -14,6 +14,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::List; +pub use vtable::ListMetadata; pub(crate) fn initialize(session: &vortex_session::VortexSession) { compute::initialize(session); diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index c55e7050351..bd508817e86 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -52,6 +52,12 @@ pub struct ListMetadata { offset_ptype: i32, } +impl ListMetadata { + pub fn elements_len(&self) -> u64 { + self.elements_len + } +} + impl ArrayHash for ListData { fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 48c70d9446a..bf2356d1058 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -77,6 +77,32 @@ fn reduce_struct_cast( )) } +pub(crate) fn struct_get_item( + array: ArrayView<'_, Struct>, + field_name: &crate::dtype::FieldName, +) -> VortexResult { + let field = array + .unmasked_field_by_name_opt(field_name) + .ok_or_else(|| { + vortex_err!( + "Field {} missing from struct array {}", + field_name, + array.struct_fields().names() + ) + })?; + + match array.validity()? { + Validity::NonNullable => Ok(field.clone()), + Validity::AllValid => field.clone().cast(field.dtype().as_nullable()), + Validity::AllInvalid => Ok(ConstantArray::new( + Scalar::null(field.dtype().as_nullable()), + field.len(), + ) + .into_array()), + Validity::Array(mask) => field.clone().mask(mask), + } +} + /// Rule to flatten get_item from struct by field name #[derive(Debug)] pub(crate) struct StructGetItemRule; @@ -90,38 +116,7 @@ impl ArrayParentReduceRule for StructGetItemRule { parent: ScalarFnArrayView<'_, GetItem>, _child_idx: usize, ) -> VortexResult> { - let field_name = parent.options; - let field = child - .unmasked_field_by_name_opt(field_name) - .ok_or_else(|| { - vortex_err!( - "Field '{}' missing from struct array {}", - field_name, - child.struct_fields().names() - ) - })?; - - match child.validity()? { - Validity::NonNullable => { - // If the struct is non-nullable, the field's validity is unchanged - Ok(Some(field.clone())) - } - Validity::AllValid => { - // If struct is nullable, field must also be nullable - field.clone().cast(field.dtype().as_nullable()).map(Some) - } - Validity::AllInvalid => { - // If everything is invalid, the field is also all invalid - Ok(Some( - ConstantArray::new(Scalar::null(field.dtype().as_nullable()), field.len()) - .into_array(), - )) - } - Validity::Array(mask) => { - // If the validity is an array, we need to combine it with the field's validity - field.clone().mask(mask).map(Some) - } - } + struct_get_item(child, parent.options).map(Some) } } diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d0590f2bf58..5f4bbd1dba6 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -9,10 +9,16 @@ use crate::ArrayRef; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::struct_::compute::rules::struct_get_item; use crate::expr::BoundExpression; use crate::expr::Expression; use crate::optimizer::ArrayOptimizer; +use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::fns::pack::Pack; +use crate::validity::Validity; impl ArrayRef { /// Apply a bound expression to this array, producing a new array in constant time. @@ -35,6 +41,24 @@ impl ArrayRef { .map(|child| self.clone().apply_bound(child)) .try_collect()?; + if let Some(field_name) = scalar_fn.as_opt::() + && let [child] = children.as_slice() + && let Some(array) = child.as_opt::() + { + return struct_get_item(array, field_name); + } + + if let Some(pack) = scalar_fn.as_opt::() { + let validity = match pack.nullability { + crate::dtype::Nullability::NonNullable => Validity::NonNullable, + crate::dtype::Nullability::Nullable => Validity::AllValid, + }; + return Ok( + StructArray::try_new(pack.names.clone(), children, self.len(), validity)? + .into_array(), + ); + } + let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); diff --git a/vortex-array/src/mask_future.rs b/vortex-array/src/mask_future.rs index a46107623c1..767809162cd 100644 --- a/vortex-array/src/mask_future.rs +++ b/vortex-array/src/mask_future.rs @@ -20,6 +20,9 @@ use vortex_mask::Mask; pub struct MaskFuture { inner: Shared>>, len: usize, + upper_bound: Option, + upper_bound_is_exact: bool, + partial_reads_allowed: bool, } impl MaskFuture { @@ -40,6 +43,9 @@ impl MaskFuture { .boxed() .shared(), len, + upper_bound: None, + upper_bound_is_exact: false, + partial_reads_allowed: false, } } @@ -55,7 +61,12 @@ impl MaskFuture { /// Create a MaskFuture from a ready mask. pub fn ready(mask: Mask) -> Self { - Self::new(mask.len(), async move { Ok(mask) }) + let upper_bound = mask.clone(); + let mut future = Self::new(mask.len(), async move { Ok(mask) }); + future.upper_bound = Some(upper_bound); + future.upper_bound_is_exact = true; + future.partial_reads_allowed = true; + future } /// Create a MaskFuture that resolves to a mask with all values set to true. @@ -72,7 +83,57 @@ impl MaskFuture { } let inner = self.inner.clone(); - Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }) + let upper_bound = self + .upper_bound + .as_ref() + .map(|upper_bound| upper_bound.slice(range.clone())); + let mut sliced = Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }); + sliced.upper_bound = upper_bound; + sliced.upper_bound_is_exact = self.upper_bound_is_exact; + sliced.partial_reads_allowed = self.partial_reads_allowed; + sliced + } + + /// Attach a conservative upper bound for the mask resolved by this future. + /// + /// Readers can use this to register I/O eagerly without waiting for filter evaluation. The + /// resolved mask must not contain a true row that is false in `upper_bound`. + pub fn with_upper_bound(mut self, upper_bound: Mask) -> Self { + assert_eq!( + upper_bound.len(), + self.len, + "MaskFuture upper bound length mismatch" + ); + self.upper_bound = Some(upper_bound); + self.upper_bound_is_exact = false; + self + } + + /// Return the conservative upper bound for this future, when one is known. + pub fn upper_bound(&self) -> Option<&Mask> { + self.upper_bound.as_ref() + } + + /// Return whether the upper bound is the exact mask returned by this future. + pub fn upper_bound_is_exact(&self) -> bool { + self.upper_bound_is_exact + } + + /// Permit readers to satisfy this selection using partial segment reads. + pub fn with_partial_reads(mut self) -> Self { + self.partial_reads_allowed = true; + self + } + + /// Prevent readers from turning this mask into partial segment reads. + pub fn without_partial_reads(mut self) -> Self { + self.partial_reads_allowed = false; + self + } + + /// Return whether readers may satisfy this selection using partial segment reads. + pub fn partial_reads_allowed(&self) -> bool { + self.partial_reads_allowed } pub fn inspect( @@ -84,6 +145,9 @@ impl MaskFuture { Self { inner: self.inner.inspect(f).boxed().shared(), len, + upper_bound: self.upper_bound, + upper_bound_is_exact: self.upper_bound_is_exact, + partial_reads_allowed: self.partial_reads_allowed, } } } @@ -119,8 +183,26 @@ mod tests { let partial = fut.slice(0..mask.len() - 1); assert_eq!(partial.len(), mask.len() - 1); + assert_eq!(partial.upper_bound(), Some(&mask.slice(0..mask.len() - 1))); assert_eq!(partial.await?, mask.slice(0..mask.len() - 1)); Ok(()) }) } + + #[test] + fn new_future_has_no_upper_bound_until_attached() { + let future = MaskFuture::new(3, async { Ok(Mask::new_false(3)) }); + assert!(future.upper_bound().is_none()); + + let upper_bound = Mask::from_indices(3, [0, 2]); + let future = future.with_upper_bound(upper_bound.clone()); + assert_eq!(future.upper_bound(), Some(&upper_bound)); + assert!(!future.upper_bound_is_exact()); + } + + #[test] + fn ready_future_has_an_exact_upper_bound() { + let future = MaskFuture::ready(Mask::from_indices(3, [0, 2])); + assert!(future.upper_bound_is_exact()); + } } diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 641efd61404..cbfd7c711bc 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -264,11 +264,18 @@ impl Patches { // Perform validation of components when they are host-resident. // This is not possible to do eagerly when the data is on GPU memory. if indices.is_host() && values.is_host() { - let max = usize::try_from(&indices.execute_scalar( - indices.len() - 1, - &mut legacy_session().create_execution_ctx(), - )?) - .map_err(|_| vortex_err!("indices must be a number"))?; + let max = if let Some(primitive) = indices.as_opt::() { + match_each_unsigned_integer_ptype!(primitive.ptype(), |T| { + NumCast::from(primitive.as_slice::()[primitive.len() - 1]) + .ok_or_else(|| vortex_err!("indices must be a number")) + }) + } else { + usize::try_from(&indices.execute_scalar( + indices.len() - 1, + &mut legacy_session().create_execution_ctx(), + )?) + .map_err(|_| vortex_err!("indices must be a number")) + }?; vortex_ensure!( max - offset < array_len, "Patch indices {max:?}, offset {offset} are longer than the array length {array_len}" diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f84ec182269..616bf46f5e3 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -5,6 +5,7 @@ use std::borrow::Cow; use std::fmt::Debug; use std::fmt::Formatter; use std::iter; +use std::ops::Range; use std::sync::Arc; use flatbuffers::FlatBufferBuilder; @@ -294,6 +295,31 @@ pub struct SerializedArray { buffers: Arc<[BufferHandle]>, } +/// Location and alignment of one serialized array buffer within its containing segment. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SerializedBuffer { + index: usize, + range: Range, + alignment: Alignment, +} + +impl SerializedBuffer { + /// Return this buffer's index in the serialized array's global buffer table. + pub fn index(&self) -> usize { + self.index + } + + /// Return this buffer's byte range within the containing segment. + pub fn range(&self) -> &Range { + &self.range + } + + /// Return the alignment required when materializing this buffer independently. + pub fn alignment(&self) -> Alignment { + self.alignment + } +} + impl Debug for SerializedArray { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("SerializedArray") @@ -516,6 +542,46 @@ impl SerializedArray { .unwrap_or_default() } + /// Return the global buffer indices referenced by this array node. + pub fn buffer_indices(&self) -> Vec { + self.flatbuffer() + .buffers() + .map(|buffers| buffers.iter().map(usize::from).collect()) + .unwrap_or_default() + } + + /// Return validated locations for all data buffers in their serialized segment. + pub fn buffer_descriptors(&self) -> VortexResult> { + let fb_array = root::(self.flatbuffer.as_ref())?; + let mut offset = 0usize; + fb_array + .buffers() + .unwrap_or_default() + .iter() + .enumerate() + .map(|(index, buffer)| { + if buffer.compression() != Compression::None { + vortex_bail!( + "Partial reads do not support serialized buffer compression {:?}", + buffer.compression() + ); + } + let start = offset + .checked_add(usize::from(buffer.padding())) + .ok_or_else(|| vortex_err!("Buffer {index} padding overflows"))?; + let end = start + .checked_add(buffer.length() as usize) + .ok_or_else(|| vortex_err!("Buffer {index} length overflows"))?; + offset = end; + Ok(SerializedBuffer { + index, + range: start..end, + alignment: Alignment::try_from_untrusted_exponent(buffer.alignment_exponent())?, + }) + }) + .collect() + } + /// Validate and align the array tree flatbuffer, returning the aligned buffer and root location. fn validate_array_tree(array_tree: impl Into) -> VortexResult<(FlatBuffer, usize)> { let fb_buffer = FlatBuffer::align_from(array_tree.into()); @@ -545,6 +611,13 @@ impl SerializedArray { }) } + /// Attach pre-resolved buffer handles to this already-validated array tree. + pub fn with_buffers(&self, buffers: Vec) -> Self { + let mut serialized = self.clone(); + serialized.buffers = buffers.into(); + serialized + } + /// Create an [`SerializedArray`] from a raw array tree flatbuffer (metadata only). /// /// This constructor creates a `SerializedArray` with no buffer data, useful for diff --git a/vortex-array/src/stats/array.rs b/vortex-array/src/stats/array.rs index d3fb7fd11e4..32398e12eb7 100644 --- a/vortex-array/src/stats/array.rs +++ b/vortex-array/src/stats/array.rs @@ -63,6 +63,11 @@ impl ArrayStats { self.inner.write().clear(stat); } + /// Clears every cached statistic. + pub fn clear_all(&self) { + *self.inner.write() = StatsSet::default(); + } + pub fn retain(&self, stats: &[Stat]) { self.inner.write().retain_only(stats); } @@ -250,6 +255,11 @@ impl StatsSetRef<'_> { self.array_stats.clear(stat); } + /// Clears every cached statistic from this array. + pub fn clear_all(&self) { + self.array_stats.clear_all(); + } + pub fn compute_min TryFrom<&'a Scalar, Error = VortexError>>( &self, ctx: &mut ExecutionCtx, diff --git a/vortex-arrow/src/executor/primitive.rs b/vortex-arrow/src/executor/primitive.rs index 6cea08c0a72..605bd8b462d 100644 --- a/vortex-arrow/src/executor/primitive.rs +++ b/vortex-arrow/src/executor/primitive.rs @@ -12,7 +12,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use crate::null_buffer::to_null_buffer; @@ -41,8 +40,15 @@ pub(super) fn to_arrow_primitive( where T::Native: NativePType, { - // We use nullable here so we can essentially ignore nullability during the cast. - let array = array.cast(DType::Primitive(T::Native::PTYPE, Nullability::Nullable))?; + // Arrow's physical primitive type is independent of field nullability. Preserve the + // array's existing nullability so already-correct encoded arrays do not pay for a recursive + // metadata-only cast before execution. + let target_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability()); + let array = if array.dtype() == &target_dtype { + array + } else { + array.cast(target_dtype)? + }; let primitive = array.execute::(ctx)?; canonical_primitive_to_arrow::(primitive, ctx) } diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 69d3851d079..f1cd4ad7be4 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -25,16 +25,19 @@ use crate::segments::RequestMetrics; pin_project! { /// Converts request lifecycle events into batches of physical reads. /// - /// Polled requests become eligible in registration order. An eligible request may absorb nearby - /// registered requests according to `coalesce_window`. Each poll emits every physical read - /// currently available, up to `batch_size`; it never waits for a full batch. + /// Takes an input stream of [`ReadRequest`]s and buffers all ready requests into local state. + /// When polled for the next request, this stream will choose the next best request based on + /// an ordering of `(has_been_polled, insertion_order)`, skipping any canceled requests, and + /// then coalescing with other nearby requests within the configured `window`. + /// + /// The output contains up to `batch_size` immediately eligible physical requests. A poll never + /// waits to fill a batch. pub(crate) struct IoRequestStream { #[pin] events: S, // True after the event source closes; buffered requests may still remain. inner_done: bool, coalesce_window: Option, - // Maximum physical reads returned by one stream item. batch_size: usize, state: State, } @@ -90,7 +93,7 @@ where } } - // Emit a partial batch immediately so the downstream driver can fill free I/O slots. + // Return up to batch_size requests that are eligible now. Do not wait to fill the batch. let mut batch = Vec::with_capacity(*this.batch_size); while batch.len() < *this.batch_size { let Some(request) = this.state.next(this.coalesce_window.as_ref()) else { @@ -146,14 +149,7 @@ impl State { fn on_event(&mut self, event: ReadEvent) { trace!(?event, "Received ReadEvent"); match event { - ReadEvent::Request(req) => { - if req.callback.is_closed() { - trace!(?req, "ReadRequest dropped before registration"); - return; - } - self.requests_by_offset.insert((req.offset, req.id)); - self.requests.insert(req.id, req); - } + ReadEvent::Request(req) => self.register(req), ReadEvent::Polled(req_id) => { if let Some(req) = self.requests.remove(&req_id) { if req.callback.is_closed() { @@ -177,6 +173,15 @@ impl State { } } + fn register(&mut self, request: ReadRequest) { + if request.callback.is_closed() { + trace!(?request, "ReadRequest dropped before registration"); + return; + } + self.requests_by_offset.insert((request.offset, request.id)); + self.requests.insert(request.id, request); + } + /// Get the next request, if any. fn next(&mut self, coalesce_window: Option<&CoalesceConfig>) -> Option { match coalesce_window { @@ -226,6 +231,10 @@ impl State { let first_req = self.next_uncoalesced()?; let mut requests = vec![first_req]; + let mut coalesce_distance = requests[0] + .coalesce_distance + .unwrap_or(window.distance) + .min(window.distance); let mut current_start = requests[0].offset; let mut current_end = requests[0].offset + requests[0].length as u64; let align = *self.coalesced_buffer_alignment as u64; @@ -241,8 +250,8 @@ impl State { found_new_requests = false; // Find the range we should scan for coalescing in this iteration - let scan_start = current_start.saturating_sub(window.distance); - let scan_end = current_end.saturating_add(window.distance); + let scan_start = current_start.saturating_sub(coalesce_distance); + let scan_end = current_end.saturating_add(coalesce_distance); // Look for requests that can be coalesced with our current range for &(req_offset, req_id) in self @@ -270,8 +279,12 @@ impl State { // Check if this request is within coalescing distance of our current range let req_end = req_offset + req.length as u64; - if (req_offset <= current_end + window.distance && req_end >= current_start) - || (req_end + window.distance >= current_start && req_offset <= current_end) + let request_distance = req + .coalesce_distance + .unwrap_or(window.distance) + .min(coalesce_distance); + if (req_offset <= current_end + request_distance && req_end >= current_start) + || (req_end + request_distance >= current_start && req_offset <= current_end) { // Calculate what the new range would be if we include this request let new_start = current_start.min(req_offset); @@ -286,6 +299,7 @@ impl State { current_start = new_start; current_end = new_end; + coalesce_distance = request_distance; let req = self .polled_requests .remove(&req_id) @@ -361,6 +375,7 @@ mod tests { offset, length, alignment: Alignment::none(), + coalesce_distance: None, callback: tx, }, rx, @@ -521,6 +536,56 @@ mod tests { } } + #[tokio::test] + async fn test_file_profile_coalesces_adjacent_pages() { + const PAGE_SIZE: usize = 64 * 1024; + let (mut req1, _rx1) = create_request(1, 0, PAGE_SIZE); + let (mut req2, _rx2) = create_request(2, PAGE_SIZE as u64, PAGE_SIZE); + req1.coalesce_distance = Some(16 * 1024); + req2.coalesce_distance = Some(16 * 1024); + + let outputs = collect_outputs( + vec![ + ReadEvent::Request(req1), + ReadEvent::Request(req2), + ReadEvent::Polled(1), + ReadEvent::Polled(2), + ], + Some(CoalesceConfig::file()), + ) + .await; + + assert_eq!(outputs.len(), 1); + assert_eq!(outputs[0].range(), 0..(2 * PAGE_SIZE) as u64); + } + + #[tokio::test] + async fn test_file_profile_does_not_cross_unrequested_page() { + const PAGE_SIZE: usize = 64 * 1024; + let (mut req1, _rx1) = create_request(1, 0, PAGE_SIZE); + let (mut req2, _rx2) = create_request(2, (2 * PAGE_SIZE) as u64, PAGE_SIZE); + req1.coalesce_distance = Some(16 * 1024); + req2.coalesce_distance = Some(16 * 1024); + + let outputs = collect_outputs( + vec![ + ReadEvent::Request(req1), + ReadEvent::Request(req2), + ReadEvent::Polled(1), + ReadEvent::Polled(2), + ], + Some(CoalesceConfig::file()), + ) + .await; + + assert_eq!(outputs.len(), 2); + assert_eq!(outputs[0].range(), 0..PAGE_SIZE as u64); + assert_eq!( + outputs[1].range(), + (2 * PAGE_SIZE) as u64..(3 * PAGE_SIZE) as u64 + ); + } + #[tokio::test] async fn test_coalesce_with_gap() { let (req1, _rx1) = create_request(1, 0, 10); @@ -558,6 +623,7 @@ mod tests { offset: 6, length: 5, alignment: Alignment::new(2), + coalesce_distance: None, callback: tx1, }; let req2 = ReadRequest { @@ -565,6 +631,7 @@ mod tests { offset: 12, length: 1, alignment: Alignment::new(4), + coalesce_distance: None, callback: tx2, }; @@ -633,6 +700,7 @@ mod tests { offset: 0, length: 10, alignment: Alignment::none(), + coalesce_distance: None, callback: tx1, }; let req2 = ReadRequest { @@ -640,6 +708,7 @@ mod tests { offset: 100, length: 10, alignment: Alignment::none(), + coalesce_distance: None, callback: tx2, }; @@ -669,6 +738,7 @@ mod tests { offset: 10, length: 4, alignment: Alignment::none(), + coalesce_distance: None, callback: tx1, }; state.on_event(ReadEvent::Request(req1)); @@ -683,6 +753,7 @@ mod tests { offset: 20, length: 8, alignment: Alignment::none(), + coalesce_distance: None, callback: tx2, }; state.on_event(ReadEvent::Request(req2)); diff --git a/vortex-file/src/read/request.rs b/vortex-file/src/read/request.rs index c4bb4bdc975..a68de2862ed 100644 --- a/vortex-file/src/read/request.rs +++ b/vortex-file/src/read/request.rs @@ -55,6 +55,17 @@ impl IoRequest { } } + /// Whether this physical request was assembled exclusively from partial segment ranges. + pub(crate) fn is_partial(&self) -> bool { + match &self.0 { + IoRequestInner::Single(request) => request.coalesce_distance.is_some(), + IoRequestInner::Coalesced(request) => request + .requests + .iter() + .all(|request| request.coalesce_distance.is_some()), + } + } + /// Resolves the request with the given result. pub fn resolve(self, result: VortexResult) { match self.0 { @@ -96,6 +107,8 @@ pub struct ReadRequest { pub(crate) offset: u64, pub(crate) length: usize, pub(crate) alignment: Alignment, + /// Optional per-request cap on the empty gap this request may coalesce across. + pub(crate) coalesce_distance: Option, pub(crate) callback: oneshot::Sender>, } @@ -106,6 +119,7 @@ impl Debug for ReadRequest { .field("offset", &self.offset) .field("length", &self.length) .field("alignment", &self.alignment) + .field("coalesce_distance", &self.coalesce_distance) .field("is_closed", &self.callback.is_closed()) .finish() } diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 33aed2fa9d2..3358d6125b6 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -4,22 +4,21 @@ use std::any::Any; use std::collections::VecDeque; use std::future::Future; +use std::ops::Range; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use futures::FutureExt; -use futures::Stream; use futures::StreamExt; use futures::channel::mpsc; use futures::future; use futures::future::BoxFuture; use futures::future::Shared; -use futures::stream::BoxStream; -use futures::stream::Fuse; use futures::stream::SelectAll; use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; @@ -30,7 +29,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_io::ReadAtRequest; -use vortex_io::ReadAtStream; use vortex_io::VortexReadAt; use vortex_io::runtime::Handle; use vortex_io::runtime::JoinOutcome; @@ -93,6 +91,26 @@ type SharedDriver = Shared>; /// observe completion takes the payload and re-raises it; later readers report a graceful error. type DriverPanic = Arc>>>; +const MAX_PARTIAL_SUBMISSION_REQUESTS: usize = 512; +const MAX_PARTIAL_SUBMISSION_BYTES: usize = 16 << 20; + +fn partial_submission_len(requests: &VecDeque) -> usize { + let mut count = 0usize; + let mut bytes = 0usize; + for request in requests.iter().take(MAX_PARTIAL_SUBMISSION_REQUESTS) { + if !request.is_partial() { + break; + } + let next_bytes = bytes.saturating_add(request.len()); + if count > 0 && next_bytes > MAX_PARTIAL_SUBMISSION_BYTES { + break; + } + count += 1; + bytes = next_bytes; + } + count +} + fn validate_read_result( request: &IoRequest, result: VortexResult, @@ -110,171 +128,6 @@ fn validate_read_result( }) } -type IoBatchStream = Fuse>>; - -enum ReadRangeResultsState { - Reading(ReadAtStream), - Missing, -} - -/// Matches streamed range results back to their logical requests. -struct ReadRangeResults { - state: ReadRangeResultsState, - remaining: Vec>, -} - -impl ReadRangeResults { - fn new(results: ReadAtStream, requests: Vec) -> Self { - Self { - state: ReadRangeResultsState::Reading(results), - remaining: requests.into_iter().map(Some).collect(), - } - } -} - -impl Stream for ReadRangeResults { - type Item = (IoRequest, VortexResult); - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - loop { - match &mut self.state { - ReadRangeResultsState::Reading(results) => match results.poll_next_unpin(cx) { - Poll::Ready(Some((request, result))) => { - let Some(position) = self.remaining.iter().position(|req| { - req.as_ref().is_some_and(|req| { - req.offset() == request.offset - && req.len() == request.length - && req.alignment() == request.alignment - }) - }) else { - tracing::warn!(?request, "reader returned an unknown range"); - continue; - }; - let req = self.remaining[position] - .take() - .vortex_expect("matched request is present"); - return Poll::Ready(Some((req, result))); - } - Poll::Ready(None) => self.state = ReadRangeResultsState::Missing, - Poll::Pending => return Poll::Pending, - }, - ReadRangeResultsState::Missing => { - let Some(req) = self.remaining.iter_mut().find_map(Option::take) else { - return Poll::Ready(None); - }; - let error = vortex_err!( - "FileSegmentSource: read_ranges ended before resolving request. {:?}", - req - ); - return Poll::Ready(Some((req, Err(error)))); - } - } - } - } -} - -/// Drives request batches while keeping the reader's concurrency slots occupied. -struct ReadDriver { - reader: Arc, - batches: IoBatchStream, - pending: VecDeque, - reads: SelectAll, - num_active: usize, - batches_done: bool, - concurrency: usize, - metrics: RequestMetrics, -} - -impl ReadDriver { - fn new( - reader: R, - batches: BoxStream<'static, Vec>, - concurrency: usize, - metrics: RequestMetrics, - ) -> Self { - Self { - reader: Arc::new(reader), - batches: batches.fuse(), - pending: VecDeque::new(), - reads: SelectAll::new(), - num_active: 0, - batches_done: false, - concurrency, - metrics, - } - } - - fn submit_pending(&mut self) { - while self.num_active < self.concurrency && !self.pending.is_empty() { - let batch_len = (self.concurrency - self.num_active).min(self.pending.len()); - let reqs = self.pending.drain(..batch_len).collect::>(); - self.num_active += batch_len; - - self.metrics.read_ranges_calls.add(1); - self.metrics.read_ranges_num_ranges.update(batch_len as f64); - if batch_len > 1 { - self.metrics.read_ranges_multi.add(1); - } - tracing::trace!( - target: "vortex_file::read_ranges", - num_ranges = batch_len, - num_active = self.num_active, - "submitting positional read batch" - ); - - let requests = reqs - .iter() - .map(|req| ReadAtRequest::new(req.offset(), req.len(), req.alignment())) - .collect::>() - .into(); - let results = self.reader.read_ranges(requests); - self.reads.push(ReadRangeResults::new(results, reqs)); - } - } -} - -impl Stream for ReadDriver { - type Item = (); - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.as_mut().get_mut(); - - // Observe every batch already available so submission can fill all free slots at once. - if !this.batches_done { - loop { - match this.batches.poll_next_unpin(cx) { - Poll::Ready(Some(batch)) => this.pending.extend(batch), - Poll::Ready(None) => { - this.batches_done = true; - break; - } - Poll::Pending => break, - } - } - } - - this.submit_pending(); - - if this.batches_done && this.num_active == 0 { - return Poll::Ready(None); - } - - match this.reads.poll_next_unpin(cx) { - Poll::Ready(Some((req, result))) => { - this.num_active -= 1; - let result = validate_read_result(&req, result); - req.resolve(result); - Poll::Ready(Some(())) - } - Poll::Ready(None) if this.num_active == 0 => Poll::Pending, - Poll::Ready(None) => { - vortex_panic!("read result streams ended with active requests") - } - Poll::Pending => Poll::Pending, - } - } -} - pub struct FileSegmentSource { segments: Arc<[SegmentSpec]>, /// A queue for sending read request events to the I/O stream. @@ -285,6 +138,8 @@ pub struct FileSegmentSource { driver_panic: DriverPanic, /// The next read request ID. next_id: Arc, + /// Preferred size of canonical byte ranges for the underlying source. + preferred_read_size: Option, } impl FileSegmentSource { @@ -299,6 +154,7 @@ impl FileSegmentSource { metrics: RequestMetrics, ) -> Self { let (send, recv) = mpsc::unbounded(); + let preferred_read_size = reader.preferred_read_size(); let max_alignment = segments .iter() @@ -324,12 +180,130 @@ impl FileSegmentSource { StreamExt::boxed(recv), coalesce_config, max_alignment, - concurrency, + MAX_PARTIAL_SUBMISSION_REQUESTS, metrics.clone(), ) .boxed(); - let drive_fut = ReadDriver::new(reader, stream, concurrency, metrics).collect::<()>(); + let drive_fut = async move { + let mut batches = stream.fuse(); + let mut pending = VecDeque::::new(); + let mut reads = SelectAll::new(); + let mut num_active = 0usize; + let mut batches_done = false; + + loop { + if !batches_done { + loop { + match batches.next().now_or_never() { + Some(Some(batch)) => pending.extend(batch), + Some(None) => { + batches_done = true; + break; + } + None => break, + } + } + } + + while num_active < concurrency && !pending.is_empty() { + // A partial batch is submitted through one `read_ranges` stream. Do not + // refill individual slots from another partial batch as each range finishes: + // that turns a queued group into one syscall submission per completion. Let + // the current group drain, then submit all ready partial ranges together. + if num_active != 0 && pending.front().is_some_and(IoRequest::is_partial) { + break; + } + let batch_len = + if num_active == 0 && pending.front().is_some_and(IoRequest::is_partial) { + partial_submission_len(&pending) + } else { + (concurrency - num_active).min(pending.len()) + }; + let reqs = pending.drain(..batch_len).collect::>(); + num_active += batch_len; + + metrics.read_ranges_calls.add(1); + metrics.read_ranges_num_ranges.update(batch_len as f64); + if batch_len > 1 { + metrics.read_ranges_multi.add(1); + } + tracing::trace!( + target: "vortex_file::read_ranges", + num_ranges = batch_len, + num_active, + "submitting positional read batch" + ); + + let requests = reqs + .iter() + .map(|req| ReadAtRequest::new(req.offset(), req.len(), req.alignment())) + .collect::>() + .into(); + let mut remaining = reqs.into_iter().map(Some).collect::>(); + let mut results = reader.read_ranges(requests); + reads.push( + async_stream::stream! { + while let Some((request, result)) = results.next().await { + let Some(position) = remaining.iter().position(|req| { + req.as_ref().is_some_and(|req| { + req.offset() == request.offset + && req.len() == request.length + && req.alignment() == request.alignment + }) + }) else { + tracing::warn!(?request, "reader returned an unknown range"); + continue; + }; + let req = remaining[position] + .take() + .vortex_expect("matched request is present"); + yield (req, result); + } + for req in remaining.into_iter().flatten() { + let error = vortex_err!( + "FileSegmentSource: read_ranges ended before resolving request. {:?}", + req + ); + yield (req, Err(error)); + } + } + .boxed(), + ); + } + + if batches_done && num_active == 0 { + break; + } + if num_active == 0 { + match batches.next().await { + Some(batch) => pending.extend(batch), + None => batches_done = true, + } + continue; + } + + let next_read = reads.next(); + let next = if batches_done { + future::Either::Left((next_read.await, batches.next())) + } else { + future::select(next_read, batches.next()).await + }; + match next { + future::Either::Left((result, _)) => { + if let Some((req, result)) = result { + num_active -= 1; + let result = validate_read_result(&req, result); + req.resolve(result); + } + } + future::Either::Right((batch, _)) => match batch { + Some(batch) => pending.extend(batch), + None => batches_done = true, + }, + } + } + }; // Spawn the driver so the runtime makes I/O progress independently of any reader. Readers // join it (below) only to surface a panic raised while driving reads. @@ -356,57 +330,156 @@ impl FileSegmentSource { driver, driver_panic, next_id: Arc::new(AtomicUsize::new(0)), + preferred_read_size, } } } impl SegmentSource for FileSegmentSource { + fn preferred_read_size(&self) -> Option { + self.preferred_read_size + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.segments + .get(*id as usize) + .map(|spec| u64::from(spec.length)) + } + fn request(&self, id: SegmentId) -> SegmentFuture { - // We eagerly register the read request here assuming the behaviour of [`FileSegmentSource`], where - // coalescing becomes effective prior to the future being polled. - let spec = *match self.segments.get(*id as usize) { - Some(spec) => spec, - None => { - return future::ready(Err(vortex_err!("Missing segment: {}", id))).boxed(); + let Some(length) = self.segment_len(id) else { + return future::ready(Err(vortex_err!("Missing segment: {}", id))).boxed(); + }; + self.request_range_with_coalesce_distance(id, 0..length, None) + } + + fn request_range(&self, segment_id: SegmentId, range: Range) -> SegmentFuture { + self.request_range_with_coalesce_distance( + segment_id, + range, + self.preferred_read_size.map(|size| size / 4), + ) + } + + fn request_ranges(&self, segment_id: SegmentId, ranges: Vec>) -> Vec { + let coalesce_distance = self.preferred_read_size.map(|size| size / 4); + let mut registered = ranges + .into_iter() + .map(|range| self.register_range(segment_id, range, coalesce_distance)) + .collect::>(); + let poll_ids: Arc<[usize]> = registered + .iter() + .filter_map(|registration| registration.as_ref().ok().map(|read| read.id)) + .collect(); + let poll_once = Arc::new(AtomicBool::new(false)); + + registered + .drain(..) + .map(|registration| match registration { + Ok(read) => self.read_future(read, Arc::clone(&poll_ids), Arc::clone(&poll_once)), + Err(error) => future::ready(Err(error)).boxed(), + }) + .collect() + } +} + +impl FileSegmentSource { + fn request_range_with_coalesce_distance( + &self, + segment_id: SegmentId, + range: Range, + coalesce_distance: Option, + ) -> SegmentFuture { + match self.register_range(segment_id, range, coalesce_distance) { + Ok(read) => { + let poll_ids = Arc::from([read.id]); + self.read_future(read, poll_ids, Arc::new(AtomicBool::new(false))) } + Err(error) => future::ready(Err(error)).boxed(), + } + } + + fn register_range( + &self, + segment_id: SegmentId, + range: Range, + coalesce_distance: Option, + ) -> VortexResult { + // We eagerly register the read request here assuming the behaviour of + // [`FileSegmentSource`], where coalescing becomes effective prior to polling. + let spec = *match self.segments.get(*segment_id as usize) { + Some(spec) => spec, + None => return Err(vortex_err!("Missing segment: {}", segment_id)), }; + if range.start > range.end || range.end > u64::from(spec.length) { + return Err(vortex_err!( + "Segment {} range {}..{} is out of bounds for a {}-byte segment", + segment_id, + range.start, + range.end, + spec.length + )); + } + let SegmentSpec { - offset, - length, - alignment, + offset, alignment, .. } = spec; + let Some(offset) = offset.checked_add(range.start) else { + return Err(vortex_err!("Segment range offset overflow")); + }; + let Ok(length) = usize::try_from(range.end - range.start) else { + return Err(vortex_err!("Segment range length does not fit usize")); + }; + let (send, recv) = oneshot::channel(); let id = self.next_id.fetch_add(1, Ordering::Relaxed); let event = ReadEvent::Request(ReadRequest { id, offset, - length: length as usize, + length, alignment, + coalesce_distance, callback: send, }); - // If we fail to submit the event, we create a future that has failed. - if let Err(e) = self.events.unbounded_send(event) { - return future::ready(Err(vortex_err!("Failed to submit read request: {e}"))).boxed(); + if let Err(error) = self.events.unbounded_send(event) { + return Err(vortex_err!("Failed to submit read request: {error}")); } - let fut = ReadFuture { + Ok(RegisteredRead { id, recv: recv.into_future(), + }) + } + + fn read_future( + &self, + read: RegisteredRead, + poll_ids: Arc<[usize]>, + poll_once: Arc, + ) -> SegmentFuture { + ReadFuture { + id: read.id, + recv: read.recv, polled: false, finished: false, + poll_ids, + poll_once, events: self.events.clone(), driver: self.driver.clone(), driver_panic: Arc::clone(&self.driver_panic), - }; - - // One allocation: we only box the returned SegmentFuture, not the inner ReadFuture. - fut.boxed() + } + .boxed() } } +struct RegisteredRead { + id: usize, + recv: oneshot::AsyncReceiver>, +} + /// A future that resolves a read request from a [`FileSegmentSource`]. /// /// See the documentation for [`FileSegmentSource`] for details on coalescing and pre-fetching. @@ -416,6 +489,8 @@ struct ReadFuture { recv: oneshot::AsyncReceiver>, polled: bool, finished: bool, + poll_ids: Arc<[usize]>, + poll_once: Arc, events: mpsc::UnboundedSender, driver: SharedDriver, driver_panic: DriverPanic, @@ -450,11 +525,16 @@ impl Future for ReadFuture { }, Poll::Pending if !self.polled => { self.polled = true; - // Notify the I/O stream that this request has been polled. - match self.events.unbounded_send(ReadEvent::Polled(self.id)) { - Ok(()) => Poll::Pending, - Err(e) => Poll::Ready(Err(vortex_err!("ReadRequest dropped by runtime: {e}"))), + if !self.poll_once.swap(true, Ordering::AcqRel) { + for &id in self.poll_ids.iter() { + if let Err(error) = self.events.unbounded_send(ReadEvent::Polled(id)) { + return Poll::Ready(Err(vortex_err!( + "ReadRequest dropped by runtime: {error}" + ))); + } + } } + Poll::Pending } _ => Poll::Pending, } @@ -482,7 +562,7 @@ pub struct RequestMetrics { pub coalesced_requests: Counter, /// Distribution of how many segment requests were merged into each physical read. pub num_requests_coalesced: Histogram, - /// Number of calls made to [`VortexReadAt::read_ranges`]. + /// Number of calls made to [`VortexReadAt::read_ranges`](vortex_io::VortexReadAt::read_ranges). pub read_ranges_calls: Counter, /// Number of `read_ranges` calls containing more than one physical range. pub read_ranges_multi: Counter, @@ -533,7 +613,20 @@ impl BufferSegmentSource { } impl SegmentSource for BufferSegmentSource { + fn segment_len(&self, id: SegmentId) -> Option { + self.segments + .get(*id as usize) + .map(|spec| u64::from(spec.length)) + } + fn request(&self, id: SegmentId) -> SegmentFuture { + let Some(length) = self.segment_len(id) else { + return future::ready(Err(vortex_err!("Missing segment: {}", id))).boxed(); + }; + self.request_range(id, 0..length) + } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { let spec = match self.segments.get(*id as usize) { Some(spec) => spec, None => { @@ -541,8 +634,19 @@ impl SegmentSource for BufferSegmentSource { } }; - let start = spec.offset as usize; - let end = start + spec.length as usize; + if range.start > range.end || range.end > u64::from(spec.length) { + return future::ready(Err(vortex_err!( + "Segment {} range {}..{} out of bounds for segment length {}", + *id, + range.start, + range.end, + spec.length + ))) + .boxed(); + } + + let start = spec.offset as usize + range.start as usize; + let end = spec.offset as usize + range.end as usize; if end > self.buffer.len() { return future::ready(Err(vortex_err!( "Segment {} range {}..{} out of bounds for buffer of length {}", @@ -554,7 +658,11 @@ impl SegmentSource for BufferSegmentSource { .boxed(); } - let slice = self.buffer.slice(start..end).aligned(spec.alignment); + let slice = if range.start == 0 { + self.buffer.slice(start..end).aligned(spec.alignment) + } else { + self.buffer.slice(start..end) + }; future::ready(Ok(BufferHandle::new_host(slice))).boxed() } } @@ -571,35 +679,68 @@ mod tests { use super::*; - fn io_request(id: RequestId, offset: u64, length: usize) -> IoRequest { - let (callback, _receiver) = oneshot::channel(); - IoRequest::new_single(ReadRequest { - id, - offset, - length, - alignment: Alignment::none(), - callback, - }) + #[derive(Clone)] + struct MissingAndUnknownReadRanges; + + impl VortexReadAt for MissingAndUnknownReadRanges { + fn concurrency(&self) -> usize { + 2 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(8) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { panic!("read_at should not be called") }.boxed() + } + + fn read_ranges(&self, _requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { + let unknown = ReadAtRequest::new(8, 4, Alignment::none()); + let returned = ReadAtRequest::new(4, 4, Alignment::none()); + let buffer = BufferHandle::new_host(ByteBuffer::from(vec![0; 4])); + futures::stream::iter([(unknown, Ok(buffer.clone())), (returned, Ok(buffer))]).boxed() + } } #[tokio::test] - async fn read_range_results_matches_results_and_reports_missing_requests() { - let requests = vec![io_request(0, 0, 4), io_request(1, 4, 4)]; - let unknown = ReadAtRequest::new(8, 4, Alignment::none()); - let returned = ReadAtRequest::new(4, 4, Alignment::none()); - let buffer = BufferHandle::new_host(ByteBuffer::from(vec![0; 4])); - let results = - futures::stream::iter([(unknown, Ok(buffer.clone())), (returned, Ok(buffer))]).boxed(); - - let resolved = ReadRangeResults::new(results, requests) - .collect::>() - .await; - - assert_eq!(resolved.len(), 2); - assert_eq!(resolved[0].0.offset(), 4); - assert!(resolved[0].1.is_ok()); - assert_eq!(resolved[1].0.offset(), 0); - assert!(resolved[1].1.is_err()); + async fn read_driver_ignores_unknown_results_and_reports_missing_requests() { + let segments: Arc<[SegmentSpec]> = Arc::from([ + SegmentSpec { + offset: 0, + length: 4, + alignment: Alignment::none(), + }, + SegmentSpec { + offset: 4, + length: 4, + alignment: Alignment::none(), + }, + ]); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + MissingAndUnknownReadRanges, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + + let results = future::join_all([ + source.request(SegmentId::from(0)), + source.request(SegmentId::from(1)), + ]) + .await; + + assert!(results[0].is_err()); + match &results[1] { + Ok(buffer) => assert_eq!(buffer.len(), 4), + Err(error) => vortex_panic!("second request must resolve: {error}"), + } } #[derive(Clone)] @@ -726,6 +867,7 @@ mod tests { #[derive(Clone)] struct ReadRangesOnly { calls: Arc, + max_batch: Arc, } impl VortexReadAt for ReadRangesOnly { @@ -746,8 +888,9 @@ mod tests { async { panic!("read_at should not be called") }.boxed() } - fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { self.calls.fetch_add(1, Ordering::Relaxed); + self.max_batch.fetch_max(requests.len(), Ordering::Relaxed); let results = requests .iter() .copied() @@ -765,6 +908,7 @@ mod tests { #[tokio::test] async fn read_driver_batches_ready_requests() -> VortexResult<()> { let calls = Arc::new(AtomicUsize::new(0)); + let max_batch = Arc::new(AtomicUsize::new(0)); let segments: Arc<[SegmentSpec]> = (0..4) .map(|i| SegmentSpec { offset: i * 4, @@ -778,6 +922,7 @@ mod tests { segments, ReadRangesOnly { calls: Arc::clone(&calls), + max_batch: Arc::clone(&max_batch), }, TokioRuntime::current(), request_metrics.clone(), @@ -789,6 +934,7 @@ mod tests { assert_eq!(result?.len(), 4); } assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(max_batch.load(Ordering::Relaxed), 4); assert_eq!(request_metrics.read_ranges_calls.value(), 1); assert_eq!(request_metrics.read_ranges_multi.value(), 1); assert_eq!(request_metrics.read_ranges_num_ranges.count(), 1); @@ -796,6 +942,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn read_driver_submits_partial_ranges_together() -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let max_batch = Arc::new(AtomicUsize::new(0)); + let segments: Arc<[SegmentSpec]> = (0..6) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + ReadRangesOnly { + calls: Arc::clone(&calls), + max_batch: Arc::clone(&max_batch), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + + let results = source.request_ranges(SegmentId::from(0), vec![0..1, 1..2, 2..3, 3..4]); + for result in future::join_all(results).await { + assert_eq!(result?.len(), 1); + } + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(max_batch.load(Ordering::Relaxed), 4); + Ok(()) + } + #[derive(Clone)] struct ControlledReadRanges { active: Arc, @@ -822,7 +999,7 @@ mod tests { async { panic!("read_at should not be called") }.boxed() } - fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { self.batch_sizes.lock().push(requests.len()); let active = self.active.fetch_add(requests.len(), Ordering::SeqCst) + requests.len(); self.max_active.fetch_max(active, Ordering::SeqCst); diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index d6ed23f96d0..62f277289ea 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -45,6 +45,9 @@ vortex-utils = { workspace = true } [target.'cfg(unix)'.dependencies] custom-labels = { workspace = true } +[target.'cfg(target_os = "linux")'.dependencies] +io-uring = { workspace = true } + [target.'cfg(not(target_arch = "wasm32"))'.dependencies] # Smol is our default impl, so we don't want it to be optional, but it cannot be part of wasm smol = { workspace = true } @@ -61,6 +64,13 @@ rstest = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["full"] } +[target.'cfg(target_os = "linux")'.dev-dependencies] +rustix = { workspace = true } + +[[bench]] +name = "uring_read_at" +harness = false + [features] object_store = ["dep:object_store", "vortex-error/object_store"] tokio = ["tokio/fs", "tokio/rt-multi-thread"] diff --git a/vortex-io/benches/uring_read_at.rs b/vortex-io/benches/uring_read_at.rs new file mode 100644 index 00000000000..ea191c4e272 --- /dev/null +++ b/vortex-io/benches/uring_read_at.rs @@ -0,0 +1,671 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fixed-workload positional-read comparison. Run with `--help` for usage. + +#[cfg(not(target_os = "linux"))] +fn main() { + eprintln!("this benchmark is Linux-only"); +} + +#[cfg(target_os = "linux")] +mod bench { + use std::env; + use std::fs::File; + use std::hint::black_box; + use std::io; + use std::os::fd::AsRawFd; + use std::os::unix::fs::FileExt; + use std::path::Path; + use std::path::PathBuf; + use std::sync::Arc; + use std::sync::Barrier; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::sync::mpsc::Receiver; + use std::sync::mpsc::SyncSender; + use std::sync::mpsc::TryRecvError; + use std::sync::mpsc::sync_channel; + use std::thread; + use std::time::Duration; + use std::time::Instant; + + use io_uring::IoUring; + use io_uring::opcode; + use io_uring::types; + use parking_lot::Mutex; + use rustix::fs::Advice; + use rustix::fs::fadvise; + use vortex_utils::aliases::hash_map::HashMap; + + pub fn main() -> io::Result<()> { + let config = Config::parse()?; + if config.help { + help(); + return Ok(()); + } + let file = Arc::new(File::open(&config.path)?); + let file_len = file.metadata()?.len(); + let max_len = config.sizes.iter().copied().max().unwrap_or(0); + if max_len == 0 || file_len < max_len as u64 { + return Err(invalid( + "input file is smaller than the largest non-zero read", + )); + } + fadvise(&*file, 0, None, Advice::Random)?; + prepare_cache(&file, file_len, config.cache)?; + + let device_before = config.device.as_deref().map(device_stats).transpose()?; + let started = Instant::now(); + let mut result = run(Arc::clone(&file), file_len, &config)?; + let elapsed = started.elapsed(); + let device_after = config.device.as_deref().map(device_stats).transpose()?; + result.latencies.sort_unstable(); + let seconds = elapsed.as_secs_f64(); + + println!( + "mode={} engine_threads={} clients={} requests={} sizes={} cpu_ns={} cache={}", + config.mode, + config.engine_threads, + config.clients, + config.requests, + config + .sizes + .iter() + .map(usize::to_string) + .collect::>() + .join(","), + config.cpu_ns, + config.cache, + ); + println!( + "elapsed_s={seconds:.6} logical_reads={} kernel_read_ops={} submission_calls={} ops_per_submit={:.2} bytes={} throughput_mib_s={:.2} reads_s={:.0}", + result.latencies.len(), + result.kernel_ops, + result.submissions, + result.kernel_ops as f64 / result.submissions.max(1) as f64, + result.bytes, + result.bytes as f64 / 1_048_576.0 / seconds, + result.latencies.len() as f64 / seconds, + ); + println!( + "latency_us_p50={:.1} latency_us_p95={:.1} latency_us_p99={:.1} latency_us_max={:.1} checksum={}", + percentile(&result.latencies, 50).as_secs_f64() * 1e6, + percentile(&result.latencies, 95).as_secs_f64() * 1e6, + percentile(&result.latencies, 99).as_secs_f64() * 1e6, + result + .latencies + .last() + .copied() + .unwrap_or_default() + .as_secs_f64() + * 1e6, + result.checksum, + ); + if let Some((before, after)) = device_before.zip(device_after) { + println!( + "device_read_ios={} device_read_mib={:.2} device_read_ms={} device_inflight_end={}", + after.read_ios.saturating_sub(before.read_ios), + after.sectors.saturating_sub(before.sectors) as f64 * 512.0 / 1_048_576.0, + after.read_ms.saturating_sub(before.read_ms), + after.inflight, + ); + } + Ok(()) + } + + fn run(file: Arc, file_len: u64, config: &Config) -> io::Result { + let engine = match config.mode { + Mode::Inline => None, + Mode::Pool => Some(Arc::new(Engine::new( + Arc::clone(&file), + EngineKind::Pread, + config.engine_threads, + config.queue_depth, + )?)), + Mode::Uring => Some(Arc::new(Engine::new( + Arc::clone(&file), + EngineKind::Uring, + config.engine_threads, + config.queue_depth, + )?)), + }; + let next = Arc::new(AtomicUsize::new(0)); + let barrier = Arc::new(Barrier::new(config.clients + 1)); + let mut joins = Vec::with_capacity(config.clients); + for client_id in 0..config.clients { + let file = Arc::clone(&file); + let engine = engine.as_ref().map(Arc::clone); + let next = Arc::clone(&next); + let barrier = Arc::clone(&barrier); + let sizes = Arc::clone(&config.sizes); + let request_count = config.requests; + let cpu_ns = config.cpu_ns; + joins.push(thread::spawn(move || -> io::Result { + let mut row = ClientRow::default(); + barrier.wait(); + loop { + let request_id = next.fetch_add(1, Ordering::Relaxed); + if request_id >= request_count { + break; + } + let len = sizes[request_id % sizes.len()]; + let offset = (random_at(request_id as u64) + % ((file_len - len as u64) / 4096 + 1)) + * 4096; + let started = Instant::now(); + let buffer = match &engine { + Some(engine) => engine.read(offset, len, client_id)?, + None => { + let mut buffer = vec![0; len]; + file.read_exact_at(&mut buffer, offset)?; + buffer + } + }; + row.latencies.push(started.elapsed()); + row.bytes += len as u64; + row.checksum = row.checksum.wrapping_add(sample(&buffer)); + busy_cpu(cpu_ns, row.checksum); + } + Ok(row) + })); + } + barrier.wait(); + let mut result = ResultRow::default(); + for join in joins { + let row = join + .join() + .map_err(|_| io::Error::other("client panicked"))??; + result.latencies.extend(row.latencies); + result.bytes += row.bytes; + result.checksum = result.checksum.wrapping_add(row.checksum); + } + result.kernel_ops = match engine { + Some(engine) => { + let stats = engine.shutdown()?; + result.submissions = stats.submissions; + stats.operations + } + None => { + result.submissions = config.requests as u64; + config.requests as u64 + } + }; + Ok(result) + } + + struct Engine { + senders: Vec>, + joins: Mutex>>>, + } + + #[derive(Clone, Copy)] + enum EngineKind { + Pread, + Uring, + } + + impl Engine { + fn new(file: Arc, kind: EngineKind, n: usize, depth: usize) -> io::Result { + if n == 0 { + return Err(invalid("engine-threads must be non-zero")); + } + let mut senders = Vec::with_capacity(n); + let mut joins = Vec::with_capacity(n); + for id in 0..n { + let (tx, rx) = sync_channel(depth); + let file = Arc::clone(&file); + joins.push( + thread::Builder::new() + .name(format!("read-engine-{id}")) + .spawn(move || match kind { + EngineKind::Pread => pread_worker(file, rx), + EngineKind::Uring => uring_worker(file, rx, depth), + })?, + ); + senders.push(tx); + } + Ok(Self { + senders, + joins: Mutex::new(joins), + }) + } + + fn read(&self, offset: u64, len: usize, shard: usize) -> io::Result> { + let (complete, receive) = sync_channel(1); + self.senders[shard % self.senders.len()] + .send(Message::Read(Request { + offset, + buffer: vec![0; len], + filled: 0, + complete, + })) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "engine stopped"))?; + receive + .recv() + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "completion dropped"))? + } + + fn shutdown(&self) -> io::Result { + for sender in &self.senders { + sender + .send(Message::Stop) + .map_err(|_| io::Error::other("engine stopped"))?; + } + let mut stats = WorkerStats::default(); + for join in self.joins.lock().drain(..) { + let worker = join + .join() + .map_err(|_| io::Error::other("engine panicked"))??; + stats.operations += worker.operations; + stats.submissions += worker.submissions; + } + Ok(stats) + } + } + + enum Message { + Read(Request), + Stop, + } + + struct Request { + offset: u64, + buffer: Vec, + filled: usize, + complete: SyncSender>>, + } + + fn pread_worker(file: Arc, rx: Receiver) -> io::Result { + let mut operations = 0; + while let Ok(message) = rx.recv() { + match message { + Message::Read(mut request) => { + let result = file + .read_exact_at(&mut request.buffer, request.offset) + .map(|()| request.buffer); + operations += 1; + drop(request.complete.send(result)); + } + Message::Stop => break, + } + } + Ok(WorkerStats { + operations, + submissions: operations, + }) + } + + fn uring_worker( + file: Arc, + rx: Receiver, + depth: usize, + ) -> io::Result { + let entries = u32::try_from(depth.next_power_of_two()).map_err(io::Error::other)?; + let mut ring: IoUring = IoUring::builder() + .setup_single_issuer() + .setup_defer_taskrun() + .build(entries)?; + let mut pending: HashMap = HashMap::with_capacity(depth); + let mut next_id = 1_u64; + let mut operations = 0; + let mut submissions = 0; + let mut stopping = false; + loop { + let completions = ring + .completion() + .map(|cqe| (cqe.user_data(), cqe.result())) + .collect::>(); + for (user_data, completion_result) in completions { + let Some(mut request) = pending.remove(&user_data) else { + return Err(io::Error::other("unknown completion")); + }; + operations += 1; + match completion_result { + result if result < 0 => drop( + request + .complete + .send(Err(io::Error::from_raw_os_error(-result))), + ), + 0 => drop(request.complete.send(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "io_uring read reached EOF", + )))), + result => { + request.filled += result as usize; + if request.filled == request.buffer.len() { + drop(request.complete.send(Ok(request.buffer))); + } else { + push(&mut ring, &file, request, &mut pending, &mut next_id)?; + } + } + } + } + if stopping && pending.is_empty() { + break; + } + let mut accepted = 0; + while pending.len() < depth { + let message = if pending.is_empty() && accepted == 0 { + rx.recv() + .map_err(|_| io::Error::other("request queue disconnected"))? + } else { + match rx.try_recv() { + Ok(message) => message, + Err(TryRecvError::Empty) => break, + Err(TryRecvError::Disconnected) => { + stopping = true; + break; + } + } + }; + match message { + Message::Read(request) => { + push(&mut ring, &file, request, &mut pending, &mut next_id)?; + accepted += 1; + } + Message::Stop => { + stopping = true; + break; + } + } + } + if !pending.is_empty() { + ring.submit_and_wait(1)?; + submissions += 1; + } + } + Ok(WorkerStats { + operations, + submissions, + }) + } + + fn push( + ring: &mut IoUring, + file: &File, + request: Request, + pending: &mut HashMap, + next_id: &mut u64, + ) -> io::Result<()> { + let id = *next_id; + *next_id = next_id.wrapping_add(1); + let remaining = request.buffer.len() - request.filled; + let len = u32::try_from(remaining.min(u32::MAX as usize)).map_err(io::Error::other)?; + let pointer = unsafe { request.buffer.as_ptr().add(request.filled).cast_mut() }; + let entry = opcode::Read::new(types::Fd(file.as_raw_fd()), pointer, len) + .offset(request.offset + request.filled as u64) + .build() + .user_data(id); + // SAFETY: `pending` owns the stable allocation until this operation's CQE is reaped. + unsafe { + ring.submission() + .push(&entry) + .map_err(|_| io::Error::new(io::ErrorKind::WouldBlock, "SQ full"))?; + } + pending.insert(id, request); + Ok(()) + } + + fn prepare_cache(file: &File, file_len: u64, mode: Cache) -> io::Result<()> { + match mode { + Cache::Keep => Ok(()), + Cache::Cold => { + file.sync_all()?; + fadvise(file, 0, None, Advice::DontNeed)?; + Ok(()) + } + Cache::Warm => { + let mut buffer = vec![0; 1024 * 1024]; + let mut offset = 0; + while offset < file_len { + let len = usize::try_from((file_len - offset).min(buffer.len() as u64)) + .map_err(io::Error::other)?; + file.read_exact_at(&mut buffer[..len], offset)?; + offset += len as u64; + } + black_box(sample(&buffer)); + Ok(()) + } + } + } + + fn busy_cpu(ns: u64, seed: u64) { + if ns == 0 { + return; + } + let start = Instant::now(); + let duration = Duration::from_nanos(ns); + let mut value = seed; + while start.elapsed() < duration { + for _ in 0..64 { + value = value.wrapping_mul(0x9e37_79b9_7f4a_7c15).rotate_left(17) + ^ 0xe703_7ed1_a0b4_28db; + } + } + black_box(value); + } + + fn sample(buffer: &[u8]) -> u64 { + u64::from(buffer[0]) + ^ (u64::from(buffer[buffer.len() / 2]) << 8) + ^ (u64::from(buffer[buffer.len() - 1]) << 16) + } + + fn percentile(values: &[Duration], p: usize) -> Duration { + values + .get((values.len().saturating_sub(1)) * p / 100) + .copied() + .unwrap_or_default() + } + + #[derive(Default)] + struct ResultRow { + latencies: Vec, + bytes: u64, + checksum: u64, + kernel_ops: u64, + submissions: u64, + } + + #[derive(Default)] + struct WorkerStats { + operations: u64, + submissions: u64, + } + + #[derive(Default)] + struct ClientRow { + latencies: Vec, + bytes: u64, + checksum: u64, + } + + fn random_at(index: u64) -> u64 { + let mut z = index.wrapping_add(0x9e37_79b9_7f4a_7c15); + z = (z ^ z >> 30).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ z >> 27).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ z >> 31 + } + + #[derive(Clone, Copy)] + enum Mode { + Inline, + Pool, + Uring, + } + impl std::fmt::Display for Mode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + match self { + Self::Inline => "pread-inline", + Self::Pool => "pread-pool", + Self::Uring => "uring", + } + ) + } + } + #[derive(Clone, Copy)] + enum Cache { + Keep, + Cold, + Warm, + } + impl std::fmt::Display for Cache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + match self { + Self::Keep => "keep", + Self::Cold => "cold", + Self::Warm => "warm", + } + ) + } + } + + struct Config { + path: PathBuf, + mode: Mode, + clients: usize, + engine_threads: usize, + queue_depth: usize, + requests: usize, + sizes: Arc<[usize]>, + cpu_ns: u64, + cache: Cache, + device: Option, + help: bool, + } + + impl Config { + fn parse() -> io::Result { + let mut c = Self { + path: PathBuf::new(), + mode: Mode::Uring, + clients: 32, + engine_threads: 1, + queue_depth: 256, + requests: 10_000, + sizes: Arc::from([64 * 1024]), + cpu_ns: 0, + cache: Cache::Keep, + device: None, + help: false, + }; + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + let value = args.next(); + match arg.as_str() { + "--help" | "-h" => c.help = true, + "--path" => c.path = value.ok_or_else(|| invalid("missing path"))?.into(), + "--mode" => { + c.mode = match value.as_deref() { + Some("pread-inline") => Mode::Inline, + Some("pread-pool") => Mode::Pool, + Some("uring") => Mode::Uring, + _ => return Err(invalid("bad mode")), + } + } + "--clients" => c.clients = number(value, &arg)?, + "--engine-threads" => c.engine_threads = number(value, &arg)?, + "--queue-depth" => c.queue_depth = number(value, &arg)?, + "--requests" => c.requests = number(value, &arg)?, + "--cpu-ns" => c.cpu_ns = number(value, &arg)?, + "--sizes" => { + c.sizes = value + .ok_or_else(|| invalid("missing sizes"))? + .split(',') + .map(size) + .collect::>>()? + .into() + } + "--cache" => { + c.cache = match value.as_deref() { + Some("keep") => Cache::Keep, + Some("cold") => Cache::Cold, + Some("warm") => Cache::Warm, + _ => return Err(invalid("bad cache")), + } + } + "--device" => c.device = value, + _ => return Err(invalid(format!("unknown argument {arg}"))), + } + } + if !c.help && c.path.as_os_str().is_empty() { + return Err(invalid("--path is required")); + } + if c.clients == 0 + || c.engine_threads == 0 + || c.queue_depth == 0 + || c.requests == 0 + || c.sizes.is_empty() + { + return Err(invalid("counts and sizes must be non-zero")); + } + Ok(c) + } + } + + fn number(value: Option, name: &str) -> io::Result { + value + .ok_or_else(|| invalid(format!("missing {name}")))? + .parse() + .map_err(|_| invalid(format!("bad {name}"))) + } + fn size(value: &str) -> io::Result { + let (n, multiplier) = match value.as_bytes().last() { + Some(b'K' | b'k') => (&value[..value.len() - 1], 1024), + Some(b'M' | b'm') => (&value[..value.len() - 1], 1024 * 1024), + _ => (value, 1), + }; + n.parse::() + .ok() + .and_then(|n| n.checked_mul(multiplier)) + .filter(|n| *n > 0) + .ok_or_else(|| invalid(format!("bad size {value}"))) + } + fn invalid(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message.into()) + } + fn help() { + println!( + "usage: uring_read_at --path FILE [--mode pread-inline|pread-pool|uring] [--clients N] [--engine-threads N] [--queue-depth N] [--requests N] [--sizes 4K,64K,1M] [--cpu-ns N] [--cache keep|cold|warm] [--device nvme0n1]" + ); + } + + #[derive(Default)] + struct DeviceStats { + read_ios: u64, + sectors: u64, + read_ms: u64, + inflight: u64, + } + fn device_stats(device: &str) -> io::Result { + let values = + std::fs::read_to_string(Path::new("/sys/class/block").join(device).join("stat"))? + .split_whitespace() + .map(|v| v.parse::().map_err(io::Error::other)) + .collect::>>()?; + if values.len() < 9 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "short block stat", + )); + } + Ok(DeviceStats { + read_ios: values[0], + sectors: values[2], + read_ms: values[3], + inflight: values[8], + }) + } +} + +#[cfg(target_os = "linux")] +fn main() -> std::io::Result<()> { + bench::main() +} diff --git a/vortex-io/src/compat/read_at.rs b/vortex-io/src/compat/read_at.rs index 3d9cc93b1a6..8cc722d24aa 100644 --- a/vortex-io/src/compat/read_at.rs +++ b/vortex-io/src/compat/read_at.rs @@ -27,6 +27,10 @@ impl VortexReadAt for Compat { self.inner().coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.inner().preferred_read_size() + } + fn concurrency(&self) -> usize { self.inner().concurrency() } diff --git a/vortex-io/src/object_store/filesystem.rs b/vortex-io/src/object_store/filesystem.rs index ca68f5f7efa..57b980975fd 100644 --- a/vortex-io/src/object_store/filesystem.rs +++ b/vortex-io/src/object_store/filesystem.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Formatter; +use std::path::PathBuf; use std::sync::Arc; use async_trait::async_trait; @@ -21,6 +22,8 @@ use crate::filesystem::FileListing; use crate::filesystem::FileSystem; use crate::object_store::ObjectStoreReadAt; use crate::runtime::Handle; +#[cfg(not(target_arch = "wasm32"))] +use crate::std_file::FileReadAt; /// A [`FileSystem`] backed by an [`ObjectStore`]. // TODO(ngates): we could consider spawning a driver task inside this file system such that we can @@ -28,6 +31,7 @@ use crate::runtime::Handle; pub struct ObjectStoreFileSystem { store: Arc, handle: Handle, + local_root: Option, } impl Debug for ObjectStoreFileSystem { @@ -41,16 +45,21 @@ impl Debug for ObjectStoreFileSystem { impl ObjectStoreFileSystem { /// Create a new filesystem backed by the given object store and runtime handle. pub fn new(store: Arc, handle: Handle) -> Self { - Self { store, handle } + Self { + store, + handle, + local_root: None, + } } /// Create a new filesystem backed by a local file system object store and the given runtime /// handle. pub fn local(handle: Handle) -> Self { - Self::new( - Arc::new(object_store::local::LocalFileSystem::new()), + Self { + store: Arc::new(object_store::local::LocalFileSystem::new()), handle, - ) + local_root: Some(PathBuf::from("/")), + } } } @@ -105,6 +114,13 @@ impl FileSystem for ObjectStoreFileSystem { } async fn open_read(&self, path: &str) -> VortexResult> { + #[cfg(not(target_arch = "wasm32"))] + if let Some(root) = &self.local_root { + return Ok(Arc::new(FileReadAt::open( + root.join(path), + self.handle.clone(), + )?)); + } Ok(Arc::new(ObjectStoreReadAt::new( Arc::clone(&self.store), to_object_path(path), @@ -153,6 +169,23 @@ mod tests { Ok(ObjectStoreFileSystem::new(store, handle)) } + #[cfg(not(target_arch = "wasm32"))] + #[tokio::test] + async fn local_files_use_file_read_settings() -> VortexResult<()> { + let file = tempfile::NamedTempFile::new()?; + let handle = Handle::find().expect("tokio runtime available within #[tokio::test]"); + let reader = ObjectStoreFileSystem::local(handle) + .open_read(file.path().to_string_lossy().as_ref()) + .await?; + + assert_eq!( + reader.coalesce_config().expect("local coalescing").distance, + 0 + ); + assert_eq!(reader.concurrency(), crate::std_file::DEFAULT_CONCURRENCY); + Ok(()) + } + /// Regression test for #6599: globbing an exact path that exists must return that one file. /// `ObjectStore::list` never yields the prefix itself, so this would return nothing if the /// exact-path branch used `list`. diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 462bc498f82..494bc650d44 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -25,6 +25,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::CoalesceConfig; +use crate::OBJECT_STORAGE_PREFERRED_READ_SIZE; use crate::ReadAtRequest; use crate::ReadAtStream; use crate::VortexReadAt; @@ -44,6 +45,7 @@ pub struct ObjectStoreReadAt { allocator: HostAllocatorRef, concurrency: usize, coalesce_config: Option, + preferred_read_size: Option, } impl ObjectStoreReadAt { @@ -68,6 +70,7 @@ impl ObjectStoreReadAt { allocator, concurrency: DEFAULT_CONCURRENCY, coalesce_config: Some(CoalesceConfig::object_storage()), + preferred_read_size: Some(OBJECT_STORAGE_PREFERRED_READ_SIZE), } } @@ -82,6 +85,12 @@ impl ObjectStoreReadAt { self.coalesce_config = Some(config); self } + + /// Set the preferred size of independently requested byte ranges for this source. + pub fn with_preferred_read_size(mut self, preferred_read_size: u64) -> Self { + self.preferred_read_size = Some(preferred_read_size); + self + } } async fn read_object_store_range( @@ -162,6 +171,10 @@ impl VortexReadAt for ObjectStoreReadAt { self.coalesce_config } + fn preferred_read_size(&self) -> Option { + self.preferred_read_size + } + fn concurrency(&self) -> usize { self.concurrency } diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index 82190bdd699..3ce0eb7d91d 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.rs @@ -22,6 +22,12 @@ use vortex_metrics::MetricBuilder; use vortex_metrics::MetricsRegistry; use vortex_metrics::Timer; +/// Preferred read size for local file sources, including SSDs. +pub const FILE_PREFERRED_READ_SIZE: u64 = 64 * 1024; + +/// Preferred read size for object storage sources. +pub const OBJECT_STORAGE_PREFERRED_READ_SIZE: u64 = 1 << 20; + /// Configuration for coalescing nearby I/O requests into single operations. #[derive(Clone, Copy, Debug)] pub struct CoalesceConfig { @@ -69,7 +75,10 @@ impl CoalesceConfig { /// Configuration appropriate for local filesystem access. pub const fn file() -> Self { - Self::new(1 << 20, 4 << 20) // 1MB distance, 4MB max + // Local random reads are cheap enough that reading gaps between segments costs more than + // issuing another operation. Adjacent and overlapping requests still coalesce, while the + // 4 MiB cap preserves useful batching for scans. + Self::new(0, 4 << 20) } /// Configuration appropriate for object storage (S3, GCS, etc.). @@ -93,6 +102,14 @@ pub trait VortexReadAt: Send + Sync + 'static { None } + /// Preferred size of independently requested byte ranges for this source. + /// + /// Layout readers can use this hint when dividing large logical segments into canonical read + /// ranges. Returning `None` asks readers to preserve whole-segment reads. + fn preferred_read_size(&self) -> Option { + None + } + /// Maximum number of concurrent I/O requests for that should be pulled from this source. /// /// This value is used to control how many [`VortexReadAt::read_at`] calls can @@ -148,6 +165,10 @@ impl VortexReadAt for Arc { self.as_ref().coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.as_ref().preferred_read_size() + } + fn concurrency(&self) -> usize { self.as_ref().concurrency() } @@ -179,6 +200,10 @@ impl VortexReadAt for Arc { self.as_ref().coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.as_ref().preferred_read_size() + } + fn concurrency(&self) -> usize { self.as_ref().concurrency() } @@ -344,6 +369,10 @@ impl VortexReadAt for InstrumentedReadAt { self.read.coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.read.preferred_read_size() + } + fn concurrency(&self) -> usize { self.read.concurrency() } @@ -439,7 +468,7 @@ mod tests { #[test] fn test_coalesce_config_file() { let config = CoalesceConfig::file(); - assert_eq!(config.distance, 1 << 20); // 1MB + assert_eq!(config.distance, 0); assert_eq!(config.max_size, 4 << 20); // 4MB } diff --git a/vortex-io/src/std_file/mod.rs b/vortex-io/src/std_file/mod.rs index c30248ed496..7023cfc5886 100644 --- a/vortex-io/src/std_file/mod.rs +++ b/vortex-io/src/std_file/mod.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod read_at; +#[cfg(target_os = "linux")] +mod uring; mod write; pub use read_at::*; diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 3d59a595f70..fb124dc19de 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -13,9 +13,16 @@ use std::os::unix::fs::FileExt; use std::os::windows::fs::FileExt; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; use futures::FutureExt; +use futures::StreamExt; +use futures::channel::mpsc; use futures::future::BoxFuture; +use futures::stream; use vortex_array::buffer::BufferHandle; use vortex_array::memory::DefaultHostAllocator; use vortex_array::memory::HostAllocatorRef; @@ -23,8 +30,12 @@ use vortex_buffer::Alignment; use vortex_error::VortexResult; use crate::CoalesceConfig; +use crate::FILE_PREFERRED_READ_SIZE; +use crate::ReadAtRequest; +use crate::ReadAtStream; use crate::VortexReadAt; use crate::runtime::Handle; +use crate::runtime::Task; /// Read exactly `buffer.len()` bytes from `file` starting at `offset`. /// This is a platform-specific helper that uses the most efficient method available. @@ -61,6 +72,54 @@ pub fn read_exact_at(file: &File, buffer: &mut [u8], offset: u64) -> io::Result< /// Default number of concurrent requests to allow for local file I/O. pub const DEFAULT_CONCURRENCY: usize = 32; +/// Local reads that complete this quickly are normally page-cache hits. Keeping the worker set +/// small avoids making allocation and task scheduling more expensive than the reads themselves. +const HOT_READ_CONCURRENCY: usize = 2; +const SMALL_READ_MAX_LENGTH: usize = 32 * 1024; +const COLD_READ_THRESHOLD: Duration = Duration::from_micros(250); + +type RangeReadResponse = (ReadAtRequest, VortexResult, Duration); + +fn spawn_range_workers( + handle: &Handle, + file: &Arc, + allocator: &HostAllocatorRef, + requests: &Arc<[ReadAtRequest]>, + next: &Arc, + send: &mpsc::UnboundedSender, + worker_count: usize, +) -> Vec> { + (0..worker_count) + .map(|_| { + let file = Arc::clone(file); + let allocator = Arc::clone(allocator); + let requests = Arc::clone(requests); + let next = Arc::clone(next); + let send = send.clone(); + handle.spawn_blocking(move || { + loop { + let index = next.fetch_add(1, Ordering::Relaxed); + let Some(request) = requests.get(index).copied() else { + break; + }; + let started = Instant::now(); + let result = (|| -> VortexResult { + let mut buffer = allocator.allocate(request.length, request.alignment)?; + read_exact_at(&file, buffer.as_mut_slice(), request.offset)?; + Ok(BufferHandle::new_host(buffer.freeze())) + })(); + if send + .unbounded_send((request, result, started.elapsed())) + .is_err() + { + break; + } + } + }) + }) + .collect() +} + /// An adapter type wrapping a [`File`] to implement [`VortexReadAt`]. pub struct FileReadAt { uri: Arc, @@ -102,6 +161,10 @@ impl VortexReadAt for FileReadAt { Some(CoalesceConfig::file()) } + fn preferred_read_size(&self) -> Option { + Some(FILE_PREFERRED_READ_SIZE) + } + fn concurrency(&self) -> usize { DEFAULT_CONCURRENCY } @@ -125,6 +188,19 @@ impl VortexReadAt for FileReadAt { let handle = self.handle.clone(); let allocator = Arc::clone(&self.allocator); async move { + #[cfg(target_os = "linux")] + if let Some(submission) = super::uring::try_admit(length) { + let buffer = allocator.allocate(length, alignment)?; + if buffer.is_empty() { + return Ok(BufferHandle::new_host(buffer.freeze())); + } + let receive = submission.read_at(Arc::clone(&file), offset, buffer); + let buffer = receive.into_future().await.map_err(|_| { + io::Error::new(io::ErrorKind::BrokenPipe, "io_uring completion dropped") + })??; + return Ok(BufferHandle::new_host(buffer.freeze())); + } + handle .spawn_blocking(move || { let mut buffer = allocator.allocate(length, alignment)?; @@ -135,4 +211,114 @@ impl VortexReadAt for FileReadAt { } .boxed() } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + if requests.is_empty() { + return stream::empty().boxed(); + } + + #[cfg(target_os = "linux")] + if let Some(reads) = super::uring::try_read_ranges( + Arc::clone(&self.file), + Arc::clone(&self.allocator), + Arc::clone(&requests), + ) { + return reads; + } + + let total_length = requests.iter().fold(0usize, |total, request| { + total.saturating_add(request.length) + }); + let average_length = total_length / requests.len(); + if average_length > SMALL_READ_MAX_LENGTH { + let worker_count = requests.len().min(DEFAULT_CONCURRENCY); + let next = Arc::new(AtomicUsize::new(0)); + let (send, recv) = mpsc::unbounded(); + let mut workers = Vec::with_capacity(worker_count); + + for _ in 0..worker_count { + let file = Arc::clone(&self.file); + let allocator = Arc::clone(&self.allocator); + let requests = Arc::clone(&requests); + let next = Arc::clone(&next); + let send = send.clone(); + workers.push(self.handle.spawn_blocking(move || { + loop { + let index = next.fetch_add(1, Ordering::Relaxed); + let Some(request) = requests.get(index).copied() else { + break; + }; + let result = (|| -> VortexResult { + let mut buffer = + allocator.allocate(request.length, request.alignment)?; + read_exact_at(&file, buffer.as_mut_slice(), request.offset)?; + Ok(BufferHandle::new_host(buffer.freeze())) + })(); + if send.unbounded_send((request, result)).is_err() { + break; + } + } + })); + } + drop(send); + + return stream::unfold((recv, workers), |(mut recv, workers)| async move { + recv.next() + .await + .map(|response| (response, (recv, workers))) + }) + .boxed(); + } + + let initial_worker_count = requests.len().min(HOT_READ_CONCURRENCY); + let next = Arc::new(AtomicUsize::new(0)); + let (send, recv) = mpsc::unbounded(); + let workers = spawn_range_workers( + &self.handle, + &self.file, + &self.allocator, + &requests, + &next, + &send, + initial_worker_count, + ); + let handle = self.handle.clone(); + let file = Arc::clone(&self.file); + let allocator = Arc::clone(&self.allocator); + let request_count = requests.len(); + + // Retaining task handles in the stream state aborts workers that have not started their + // next range if the consumer drops the response stream. + async_stream::stream! { + let mut recv = recv; + let mut workers = workers; + let mut received = 0usize; + while let Some((request, result, read_latency)) = recv.next().await { + received += 1; + if received == 1 + && initial_worker_count < DEFAULT_CONCURRENCY + && read_latency >= COLD_READ_THRESHOLD + { + let additional_workers = request_count + .saturating_sub(initial_worker_count) + .min(DEFAULT_CONCURRENCY - initial_worker_count); + workers.extend(spawn_range_workers( + &handle, + &file, + &allocator, + &requests, + &next, + &send, + additional_workers, + )); + } + yield (request, result); + if received == request_count { + break; + } + } + drop(workers); + } + .boxed() + } } diff --git a/vortex-io/src/std_file/uring.rs b/vortex-io/src/std_file/uring.rs new file mode 100644 index 00000000000..2efdf3d55e7 --- /dev/null +++ b/vortex-io/src/std_file/uring.rs @@ -0,0 +1,518 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Process-wide `io_uring` engine for local positional reads. + +use std::collections::VecDeque; +use std::env; +use std::fs::File; +use std::io; +use std::os::fd::AsRawFd; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use std::sync::mpsc::Receiver; +use std::thread; + +use futures::FutureExt; +use futures::StreamExt; +use futures::future; +use futures::future::BoxFuture; +use futures::stream; +use io_uring::IoUring; +use io_uring::opcode; +use io_uring::types; +use vortex_array::buffer::BufferHandle; +use vortex_array::memory::HostAllocatorRef; +use vortex_array::memory::WritableHostBuffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::parallelism::get_available_parallelism; + +use crate::ReadAtRequest; +use crate::ReadAtStream; + +const DEFAULT_QUEUE_DEPTH: usize = 256; +const DEFAULT_MIN_READ_SIZE: usize = 1024 * 1024; +const MAX_RINGS: usize = 4; +const MAX_BATCH_REQUESTS: usize = 512; +const MAX_BATCH_BYTES: usize = 16 << 20; + +type Completion = oneshot::Sender>; + +static ENGINE: OnceLock>> = OnceLock::new(); + +fn engine() -> Option<&'static Arc> { + ENGINE + .get_or_init(|| match UringEngine::from_environment() { + Ok(engine) => engine.map(Arc::new), + Err(error) => { + tracing::debug!(%error, "io_uring unavailable; using blocking positional reads"); + None + } + }) + .as_ref() +} + +/// Submit a positional read to the shared engine. +/// +/// `None` means that io_uring is disabled or unavailable and the caller should use its portable +/// blocking-I/O path. Setting `VORTEX_IO_URING=1` enables the engine. The ring and queue counts can +/// be overridden for benchmarking with `VORTEX_IO_URING_RINGS` and +/// `VORTEX_IO_URING_QUEUE_DEPTH`; `VORTEX_IO_URING_MAX_IN_FLIGHT` controls when excess requests +/// spill back to the blocking-I/O path and `VORTEX_IO_URING_MIN_READ_SIZE` controls the minimum +/// request size. +pub(super) fn try_admit(length: usize) -> Option { + let engine = engine()?; + let admission = engine.try_admit(length)?; + + Some(Submission { + engine: Arc::clone(engine), + admission, + }) +} + +/// Submit a complete small-range batch to the ring before waiting for any completion. +/// +/// The batch limits match the file segment driver's partial-read limits, bounding both the +/// submission queue and the memory allocated before I/O begins. +pub(super) fn try_read_ranges( + file: Arc, + allocator: HostAllocatorRef, + requests: Arc<[ReadAtRequest]>, +) -> Option { + let engine = engine()?; + let total_bytes = requests + .iter() + .try_fold(0usize, |sum, request| sum.checked_add(request.length))?; + if requests.len() > MAX_BATCH_REQUESTS + || total_bytes > MAX_BATCH_BYTES + || requests + .iter() + .any(|request| request.length < engine.min_read_size) + { + return None; + } + let mut admissions = engine.try_admit_count(requests.len())?.into_iter(); + + let mut batches = (0..engine.senders.len()) + .map(|_| Vec::new()) + .collect::>(); + let mut responses: Vec)>> = + Vec::with_capacity(requests.len()); + let first_sender = engine.next.fetch_add(1, Ordering::Relaxed) % engine.senders.len(); + + for (index, request) in requests.iter().copied().enumerate() { + let admission = admissions + .next() + .vortex_expect("one admission reserved per range"); + let buffer = match allocator.allocate(request.length, request.alignment) { + Ok(buffer) => buffer, + Err(error) => { + responses.push(future::ready((request, Err(error))).boxed()); + continue; + } + }; + if buffer.is_empty() { + responses.push( + future::ready((request, Ok(BufferHandle::new_host(buffer.freeze())))).boxed(), + ); + continue; + } + + let (complete, receive) = oneshot::channel(); + let sender_index = (first_sender + index) % batches.len(); + batches[sender_index].push(Request { + file: Arc::clone(&file), + offset: request.offset, + buffer, + filled: 0, + complete, + _admission: Some(admission), + }); + responses.push( + async move { + let result = match receive.into_future().await { + Ok(Ok(buffer)) => Ok(BufferHandle::new_host(buffer.freeze())), + Ok(Err(error)) => Err(error.into()), + Err(_) => Err(vortex_err!("io_uring completion dropped")), + }; + (request, result) + } + .boxed(), + ); + } + + for (sender, batch) in engine.senders.iter().zip(batches) { + if batch.is_empty() { + continue; + } + if let Err(error) = sender.send(batch) { + for request in error.0 { + drop(request.complete.send(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "io_uring worker stopped", + )))); + } + } + } + + Some( + stream::iter(responses) + .buffer_unordered(requests.len().max(1)) + .boxed(), + ) +} + +pub(super) struct Submission { + engine: Arc, + admission: Admission, +} + +impl Submission { + pub(super) fn read_at( + self, + file: Arc, + offset: u64, + buffer: WritableHostBuffer, + ) -> oneshot::Receiver> { + let (complete, receive) = oneshot::channel(); + let request = Request { + file, + offset, + buffer, + filled: 0, + complete, + _admission: Some(self.admission), + }; + let sender_index = + self.engine.next.fetch_add(1, Ordering::Relaxed) % self.engine.senders.len(); + if let Err(error) = self.engine.senders[sender_index].send(vec![request]) { + for request in error.0 { + drop(request.complete.send(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "io_uring worker stopped", + )))); + } + } + receive + } +} + +struct UringEngine { + senders: Vec>>, + next: AtomicUsize, + in_flight: Arc, + max_in_flight: usize, + min_read_size: usize, +} + +impl UringEngine { + fn from_environment() -> io::Result> { + if !env::var("VORTEX_IO_URING") + .is_ok_and(|value| matches!(value.as_str(), "1" | "true" | "on" | "yes")) + { + return Ok(None); + } + + let available = get_available_parallelism().unwrap_or(1); + // One owner per four available CPUs retained the batching advantage without making the + // owner thread a page-cache bottleneck. Storage-bound workloads naturally need fewer. + let default_rings = available.div_ceil(4).clamp(1, MAX_RINGS); + let rings = read_env_usize("VORTEX_IO_URING_RINGS", default_rings)?.clamp(1, 64); + let depth = read_env_usize("VORTEX_IO_URING_QUEUE_DEPTH", DEFAULT_QUEUE_DEPTH)? + .clamp(8, 32_768) + .next_power_of_two(); + let max_in_flight = + read_env_usize("VORTEX_IO_URING_MAX_IN_FLIGHT", rings)?.clamp(1, rings * depth); + let min_read_size = read_env_usize("VORTEX_IO_URING_MIN_READ_SIZE", DEFAULT_MIN_READ_SIZE)?; + + let mut senders = Vec::with_capacity(rings); + for id in 0..rings { + let (send, receive) = mpsc::channel(); + let (ready_send, ready_receive) = mpsc::sync_channel(1); + thread::Builder::new() + .name(format!("vortex-io-uring-{id}")) + .spawn(move || match new_ring(depth) { + Ok(ring) => { + drop(ready_send.send(Ok(()))); + worker(ring, receive, depth); + } + Err(error) => { + let startup_error = io::Error::new(error.kind(), error.to_string()); + drop(ready_send.send(Err(startup_error))); + } + })?; + // SINGLE_ISSUER and DEFER_TASKRUN bind the ring to its owner task, so ring setup must + // happen inside the owner thread. This handshake still detects setup failure before + // publishing the engine. + ready_receive.recv().map_err(|_| { + io::Error::new(io::ErrorKind::BrokenPipe, "io_uring worker failed to start") + })??; + senders.push(send); + } + tracing::debug!( + rings, + depth, + max_in_flight, + min_read_size, + "started local-file io_uring engine" + ); + Ok(Some(Self { + senders, + next: AtomicUsize::new(0), + in_flight: Arc::new(AtomicUsize::new(0)), + max_in_flight, + min_read_size, + })) + } + + fn try_admit(&self, length: usize) -> Option { + if length < self.min_read_size { + return None; + } + self.try_admit_count(1)?.pop() + } + + fn try_admit_count(&self, count: usize) -> Option> { + self.in_flight + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + current + .checked_add(count) + .filter(|&next| next <= self.max_in_flight) + }) + .ok()?; + Some( + (0..count) + .map(|_| Admission(Arc::clone(&self.in_flight))) + .collect(), + ) + } +} + +struct Admission(Arc); + +impl Drop for Admission { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} + +fn read_env_usize(name: &str, default: usize) -> io::Result { + match env::var(name) { + Ok(value) => value.parse().map_err(|error| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("invalid {name}={value:?}: {error}"), + ) + }), + Err(env::VarError::NotPresent) => Ok(default), + Err(error) => Err(io::Error::new(io::ErrorKind::InvalidInput, error)), + } +} + +fn new_ring(depth: usize) -> io::Result { + let entries = u32::try_from(depth).map_err(io::Error::other)?; + IoUring::builder() + .setup_single_issuer() + .setup_defer_taskrun() + .build(entries) + .or_else(|_| IoUring::new(entries)) +} + +struct Request { + file: Arc, + offset: u64, + buffer: WritableHostBuffer, + filled: usize, + complete: Completion, + _admission: Option, +} + +fn worker(ring: IoUring, receive: Receiver>, depth: usize) { + let result = run_worker(ring, &receive, depth); + if let Err(error) = result { + tracing::warn!(%error, "local-file io_uring worker stopped"); + } +} + +fn run_worker(ring: IoUring, receive: &Receiver>, depth: usize) -> io::Result<()> { + let mut pending: HashMap = HashMap::with_capacity(depth); + let mut queued = VecDeque::new(); + // Declared after `pending` so the ring is closed (and the kernel has released all requests) + // before any in-flight buffers are dropped on an error return. + let mut ring = ring; + let mut next_id = 1_u64; + + loop { + let completions = ring + .completion() + .map(|cqe| (cqe.user_data(), cqe.result())) + .collect::>(); + for (id, result) in completions { + let Some(mut request) = pending.remove(&id) else { + return Err(io::Error::other("io_uring returned an unknown completion")); + }; + match result { + result if result < 0 => { + drop( + request + .complete + .send(Err(io::Error::from_raw_os_error(-result))), + ); + } + 0 => { + drop(request.complete.send(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "io_uring read reached EOF", + )))); + } + result => { + request.filled += result as usize; + if request.filled == request.buffer.len() { + drop(request.complete.send(Ok(request.buffer))); + } else { + push(&mut ring, request, &mut pending, &mut next_id)?; + } + } + } + } + + if pending.is_empty() && queued.is_empty() { + match receive.recv() { + Ok(batch) => queued.extend(batch), + Err(_) => return Ok(()), + } + } + while let Ok(batch) = receive.try_recv() { + queued.extend(batch); + } + while pending.len() < depth { + let Some(request) = queued.pop_front() else { + break; + }; + push(&mut ring, request, &mut pending, &mut next_id)?; + } + + if !pending.is_empty() { + ring.submit_and_wait(1)?; + } + } +} + +fn push( + ring: &mut IoUring, + mut request: Request, + pending: &mut HashMap, + next_id: &mut u64, +) -> io::Result<()> { + let id = *next_id; + *next_id = next_id.wrapping_add(1); + let remaining = request.buffer.len() - request.filled; + let length = u32::try_from(remaining.min(u32::MAX as usize)).map_err(io::Error::other)?; + let pointer = request.buffer.as_mut_slice()[request.filled..].as_mut_ptr(); + let entry = opcode::Read::new(types::Fd(request.file.as_raw_fd()), pointer, length) + .offset(request.offset + request.filled as u64) + .build() + .user_data(id); + // SAFETY: `pending` retains the file and stable buffer allocation until the CQE is reaped. + unsafe { + ring.submission() + .push(&entry) + .map_err(|_| io::Error::new(io::ErrorKind::WouldBlock, "io_uring SQ is full"))?; + } + pending.insert(id, request); + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use vortex_array::memory::DefaultHostAllocator; + use vortex_array::memory::HostAllocator; + use vortex_buffer::Alignment; + + use super::*; + + #[test] + fn reads_into_owned_host_buffer() -> anyhow::Result<()> { + let mut file = tempfile::tempfile()?; + file.write_all(b"abcdefgh")?; + let file = Arc::new(file); + let (send, receive) = mpsc::channel(); + let owner = thread::spawn(move || { + let ring = new_ring(8)?; + run_worker(ring, &receive, 8) + }); + + let buffer = DefaultHostAllocator.allocate(4, Alignment::none())?; + let (complete, completed) = oneshot::channel(); + send.send(vec![Request { + file, + offset: 2, + buffer, + filled: 0, + complete, + _admission: Some(Admission(Arc::new(AtomicUsize::new(1)))), + }]) + .map_err(|_| anyhow::anyhow!("io_uring request channel closed"))?; + + let buffer = futures::executor::block_on(completed.into_future()) + .map_err(|_| anyhow::anyhow!("io_uring completion channel closed"))??; + assert_eq!(buffer.freeze().as_slice(), b"cdef"); + drop(send); + owner + .join() + .map_err(|_| anyhow::anyhow!("io_uring owner panicked"))??; + Ok(()) + } + + #[test] + fn reads_batch_larger_than_queue_depth() -> anyhow::Result<()> { + let mut file = tempfile::tempfile()?; + let data = (0_u8..32).collect::>(); + file.write_all(&data)?; + let file = Arc::new(file); + let (send, receive) = mpsc::channel(); + let owner = thread::spawn(move || { + let ring = new_ring(8)?; + run_worker(ring, &receive, 8) + }); + + let in_flight = Arc::new(AtomicUsize::new(16)); + let mut requests = Vec::new(); + let mut completions = Vec::new(); + for offset in (0_u64..32).step_by(2) { + let buffer = DefaultHostAllocator.allocate(2, Alignment::none())?; + let (complete, completed) = oneshot::channel(); + requests.push(Request { + file: Arc::clone(&file), + offset, + buffer, + filled: 0, + complete, + _admission: Some(Admission(Arc::clone(&in_flight))), + }); + completions.push((offset as usize, completed)); + } + send.send(requests) + .map_err(|_| anyhow::anyhow!("io_uring worker stopped"))?; + + for (offset, completed) in completions { + let buffer = futures::executor::block_on(completed.into_future()) + .map_err(|_| anyhow::anyhow!("io_uring completion channel closed"))??; + assert_eq!(buffer.freeze().as_slice(), &data[offset..offset + 2]); + } + assert_eq!(in_flight.load(Ordering::Relaxed), 0); + drop(send); + owner + .join() + .map_err(|_| anyhow::anyhow!("io_uring owner panicked"))??; + Ok(()) + } +} diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..63a69c78bf0 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -40,11 +40,13 @@ termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } vortex-array = { workspace = true } +vortex-alp = { workspace = true } vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["layout"] } +vortex-fastlanes = { workspace = true } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-metrics = { workspace = true } diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index c3743a1df45..0fdd28e0b00 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -346,7 +346,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 #[test] fn test_display_tree_with_segment_source() { if std::env::var("NEXTEST_RUN_ID").is_ok() { - temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", None::<&str>, || { + temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("0"), || { block_on(|handle| async move { let session = new_session().with_handle(handle); let ctx = ArrayContext::empty(); diff --git a/vortex-layout/src/layouts/flat/mod.rs b/vortex-layout/src/layouts/flat/mod.rs index 6913d2fc0c4..eb2eb6a81ca 100644 --- a/vortex-layout/src/layouts/flat/mod.rs +++ b/vortex-layout/src/layouts/flat/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod partial; mod reader; pub mod writer; @@ -34,7 +35,7 @@ use crate::segments::SegmentSource; /// Check if inline array node is enabled. pub(super) fn flat_layout_inline_array_node() -> bool { static FLAT_LAYOUT_INLINE_ARRAY_NODE: LazyLock = - LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").is_ok_and(|v| v == "1")); + LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").map_or(true, |v| v != "0")); *FLAT_LAYOUT_INLINE_ARRAY_NODE } diff --git a/vortex-layout/src/layouts/flat/partial.rs b/vortex-layout/src/layouts/flat/partial.rs new file mode 100644 index 00000000000..1dc6d977782 --- /dev/null +++ b/vortex-layout/src/layouts/flat/partial.rs @@ -0,0 +1,1065 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::BTreeSet; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::try_join_all; +use prost::Message; +use vortex_alp::ALPRD; +use vortex_alp::ALPRDMetadata; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Struct; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::patches::Patches; +use vortex_array::patches::PatchesMetadata; +use vortex_array::serde::SerializedArray; +use vortex_array::serde::SerializedBuffer; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedMetadata; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::layouts::flat::FlatLayout; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +#[derive(Clone)] +pub(super) struct PartialReadPlan { + array_tree: ByteBuffer, + bytes_per_row: usize, + row_granularity: usize, + kind: PartialReadKind, +} + +#[derive(Clone)] +enum PartialReadKind { + Fixed(Arc<[PlannedBuffer]>), + Alprd(Box), +} + +#[derive(Clone)] +struct PlannedBuffer { + descriptor: SerializedBuffer, + bytes_per_row: usize, + row_granularity: usize, + bytes_per_granule: usize, +} + +#[derive(Clone)] +struct ALPRDReadPlan { + serialized: SerializedArray, + descriptors: Arc<[SerializedBuffer]>, + left: BitPackedReadPlan, + right: BitPackedReadPlan, + patch_buffers: Arc<[SerializedBuffer]>, + patch_metadata: PatchesMetadata, + patch_indices_dtype: DType, + left_parts_dtype: DType, + left_parts_dictionary: Buffer, + right_bit_width: u8, + element_dtype: DType, + list_size: u32, + row_count: usize, +} + +struct PageResolveContext<'a> { + dtype: &'a DType, + row_range: &'a Range, + mask: &'a Mask, + ctx: &'a ReadContext, + session: &'a VortexSession, +} + +#[derive(Clone)] +struct BitPackedReadPlan { + descriptor: SerializedBuffer, + ptype: vortex_array::dtype::PType, + bit_width: u8, + offset: u16, +} + +pub(super) struct RegisteredPartialRead { + array_tree: ByteBuffer, + kind: RegisteredReadKind, +} + +enum RegisteredReadKind { + Fixed { + pages: Vec, + }, + Alprd { + pages: Vec, + patch_buffers: Vec<(SegmentFuture, SerializedBuffer)>, + plan: ALPRDReadPlan, + }, +} + +struct RegisteredALPRDPage { + rows: Range, + left: SegmentFuture, + right: SegmentFuture, +} + +struct RegisteredPage { + rows: Range, + buffers: Vec<(SegmentFuture, SerializedBuffer)>, +} + +impl PartialReadPlan { + pub(super) fn supports_mask(mask: &Mask) -> bool { + !mask.all_true() + } + + pub(super) fn try_new(layout: &FlatLayout) -> VortexResult> { + let Some(array_tree) = layout.array_tree().cloned() else { + return Ok(None); + }; + let serialized = SerializedArray::from_array_tree(array_tree.clone())?; + let descriptors: Arc<[SerializedBuffer]> = serialized.buffer_descriptors()?.into(); + let row_count = usize::try_from(layout.row_count())?; + + if let Some((plan, bytes_per_row)) = try_alprd_plan( + &serialized, + layout.dtype(), + layout.array_ctx(), + row_count, + Arc::clone(&descriptors), + )? { + return Ok(Some(Self { + array_tree, + bytes_per_row, + row_granularity: 1, + kind: PartialReadKind::Alprd(Box::new(plan)), + })); + } + + let mut planned = Vec::new(); + if !collect_raw_buffers( + &serialized, + layout.dtype(), + layout.array_ctx(), + 1, + row_count, + &descriptors, + &mut planned, + )? { + return Ok(None); + } + planned.sort_unstable_by_key(|buffer| buffer.descriptor.index()); + if planned.len() != descriptors.len() + || planned + .iter() + .enumerate() + .any(|(index, buffer)| buffer.descriptor.index() != index) + { + return Ok(None); + } + for buffer in &planned { + let expected = row_count + .div_ceil(buffer.row_granularity) + .checked_mul(buffer.bytes_per_granule) + .ok_or_else(|| vortex_err!("Partial buffer length overflow"))?; + if buffer.descriptor.range().len() != expected { + return Ok(None); + } + } + let bytes_per_row = planned.iter().try_fold(0usize, |sum, buffer| { + sum.checked_add(buffer.bytes_per_row) + .ok_or_else(|| vortex_err!("Partial row width overflow")) + })?; + if bytes_per_row == 0 { + return Ok(None); + } + let row_granularity = planned + .iter() + .map(|buffer| buffer.row_granularity) + .try_fold(1usize, checked_lcm)?; + Ok(Some(Self { + array_tree, + bytes_per_row, + row_granularity, + kind: PartialReadKind::Fixed(planned.into()), + })) + } + + pub(super) fn register( + &self, + source: &Arc, + segment_id: SegmentId, + layout_len: usize, + row_range: &Range, + mask: &Mask, + ) -> Option { + if !Self::supports_mask(mask) { + return None; + } + let preferred_read_size = usize::try_from(source.preferred_read_size()?).ok()?; + let segment_len = usize::try_from(source.segment_len(segment_id)?).ok()?; + let desired_rows = (preferred_read_size / self.bytes_per_row).max(1); + let page_rows = desired_rows + .div_ceil(self.row_granularity) + .saturating_mul(self.row_granularity); + let pages = selected_pages(page_rows, layout_len, row_range, mask)?; + let pages = match &self.kind { + PartialReadKind::Alprd(_) => selected_page_runs(&pages, row_range, mask)?, + PartialReadKind::Fixed(_) => pages, + }; + let (partial_bytes, request_count) = self.estimated_partial_io(&pages, layout_len)?; + let partial_cost = partial_bytes.checked_add( + request_count + .saturating_sub(1) + .checked_mul(preferred_read_size)?, + )?; + if partial_cost >= segment_len { + tracing::trace!( + layout_len, + page_rows, + page_count = pages.len(), + partial_bytes, + request_count, + partial_cost, + segment_len, + "Flat partial read rejected by I/O cost" + ); + return None; + } + tracing::trace!( + layout_len, + page_rows, + page_count = pages.len(), + partial_bytes, + request_count, + partial_cost, + segment_len, + "Flat partial read registered" + ); + + let kind = match &self.kind { + PartialReadKind::Fixed(buffers) => { + let page_specs = pages + .into_iter() + .map(|rows| { + let ranges = buffers + .iter() + .map(|buffer| { + let start = buffer.descriptor.range().start + + (rows.start / buffer.row_granularity) + * buffer.bytes_per_granule; + let end = buffer.descriptor.range().start + + rows.end.div_ceil(buffer.row_granularity) + * buffer.bytes_per_granule; + Some(( + u64::try_from(start).ok()?..u64::try_from(end).ok()?, + buffer.descriptor.clone(), + )) + }) + .collect::>>()?; + Some((rows, ranges)) + }) + .collect::>>()?; + let requests = source.request_ranges( + segment_id, + page_specs + .iter() + .flat_map(|(_, ranges)| ranges.iter().map(|(range, _)| range.clone())) + .collect(), + ); + let mut requests = requests.into_iter(); + let pages = page_specs + .into_iter() + .map(|(rows, ranges)| { + let buffers = ranges + .into_iter() + .map(|(_, descriptor)| Some((requests.next()?, descriptor))) + .collect::>>()?; + Some(RegisteredPage { rows, buffers }) + }) + .collect::>>()?; + RegisteredReadKind::Fixed { pages } + } + PartialReadKind::Alprd(plan) => { + let values_per_row = usize::try_from(plan.list_size).ok()?; + let page_specs = pages + .into_iter() + .map(|rows| { + let inner_start = rows.start.checked_mul(values_per_row)?; + let inner_end = rows.end.checked_mul(values_per_row)?; + let left = bitpacked_range(&plan.left, inner_start..inner_end)?; + let right = bitpacked_range(&plan.right, inner_start..inner_end)?; + Some(( + rows, + u64::try_from(left.start).ok()?..u64::try_from(left.end).ok()?, + u64::try_from(right.start).ok()?..u64::try_from(right.end).ok()?, + )) + }) + .collect::>>()?; + let patch_specs = plan + .patch_buffers + .iter() + .map(|descriptor| { + Some(( + u64::try_from(descriptor.range().start).ok()? + ..u64::try_from(descriptor.range().end).ok()?, + descriptor.clone(), + )) + }) + .collect::>>()?; + let ranges = page_specs + .iter() + .flat_map(|(_, left, right)| [left.clone(), right.clone()]) + .chain(patch_specs.iter().map(|(range, _)| range.clone())) + .collect(); + let mut requests = source.request_ranges(segment_id, ranges).into_iter(); + let pages = page_specs + .into_iter() + .map(|(rows, ..)| { + Some(RegisteredALPRDPage { + rows, + left: requests.next()?, + right: requests.next()?, + }) + }) + .collect::>>()?; + let patch_buffers = patch_specs + .into_iter() + .map(|(_, descriptor)| Some((requests.next()?, descriptor))) + .collect::>>()?; + RegisteredReadKind::Alprd { + pages, + patch_buffers, + plan: plan.as_ref().clone(), + } + } + }; + + Some(RegisteredPartialRead { + array_tree: self.array_tree.clone(), + kind, + }) + } + + fn estimated_partial_io( + &self, + pages: &[Range], + _layout_len: usize, + ) -> Option<(usize, usize)> { + match &self.kind { + PartialReadKind::Fixed(buffers) => { + let bytes = pages.iter().try_fold(0usize, |total, rows| { + buffers.iter().try_fold(total, |total, buffer| { + let granules = rows + .end + .div_ceil(buffer.row_granularity) + .checked_sub(rows.start / buffer.row_granularity)?; + total.checked_add(granules.checked_mul(buffer.bytes_per_granule)?) + }) + })?; + Some((bytes, pages.len().checked_mul(buffers.len())?)) + } + PartialReadKind::Alprd(plan) => { + let values_per_row = usize::try_from(plan.list_size).ok()?; + let page_bytes = pages.iter().try_fold(0usize, |total, rows| { + let values = rows.start.checked_mul(values_per_row)? + ..rows.end.checked_mul(values_per_row)?; + let left = bitpacked_range(&plan.left, values.clone())?; + let right = bitpacked_range(&plan.right, values)?; + total.checked_add(left.len())?.checked_add(right.len()) + })?; + let patch_bytes = plan + .patch_buffers + .iter() + .try_fold(0usize, |total, buffer| { + total.checked_add(buffer.range().len()) + })?; + Some(( + page_bytes.checked_add(patch_bytes)?, + pages + .len() + .checked_mul(2)? + .checked_add(plan.patch_buffers.len())?, + )) + } + } + } +} + +impl RegisteredPartialRead { + pub(super) async fn resolve( + self, + dtype: &DType, + row_range: &Range, + mask: &Mask, + ctx: &ReadContext, + session: &VortexSession, + ) -> VortexResult { + let chunks = match self.kind { + RegisteredReadKind::Fixed { pages } => { + resolve_fixed_pages( + self.array_tree, + pages, + PageResolveContext { + dtype, + row_range, + mask, + ctx, + session, + }, + ) + .await? + } + RegisteredReadKind::Alprd { + pages, + patch_buffers, + plan, + } => { + resolve_alprd_pages( + pages, + patch_buffers, + plan, + dtype, + row_range, + mask, + ctx, + session, + ) + .await? + } + }; + finish_chunks(chunks, dtype, session) + } +} + +async fn resolve_fixed_pages( + array_tree: ByteBuffer, + pages: Vec, + context: PageResolveContext<'_>, +) -> VortexResult> { + let mut page_futures = Vec::new(); + for page in pages { + let local_mask = page_mask(&page.rows, context.row_range, context.mask.indices())?; + if local_mask.all_false() { + continue; + } + let array_tree = array_tree.clone(); + let dtype = context.dtype.clone(); + let ctx = context.ctx.clone(); + let session = context.session.clone(); + page_futures.push(async move { + let buffers = try_join_all(page.buffers.into_iter().map( + |(future, descriptor)| async move { + future.await?.ensure_aligned(descriptor.alignment()) + }, + )) + .await?; + let array = SerializedArray::from_flatbuffer_with_buffers(array_tree, buffers)? + .decode(&dtype, page.rows.len(), &ctx, &session)?; + clear_stats(&array); + apply_page_mask(array, local_mask) + }); + } + try_join_all(page_futures).await +} + +#[allow(clippy::too_many_arguments)] +async fn resolve_alprd_pages( + pages: Vec, + patch_requests: Vec<(SegmentFuture, SerializedBuffer)>, + plan: ALPRDReadPlan, + dtype: &DType, + row_range: &Range, + mask: &Mask, + ctx: &ReadContext, + session: &VortexSession, +) -> VortexResult> { + let patch_handles = try_join_all(patch_requests.into_iter().map( + |(future, descriptor)| async move { + Ok::<_, vortex_error::VortexError>(( + descriptor.index(), + future.await?.ensure_aligned(descriptor.alignment())?, + )) + }, + )); + let page_handles = try_join_all(pages.into_iter().filter_map(|page| { + let local_mask = match page_mask(&page.rows, row_range, mask.indices()) { + Ok(local_mask) if !local_mask.all_false() => local_mask, + Ok(_) => return None, + Err(error) => return Some(futures::future::ready(Err(error)).left_future()), + }; + let left_alignment = plan.left.descriptor.alignment(); + let right_alignment = plan.right.descriptor.alignment(); + Some( + async move { + let (left, right) = futures::try_join!(page.left, page.right)?; + Ok::<_, vortex_error::VortexError>(( + page.rows, + local_mask, + left.ensure_aligned(left_alignment)?, + right.ensure_aligned(right_alignment)?, + )) + } + .right_future(), + ) + })); + + // Every range for this Flat layout is registered before resolution. Poll the complete set + // together so the driver can issue and coalesce patch, left-part, and right-part reads in one + // I/O round; array reconstruction starts only after that set has resolved. + let (patch_handles, page_handles) = futures::try_join!(patch_handles, page_handles)?; + let mut handles = empty_handles(plan.descriptors.len()); + for (index, handle) in patch_handles { + handles[index] = handle; + } + let serialized = plan.serialized.with_buffers(handles); + let alprd = serialized.child(0); + let patch_len = plan.patch_metadata.len()?; + let patch_indices = + alprd + .child(2) + .decode(&plan.patch_indices_dtype, patch_len, ctx, session)?; + let patch_indices = patch_indices + .execute::(&mut session.create_execution_ctx())? + .into_array(); + let patch_values = alprd.child(3).decode( + &plan.left_parts_dtype.as_nonnullable(), + patch_len, + ctx, + session, + )?; + let full_inner_len = usize::try_from(plan.list_size)? + .checked_mul(plan.row_count) + .ok_or_else(|| vortex_err!("ALPRD inner length overflow"))?; + let full_patches = Patches::new( + full_inner_len, + plan.patch_metadata.offset()?, + patch_indices, + patch_values, + None, + )?; + + page_handles + .into_iter() + .map(|(rows, local_mask, left, right)| { + let inner_start = rows.start * plan.list_size as usize; + let inner_end = rows.end * plan.list_size as usize; + let inner_len = inner_end - inner_start; + let left = BitPacked::try_new( + left, + plan.left.ptype, + Validity::from(plan.left_parts_dtype.nullability()), + None, + plan.left.bit_width, + inner_len, + 0, + )? + .into_array(); + let right = BitPacked::try_new( + right, + plan.right.ptype, + Validity::NonNullable, + None, + plan.right.bit_width, + inner_len, + 0, + )? + .into_array(); + let patches = full_patches.slice(inner_start..inner_end)?; + let elements = ALPRD::try_new( + plan.element_dtype.clone(), + left, + plan.left_parts_dictionary.clone(), + right, + plan.right_bit_width, + patches, + )? + .into_array(); + let array = FixedSizeListArray::try_new( + elements, + plan.list_size, + Validity::from(dtype.nullability()), + rows.len(), + )? + .into_array(); + apply_page_mask(array, local_mask) + }) + .collect() +} + +fn finish_chunks( + mut chunks: Vec, + dtype: &DType, + session: &VortexSession, +) -> VortexResult { + match chunks.len() { + 0 => Ok(Canonical::empty(dtype).into_array()), + 1 => Ok(chunks.remove(0)), + _ => { + let chunks = ChunkedArray::try_new(chunks, dtype.clone())?.into_array(); + let mut ctx = session.create_execution_ctx(); + Ok(chunks.execute::(&mut ctx)?.into_array()) + } + } +} + +fn apply_page_mask(array: ArrayRef, mask: Mask) -> VortexResult { + if mask.all_true() { + Ok(array) + } else if let AllOr::Some([(start, end)]) = mask.slices() { + array.slice(*start..*end) + } else { + array.filter(mask) + } +} + +fn clear_stats(array: &ArrayRef) { + for child in array.depth_first_traversal() { + child.statistics().clear_all(); + } +} + +fn empty_handles(len: usize) -> Vec { + (0..len) + .map(|_| BufferHandle::new_host(ByteBuffer::empty())) + .collect() +} + +fn selected_pages( + page_rows: usize, + layout_len: usize, + row_range: &Range, + mask: &Mask, +) -> Option>> { + let mut page_indices = BTreeSet::new(); + match mask.slices() { + AllOr::All => return None, + AllOr::None => {} + AllOr::Some(slices) => { + for &(start, end) in slices { + if start >= end { + continue; + } + let global_start = row_range.start.checked_add(start)?; + let global_end = row_range.start.checked_add(end)?; + if global_end > row_range.end || global_end > layout_len { + return None; + } + page_indices.extend(global_start / page_rows..=(global_end - 1) / page_rows); + } + } + } + Some( + page_indices + .into_iter() + .map(|page_index| { + let start = page_index * page_rows; + start..start.saturating_add(page_rows).min(layout_len) + }) + .collect(), + ) +} + +fn selected_page_runs( + pages: &[Range], + row_range: &Range, + mask: &Mask, +) -> Option>> { + let selected = mask.slices(); + let mut runs = Vec::new(); + for page in pages { + match selected { + AllOr::All => { + let start = page.start.max(row_range.start); + let end = page.end.min(row_range.end); + if start < end { + runs.push(start..end); + } + } + AllOr::None => {} + AllOr::Some(slices) => { + for &(start, end) in slices { + let global_start = row_range.start.checked_add(start)?; + let global_end = row_range.start.checked_add(end)?; + let run_start = page.start.max(global_start); + let run_end = page.end.min(global_end); + if run_start < run_end { + runs.push(run_start..run_end); + } + } + } + } + } + Some(runs) +} + +fn page_mask( + page_rows: &Range, + row_range: &Range, + selected: AllOr<&[usize]>, +) -> VortexResult { + match selected { + AllOr::None => Ok(Mask::new_false(page_rows.len())), + AllOr::All => { + let start = page_rows.start.max(row_range.start); + let end = page_rows.end.min(row_range.end); + Ok(Mask::from_indices( + page_rows.len(), + (start..end).map(|row| row - page_rows.start), + )) + } + AllOr::Some(indices) => Ok(Mask::from_indices( + page_rows.len(), + indices.iter().filter_map(|&index| { + let row = row_range.start.checked_add(index)?; + page_rows.contains(&row).then(|| row - page_rows.start) + }), + )), + } +} + +fn try_alprd_plan( + node: &SerializedArray, + dtype: &DType, + ctx: &ReadContext, + row_count: usize, + descriptors: Arc<[SerializedBuffer]>, +) -> VortexResult> { + if ctx.resolve(node.encoding_id()) != Some(FixedSizeList.id()) + || node.nbuffers() != 0 + || node.nchildren() != 1 + { + return Ok(None); + } + let DType::FixedSizeList(element_dtype, list_size, _) = dtype else { + return Ok(None); + }; + let list_size_usize = usize::try_from(*list_size)?; + if list_size_usize == 0 || !list_size_usize.is_multiple_of(1024) { + return Ok(None); + } + let DType::Primitive(element_ptype, element_nullability) = element_dtype.as_ref() else { + return Ok(None); + }; + if !matches!( + element_ptype, + vortex_array::dtype::PType::F32 | vortex_array::dtype::PType::F64 + ) { + return Ok(None); + } + + let alprd = node.child(0); + if ctx + .resolve(alprd.encoding_id()) + .is_none_or(|id| id.as_str() != "vortex.alprd") + || alprd.nbuffers() != 0 + || alprd.nchildren() != 4 + { + return Ok(None); + } + let metadata = ALPRDMetadata::decode(alprd.metadata())?; + let Some(patch_metadata) = metadata.patches().copied() else { + return Ok(None); + }; + let left_parts_dtype = DType::Primitive(metadata.left_parts_ptype(), *element_nullability); + let right_ptype = match element_ptype { + vortex_array::dtype::PType::F32 => vortex_array::dtype::PType::U32, + vortex_array::dtype::PType::F64 => vortex_array::dtype::PType::U64, + _ => unreachable!(), + }; + let inner_len = row_count + .checked_mul(list_size_usize) + .ok_or_else(|| vortex_err!("ALPRD inner length overflow"))?; + let left = try_bitpacked_plan( + &alprd.child(0), + left_parts_dtype.as_ptype(), + ctx, + inner_len, + &descriptors, + )?; + let right = try_bitpacked_plan(&alprd.child(1), right_ptype, ctx, inner_len, &descriptors)?; + let (Some(left), Some(right)) = (left, right) else { + return Ok(None); + }; + + let mut patch_indices = BTreeSet::new(); + collect_buffer_indices(&alprd.child(2), &mut patch_indices); + collect_buffer_indices(&alprd.child(3), &mut patch_indices); + let Some(patch_buffers) = patch_indices + .into_iter() + .map(|index| descriptors.get(index).cloned()) + .collect::>>() + else { + return Ok(None); + }; + if patch_buffers.is_empty() { + return Ok(None); + } + let expected_indices: BTreeSet<_> = [left.descriptor.index(), right.descriptor.index()] + .into_iter() + .chain(patch_buffers.iter().map(SerializedBuffer::index)) + .collect(); + if expected_indices.len() != descriptors.len() + || expected_indices.iter().copied().ne(0..descriptors.len()) + { + return Ok(None); + } + + let blocks_per_row = list_size_usize / 1024; + let bytes_per_row = blocks_per_row + .checked_mul(128) + .and_then(|value| value.checked_mul(left.bit_width as usize + right.bit_width as usize)) + .ok_or_else(|| vortex_err!("ALPRD row width overflow"))?; + Ok(Some(( + ALPRDReadPlan { + serialized: node.clone(), + descriptors, + left, + right, + patch_buffers: patch_buffers.into(), + patch_metadata, + patch_indices_dtype: patch_metadata.indices_dtype()?, + left_parts_dtype, + left_parts_dictionary: metadata.left_parts_dictionary()?, + right_bit_width: metadata.right_bit_width()?, + element_dtype: element_dtype.as_ref().clone(), + list_size: *list_size, + row_count, + }, + bytes_per_row, + ))) +} + +fn try_bitpacked_plan( + node: &SerializedArray, + ptype: vortex_array::dtype::PType, + ctx: &ReadContext, + len: usize, + descriptors: &[SerializedBuffer], +) -> VortexResult> { + if ctx + .resolve(node.encoding_id()) + .is_none_or(|id| id.as_str() != "fastlanes.bitpacked") + || node.nchildren() != 0 + || node.buffer_indices().len() != 1 + { + return Ok(None); + } + let metadata = BitPackedMetadata::decode(node.metadata())?; + if metadata.patches().is_some() || metadata.offset()? != 0 { + return Ok(None); + } + let bit_width = metadata.bit_width()?; + let Some(descriptor) = descriptors.get(node.buffer_indices()[0]).cloned() else { + return Ok(None); + }; + let expected_len = len + .div_ceil(1024) + .checked_mul(128 * bit_width as usize) + .ok_or_else(|| vortex_err!("Bit-packed buffer length overflow"))?; + if descriptor.range().len() != expected_len { + return Ok(None); + } + Ok(Some(BitPackedReadPlan { + descriptor, + ptype, + bit_width, + offset: 0, + })) +} + +fn bitpacked_range(plan: &BitPackedReadPlan, values: Range) -> Option> { + if plan.offset != 0 || !values.start.is_multiple_of(1024) || !values.end.is_multiple_of(1024) { + return None; + } + let bytes_per_block = 128usize.checked_mul(plan.bit_width as usize)?; + let start = plan + .descriptor + .range() + .start + .checked_add((values.start / 1024).checked_mul(bytes_per_block)?)?; + let end = plan + .descriptor + .range() + .start + .checked_add((values.end / 1024).checked_mul(bytes_per_block)?)?; + Some(start..end) +} + +fn collect_buffer_indices(node: &SerializedArray, output: &mut BTreeSet) { + output.extend(node.buffer_indices()); + for index in 0..node.nchildren() { + collect_buffer_indices(&node.child(index), output); + } +} + +fn collect_raw_buffers( + node: &SerializedArray, + dtype: &DType, + ctx: &ReadContext, + row_multiplier: usize, + root_row_count: usize, + descriptors: &[SerializedBuffer], + output: &mut Vec, +) -> VortexResult { + let Some(id) = ctx.resolve(node.encoding_id()) else { + return Ok(false); + }; + + if id == Primitive.id() { + let DType::Primitive(ptype, _) = dtype else { + return Ok(false); + }; + if node.nchildren() != 0 || node.buffer_indices().len() != 1 { + return Ok(false); + } + let index = node.buffer_indices()[0]; + let Some(descriptor) = descriptors.get(index) else { + return Ok(false); + }; + output.push(PlannedBuffer { + descriptor: descriptor.clone(), + bytes_per_row: row_multiplier + .checked_mul(ptype.byte_width()) + .ok_or_else(|| vortex_err!("Partial primitive row width overflow"))?, + row_granularity: 1, + bytes_per_granule: row_multiplier + .checked_mul(ptype.byte_width()) + .ok_or_else(|| vortex_err!("Partial primitive row width overflow"))?, + }); + return Ok(true); + } + + if id == FixedSizeList.id() { + let DType::FixedSizeList(element_dtype, list_size, _) = dtype else { + return Ok(false); + }; + if node.nbuffers() != 0 || node.nchildren() != 1 { + return Ok(false); + } + let multiplier = row_multiplier + .checked_mul(*list_size as usize) + .ok_or_else(|| vortex_err!("Partial fixed-size-list width overflow"))?; + return collect_raw_buffers( + &node.child(0), + element_dtype, + ctx, + multiplier, + root_row_count, + descriptors, + output, + ); + } + + if id == Struct.id() { + let DType::Struct(fields, _) = dtype else { + return Ok(false); + }; + if node.nbuffers() != 0 || node.nchildren() != fields.nfields() { + return Ok(false); + } + for (index, field_dtype) in fields.fields().enumerate() { + if !collect_raw_buffers( + &node.child(index), + &field_dtype, + ctx, + row_multiplier, + root_row_count, + descriptors, + output, + )? { + return Ok(false); + } + } + return Ok(true); + } + + if id.as_str() == "vortex.alprd" { + if !matches!(dtype, DType::Primitive(_, _)) + || node.nbuffers() != 0 + || node.nchildren() != 2 + || row_multiplier == 0 + || root_row_count == 0 + { + return Ok(false); + } + let granularity = 1024 / gcd(1024, row_multiplier); + for child_index in 0..2 { + let child = node.child(child_index); + let Some(child_id) = ctx.resolve(child.encoding_id()) else { + return Ok(false); + }; + if child_id.as_str() != "fastlanes.bitpacked" + || child.nchildren() != 0 + || child.buffer_indices().len() != 1 + { + return Ok(false); + } + let index = child.buffer_indices()[0]; + let Some(descriptor) = descriptors.get(index) else { + return Ok(false); + }; + let granules = root_row_count.div_ceil(granularity); + if descriptor.range().len() % granules != 0 { + return Ok(false); + } + let bytes_per_granule = descriptor.range().len() / granules; + output.push(PlannedBuffer { + descriptor: descriptor.clone(), + bytes_per_row: bytes_per_granule.div_ceil(granularity), + row_granularity: granularity, + bytes_per_granule, + }); + } + return Ok(true); + } + + Ok(false) +} + +fn gcd(mut left: usize, mut right: usize) -> usize { + while right != 0 { + (left, right) = (right, left % right); + } + left +} + +fn checked_lcm(left: usize, right: usize) -> VortexResult { + left.checked_div(gcd(left, right)) + .and_then(|value| value.checked_mul(right)) + .ok_or_else(|| vortex_err!("Partial row granularity overflow")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn selected_page_runs_are_exact_and_coalesced() { + let pages = [10..14, 14..18]; + let mask = Mask::from_indices(8, [1, 2, 6]); + + assert_eq!( + selected_page_runs(&pages, &(10..18), &mask), + Some(vec![11..13, 16..17]) + ); + } +} diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index aa7609f1659..f3d89f78ef4 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -4,6 +4,7 @@ use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; +use std::sync::OnceLock; use futures::FutureExt; use futures::future::BoxFuture; @@ -22,6 +23,8 @@ use vortex_session::VortexSession; use crate::layouts::SharedArrayFuture; use crate::layouts::flat::FlatLayout; +use crate::layouts::flat::partial::PartialReadPlan; +use crate::layouts::flat::partial::RegisteredPartialRead; use crate::reader::LayoutReader; use crate::reader::RowSplits; use crate::reader::SplitRange; @@ -34,11 +37,13 @@ use crate::segments::SegmentSource; // actual expression? Perhaps all expressions are given a selection mask to decide for themselves? const EXPR_EVAL_THRESHOLD: f64 = 0.2; +#[derive(Clone)] pub struct FlatReader { layout: FlatLayout, name: Arc, segment_source: Arc, session: VortexSession, + partial_plan: Arc>>, } impl FlatReader { @@ -53,9 +58,36 @@ impl FlatReader { name, segment_source, session, + partial_plan: Arc::new(OnceLock::new()), } } + fn register_partial( + &self, + row_range: &Range, + mask: &Mask, + ) -> Option { + if !PartialReadPlan::supports_mask(mask) { + return None; + } + let plan = self + .partial_plan + .get_or_init(|| match PartialReadPlan::try_new(&self.layout) { + Ok(plan) => plan, + Err(error) => { + tracing::debug!("Flat partial-read plan disabled: {error}"); + None + } + }); + plan.as_ref()?.register( + &self.segment_source, + self.layout.segment_id(), + usize::try_from(self.layout.row_count()).ok()?, + row_range, + mask, + ) + } + /// Register the segment request and return a future that would resolve into the deserialised array. fn array_future(&self) -> SharedArrayFuture { let row_count = @@ -131,18 +163,87 @@ impl LayoutReader for FlatReader { .vortex_expect("Row range begin must fit within FlatLayout size") ..usize::try_from(row_range.end) .vortex_expect("Row range end must fit within FlatLayout size"); + if !mask.partial_reads_allowed() { + let name = Arc::clone(&self.name); + let array = self.array_future(); + let expr = expr.clone(); + let session = self.session.clone(); + + return Ok(MaskFuture::new(mask.len(), async move { + let mut array = array.await?; + let mask = mask.await?; + + if row_range.start > 0 || row_range.end < array.len() { + array = array.slice(row_range.clone())?; + } + + let mask_density = mask.density(); + let array_mask = if mask_density < EXPR_EVAL_THRESHOLD { + let array = array.apply_bound(&expr)?; + let array = array.filter(mask.clone())?; + let mut ctx = session.create_execution_ctx(); + let array_mask = array.null_as_false().execute(&mut ctx)?; + mask.intersect_by_rank(&array_mask) + } else { + let array = array.apply_bound(&expr)?; + let mut ctx = session.create_execution_ctx(); + let array_mask = array.null_as_false().execute(&mut ctx)?; + mask.bitand(&array_mask) + }; + + trace!( + "Flat mask evaluation {} - {} (mask = {}) => {}", + name, + expr, + mask_density, + array_mask.density(), + ); + Ok(array_mask) + })); + } let name = Arc::clone(&self.name); - let array = self.array_future(); let expr = expr.clone(); let session = self.session.clone(); + let reader = self.clone(); + let partial_reads_allowed = mask.partial_reads_allowed(); + let registered = partial_reads_allowed + .then(|| mask.upper_bound()) + .flatten() + .and_then(|upper_bound| self.register_partial(&row_range, upper_bound)); + let eager_array = + (mask.upper_bound_is_exact() && registered.is_none()).then(|| self.array_future()); Ok(MaskFuture::new(mask.len(), async move { // TODO(ngates): if the mask density is low enough, or if the mask is dense within a range // (as often happens with zone map pruning), then we could slice/filter the array prior // to evaluating the expression. - let mut array = array.clone().await?; let mask = mask.await?; + if let Some(registered) = registered.or_else(|| { + partial_reads_allowed + .then(|| reader.register_partial(&row_range, &mask)) + .flatten() + }) { + let array = registered + .resolve( + reader.layout.dtype(), + &row_range, + &mask, + reader.layout.array_ctx(), + &session, + ) + .await?; + let array = array.apply_bound(&expr)?; + let mut ctx = session.create_execution_ctx(); + let array_mask = array.null_as_false().execute(&mut ctx)?; + return Ok(mask.intersect_by_rank(&array_mask)); + } + + let mut array = match eager_array { + Some(array) => array.await?, + None => reader.array_future().await?, + }; + // Slice the array based on the row mask. if row_range.start > 0 || row_range.end < array.len() { array = array.slice(row_range.clone())?; @@ -190,16 +291,64 @@ impl LayoutReader for FlatReader { .vortex_expect("Row range begin must fit within FlatLayout size") ..usize::try_from(row_range.end) .vortex_expect("Row range end must fit within FlatLayout size"); + if !mask.partial_reads_allowed() { + let name = Arc::clone(&self.name); + let array = self.array_future(); + let expr = expr.clone(); + + return Ok(async move { + trace!("Flat array evaluation {} - {}", name, expr); + + let mut array = array.await?; + let mask = mask.await?; + + if row_range.start > 0 || row_range.end < array.len() { + array = array.slice(row_range.clone())?; + } + if !mask.all_true() { + array = array.filter(mask)?; + } + array = array.apply_bound(&expr)?; + Ok(array) + } + .boxed()); + } let name = Arc::clone(&self.name); - let array = self.array_future(); let expr = expr.clone(); + let reader = self.clone(); + let partial_reads_allowed = mask.partial_reads_allowed(); + let registered = partial_reads_allowed + .then(|| mask.upper_bound()) + .flatten() + .and_then(|upper_bound| self.register_partial(&row_range, upper_bound)); + let eager_array = ((!partial_reads_allowed || mask.upper_bound_is_exact()) + && registered.is_none()) + .then(|| self.array_future()); Ok(async move { trace!("Flat array evaluation {} - {}", name, expr); - let mut array = array.clone().await?; let mask = mask.await?; + if let Some(registered) = registered { + let mut array = registered + .resolve( + reader.layout.dtype(), + &row_range, + &mask, + reader.layout.array_ctx(), + &reader.session, + ) + .await?; + array = array.apply_bound(&expr)?; + return Ok(array); + } + + let mut array = match eager_array { + Some(array) => array.await?, + None => reader.array_future().await?, + }; + // Slice the array based on the row mask. if row_range.start > 0 || row_range.end < array.len() { array = array.slice(row_range.clone())?; @@ -228,8 +377,12 @@ impl LayoutReader for FlatReader { #[cfg(test)] mod test { + use std::ops::Range; use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use parking_lot::Mutex; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; @@ -245,14 +398,46 @@ mod test { use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; + use vortex_mask::Mask; use crate::LayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; + use crate::segments::SegmentSource; + use crate::segments::SharedSegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::new_session; + #[derive(Clone, Default)] + struct RangedTestSource { + inner: Arc, + ranges: Arc>>>, + whole_requests: Arc, + } + + impl SegmentSource for RangedTestSource { + fn preferred_read_size(&self) -> Option { + Some(16) + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.inner.segment_len(id) + } + + fn request(&self, id: SegmentId) -> SegmentFuture { + self.whole_requests.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + self.ranges.lock().push(range.clone()); + self.inner.request_range(id, range) + } + } + #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { @@ -370,4 +555,154 @@ mod test { assert_arrays_eq!(result, expected, &mut ctx); }) } + + #[test] + fn sparse_projection_reads_only_virtual_pages() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0i32..64).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(source.clone()), + &session, + &Default::default(), + )?; + let expr = root().bind(reader.dtype())?; + let result = reader + .projection_evaluation( + &(0..64), + &expr, + MaskFuture::ready(Mask::from_indices(64, [1, 10])), + )? + .await?; + + let expected = PrimitiveArray::from_iter([1i32, 10]).into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!(source.whole_requests.load(Ordering::Relaxed), 0); + assert_eq!(*source.ranges.lock(), [0..16, 32..48]); + + let result = reader + .projection_evaluation( + &(0..64), + &expr, + MaskFuture::ready(Mask::from_indices(64, [1, 10])), + )? + .await?; + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!( + *source.ranges.lock(), + [0..16, 32..48, 0..16, 32..48], + "separate evaluations must not retain page data" + ); + Ok(()) + }) + } + + #[test] + fn dense_projection_chooses_whole_segment() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0i32..64).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(source.clone()), + &session, + &Default::default(), + )?; + let expr = root().bind(reader.dtype())?; + let mask = Mask::from_indices(64, (0..64).step_by(2)); + let result = reader + .projection_evaluation(&(0..64), &expr, MaskFuture::ready(mask.clone()))? + .await?; + + assert_arrays_eq!(result, array.filter(mask)?, &mut ctx); + assert_eq!(source.whole_requests.load(Ordering::Relaxed), 1); + assert!(source.ranges.lock().is_empty()); + Ok(()) + }) + } + + #[test] + fn filter_and_projection_share_pages_while_scan_is_in_flight() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0i32..64).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(SharedSegmentSource::new(source.clone())), + &session, + &Default::default(), + )?; + let projection_expr = root().bind(reader.dtype())?; + let filter_expr = gt(root(), lit(-1i32)).bind(reader.dtype())?; + let mask = Mask::from_indices(64, [1, 10]); + + let filter = reader.filter_evaluation( + &(0..64), + &filter_expr, + MaskFuture::ready(mask.clone()), + )?; + let projection = reader.projection_evaluation( + &(0..64), + &projection_expr, + MaskFuture::ready(mask), + )?; + + let (filter_mask, result) = futures::try_join!(filter, projection)?; + assert_eq!( + filter_mask.indices(), + Mask::from_indices(64, [1, 10]).indices() + ); + let expected = PrimitiveArray::from_iter([1i32, 10]).into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!( + *source.ranges.lock(), + [0..16, 32..48], + "one in-flight request should serve filter and projection" + ); + Ok(()) + }) + } } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4d8565fdd10..a8cecc4cdf4 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -399,8 +399,8 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32), children: 2 - ├── elements: vortex.flat, dtype: i32, segment: 0 - └── offsets: vortex.flat, dtype: u64, segment: 1 + ├── elements: vortex.flat, dtype: i32, segment 0, buffers=[20B], total=20B + └── offsets: vortex.flat, dtype: u64, segment 1, buffers=[32B], total=32B "); Ok(()) } @@ -416,9 +416,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32)?, children: 3 - ├── elements: vortex.flat, dtype: i32, segment: 0 - ├── offsets: vortex.flat, dtype: u64, segment: 1 - └── validity: vortex.flat, dtype: bool, segment: 2 + ├── elements: vortex.flat, dtype: i32, segment 0, buffers=[20B], total=20B + ├── offsets: vortex.flat, dtype: u64, segment 1, buffers=[32B], total=32B + └── validity: vortex.flat, dtype: bool, segment 2, buffers=[1B], total=1B "); Ok(()) } @@ -428,7 +428,7 @@ mod tests { async fn non_list_input_routes_to_fallback() -> VortexResult<()> { let primitive = buffer![1i32, 2, 3].into_array(); let layout = write(&flat_list_strategy(), primitive).await?; - insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B"); Ok(()) } @@ -478,9 +478,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list({a=i32, b=i32}), children: 2 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 1 - │ └── b: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── a: vortex.flat, dtype: i32, segment 1, buffers=[20B], total=20B + │ └── b: vortex.flat, dtype: i32, segment 2, buffers=[20B], total=20B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[32B], total=32B "); Ok(()) } @@ -506,9 +506,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── elements: vortex.flat, dtype: i32, segment 1, buffers=[24B], total=24B + │ └── offsets: vortex.flat, dtype: u64, segment 2, buffers=[40B], total=40B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[24B], total=24B "); Ok(()) } @@ -539,10 +539,10 @@ mod tests { vortex.list, dtype: list(list(list(i32))), children: 2 ├── elements: vortex.list, dtype: list(list(i32)), children: 2 │ ├── elements: vortex.list, dtype: list(i32), children: 2 - │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 - │ │ └── offsets: vortex.flat, dtype: u64, segment: 3 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ │ ├── elements: vortex.flat, dtype: i32, segment 2, buffers=[16B], total=16B + │ │ └── offsets: vortex.flat, dtype: u64, segment 3, buffers=[24B], total=24B + │ └── offsets: vortex.flat, dtype: u64, segment 1, buffers=[16B], total=16B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[16B], total=16B "); Ok(()) } @@ -572,11 +572,11 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.chunked, dtype: list(i32), children: 2 ├── [0]: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 0 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 + │ ├── elements: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + │ └── offsets: vortex.flat, dtype: u64, segment 1, buffers=[24B], total=24B └── [1]: vortex.list, dtype: list(i32), children: 2 - ├── elements: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 3 + ├── elements: vortex.flat, dtype: i32, segment 2, buffers=[16B], total=16B + └── offsets: vortex.flat, dtype: u64, segment 3, buffers=[24B], total=24B "); Ok(()) } diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..c70b93909d9 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -406,10 +406,10 @@ mod tests { .into_array(); let layout = write(&flat_table(), struct_array).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 - ├── a: vortex.flat, dtype: i32, segment: 0 - └── b: vortex.flat, dtype: i32, segment: 1 + ├── a: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + └── b: vortex.flat, dtype: i32, segment 1, buffers=[12B], total=12B "); Ok(()) } @@ -432,12 +432,12 @@ mod tests { .into_array(); let layout = write(&flat_table().with_list_layout(), outer).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── elements: vortex.flat, dtype: i32, segment 1, buffers=[24B], total=24B + │ └── offsets: vortex.flat, dtype: u64, segment 2, buffers=[40B], total=40B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[24B], total=24B "); Ok(()) } @@ -463,14 +463,14 @@ mod tests { let st = StructArray::from_fields([("items", items)].as_slice())?.into_array(); let layout = write(&flat_table().with_list_layout(), st).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1 └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 2 - │ └── b: vortex.flat, dtype: i32, segment: 3 - ├── offsets: vortex.flat, dtype: u64, segment: 0 - └── validity: vortex.flat, dtype: bool, segment: 1 + │ ├── a: vortex.flat, dtype: i32, segment 2, buffers=[20B], total=20B + │ └── b: vortex.flat, dtype: i32, segment 3, buffers=[20B], total=20B + ├── offsets: vortex.flat, dtype: u64, segment 0, buffers=[32B], total=32B + └── validity: vortex.flat, dtype: bool, segment 1, buffers=[1B], total=1B "); Ok(()) } @@ -502,14 +502,14 @@ mod tests { ) .with_list_layout(); let layout = write(&dispatcher, chunked).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32), children: 2 ├── elements: vortex.chunked, dtype: i32, children: 2 - │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 + │ ├── [0]: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + │ └── [1]: vortex.flat, dtype: i32, segment 1, buffers=[16B], total=16B └── offsets: vortex.chunked, dtype: u64, children: 2 - ├── [0]: vortex.flat, dtype: u64, segment: 2 - └── [1]: vortex.flat, dtype: u64, segment: 3 + ├── [0]: vortex.flat, dtype: u64, segment 2, buffers=[24B], total=24B + └── [1]: vortex.flat, dtype: u64, segment 3, buffers=[16B], total=16B "); Ok(()) } @@ -569,7 +569,7 @@ mod tests { async fn non_struct_input_uses_leaf() -> VortexResult<()> { let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let layout = write(&flat_table(), primitive).await?; - insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B"); Ok(()) } @@ -601,14 +601,14 @@ mod tests { let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array(); let layout = write(&dispatcher, chunked).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 ├── a: vortex.chunked, dtype: i32, children: 2 - │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 + │ ├── [0]: vortex.flat, dtype: i32, segment 0, buffers=[8B], total=8B + │ └── [1]: vortex.flat, dtype: i32, segment 1, buffers=[4B], total=4B └── b: vortex.chunked, dtype: i32, children: 2 - ├── [0]: vortex.flat, dtype: i32, segment: 2 - └── [1]: vortex.flat, dtype: i32, segment: 3 + ├── [0]: vortex.flat, dtype: i32, segment 2, buffers=[8B], total=8B + └── [1]: vortex.flat, dtype: i32, segment 3, buffers=[4B], total=4B "); Ok(()) } @@ -628,10 +628,10 @@ mod tests { let strategy = flat_table().with_field_writer(field_path!(a), Arc::new(FlatLayoutStrategy::default())); let layout = write(&strategy, struct_array).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 - ├── a: vortex.flat, dtype: i32, segment: 0 - └── b: vortex.flat, dtype: i32, segment: 1 + ├── a: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + └── b: vortex.flat, dtype: i32, segment 1, buffers=[12B], total=12B "); Ok(()) } diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 218efb64a0d..d7e8342c664 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -69,8 +69,10 @@ pub fn split_exec( let reader = Arc::clone(&ctx.reader); let filter = Arc::clone(filter); let row_range = row_range.clone(); + let filter_upper_bound = row_mask.clone(); + let partial_reads_allowed = !row_mask.all_true(); - MaskFuture::new(row_mask.len(), async move { + let filter_mask = MaskFuture::new(row_mask.len(), async move { let mut mask = row_mask; let mut dynamic_versions = vec![None; filter.conjuncts().len()]; @@ -117,8 +119,13 @@ pub fn split_exec( return Ok(mask); } + let mask_future = if partial_reads_allowed { + MaskFuture::ready(mask) + } else { + MaskFuture::ready(mask).without_partial_reads() + }; let conjunct_mask = reader - .filter_evaluation(&row_range, conjunct, MaskFuture::ready(mask))? + .filter_evaluation(&row_range, conjunct, mask_future)? .await?; filter.report_selectivity(idx, conjunct_mask.density()); @@ -128,6 +135,12 @@ pub fn split_exec( Ok(mask) }) + .with_upper_bound(filter_upper_bound); + if partial_reads_allowed { + filter_mask.with_partial_reads() + } else { + filter_mask + } } }; diff --git a/vortex-layout/src/segments/cache.rs b/vortex-layout/src/segments/cache.rs index 1f7f5e91f5c..8d24ed3aeb6 100644 --- a/vortex-layout/src/segments/cache.rs +++ b/vortex-layout/src/segments/cache.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use async_trait::async_trait; @@ -146,6 +147,14 @@ impl SegmentCacheSourceAdapter { } impl SegmentSource for SegmentCacheSourceAdapter { + fn preferred_read_size(&self) -> Option { + self.source.preferred_read_size() + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.source.segment_len(id) + } + fn request(&self, id: SegmentId) -> SegmentFuture { let cache = Arc::clone(&self.cache); let delegate = self.source.request(id); @@ -166,4 +175,59 @@ impl SegmentSource for SegmentCacheSourceAdapter { } .boxed() } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + let cache = Arc::clone(&self.cache); + let delegate = self.source.request_range(id, range.clone()); + + async move { + if let Ok(Some(segment)) = cache.get(id).await { + let start = usize::try_from(range.start)?; + let end = usize::try_from(range.end)?; + if start > end || end > segment.len() { + return Err(vortex_error::vortex_err!( + "Segment {} range {}..{} is out of bounds for cached segment length {}", + id, + range.start, + range.end, + segment.len() + )); + } + tracing::debug!("Resolved segment {} range {:?} from cache", id, range); + return Ok(BufferHandle::new_host(segment.slice(start..end))); + } + delegate.await + } + .boxed() + } + + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + let delegates = self.source.request_ranges(id, ranges.clone()); + ranges + .into_iter() + .zip(delegates) + .map(|(range, delegate)| { + let cache = Arc::clone(&self.cache); + async move { + if let Ok(Some(segment)) = cache.get(id).await { + let start = usize::try_from(range.start)?; + let end = usize::try_from(range.end)?; + if start > end || end > segment.len() { + return Err(vortex_error::vortex_err!( + "Segment {} range {}..{} is out of bounds for cached segment length {}", + id, + range.start, + range.end, + segment.len() + )); + } + tracing::debug!("Resolved segment {} range {:?} from cache", id, range); + return Ok(BufferHandle::new_host(segment.slice(start..end))); + } + delegate.await + } + .boxed() + }) + .collect() + } } diff --git a/vortex-layout/src/segments/shared.rs b/vortex-layout/src/segments/shared.rs index c794daf608e..c2355163338 100644 --- a/vortex-layout/src/segments/shared.rs +++ b/vortex-layout/src/segments/shared.rs @@ -1,12 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use futures::FutureExt; use futures::TryFutureExt; use futures::future::BoxFuture; +use futures::future::Shared; use futures::future::WeakShared; +use futures::future::join_all; use vortex_array::buffer::BufferHandle; use vortex_error::SharedVortexResult; use vortex_error::VortexError; @@ -22,40 +25,172 @@ use crate::segments::SegmentSource; /// request. pub struct SharedSegmentSource { inner: S, - in_flight: DashMap>, + in_flight: Arc>, +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum SegmentRequest { + Whole(SegmentId), + Range(SegmentId, Range), + Ranges(SegmentId, Arc<[Range]>), } type SharedSegmentFuture = BoxFuture<'static, SharedVortexResult>; +type SharedBatchResult = Arc<[SharedVortexResult]>; +type SharedBatchFuture = BoxFuture<'static, SharedBatchResult>; + +enum WeakInFlight { + Segment(WeakShared), + Batch(WeakShared), +} + +impl WeakInFlight { + fn upgrade_segment(&self) -> Option { + let Self::Segment(future) = self else { + return None; + }; + future + .upgrade() + .map(|future| future.map_err(VortexError::from).boxed()) + } + + fn upgrade_batch(&self) -> Option> { + let Self::Batch(future) = self else { + return None; + }; + future.upgrade() + } +} + +struct InFlightGuard { + in_flight: Arc>, + request: SegmentRequest, +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + self.in_flight.remove(&self.request); + } +} impl SharedSegmentSource { /// Create a new `SharedSegmentSource` wrapping the provided inner source. pub fn new(inner: S) -> Self { Self { inner, - in_flight: DashMap::default(), + in_flight: Arc::default(), } } } impl SegmentSource for SharedSegmentSource { + fn preferred_read_size(&self) -> Option { + self.inner.preferred_read_size() + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.inner.segment_len(id) + } + fn request(&self, id: SegmentId) -> SegmentFuture { + self.request_shared(SegmentRequest::Whole(id)) + } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + self.request_shared(SegmentRequest::Range(id, range)) + } + + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + if ranges.is_empty() { + return Vec::new(); + } + let len = ranges.len(); + let request = SegmentRequest::Ranges(id, Arc::from(ranges.clone())); + + loop { + match self.in_flight.entry(request.clone()) { + Entry::Occupied(entry) => { + if let Some(batch) = entry.get().upgrade_batch() { + return batch_outputs(batch, len); + } + entry.remove(); + } + Entry::Vacant(entry) => { + let futures = self.inner.request_ranges(id, ranges); + let guard = InFlightGuard { + in_flight: Arc::clone(&self.in_flight), + request, + }; + let batch = async move { + let _guard = guard; + Arc::from( + join_all( + futures + .into_iter() + .map(|future| async move { future.await.map_err(Arc::new) }), + ) + .await, + ) + } + .boxed() + .shared(); + entry.insert(WeakInFlight::Batch( + batch + .downgrade() + .vortex_expect("new shared batch cannot be complete"), + )); + return batch_outputs(batch, len); + } + } + } + } +} + +fn batch_outputs(batch: Shared, len: usize) -> Vec { + (0..len) + .map(|index| { + let batch = batch.clone(); + async move { batch.await[index].clone().map_err(VortexError::from) }.boxed() + as SegmentFuture + }) + .collect() +} + +impl SharedSegmentSource { + fn request_shared(&self, request: SegmentRequest) -> SegmentFuture { loop { - match self.in_flight.entry(id) { + match self.in_flight.entry(request.clone()) { Entry::Occupied(e) => { - if let Some(shared_future) = e.get().upgrade() { - return shared_future.map_err(VortexError::from).boxed(); + if let Some(shared_future) = e.get().upgrade_segment() { + return shared_future; } else { // The future has been dropped, remove the entry and try again. e.remove(); } } Entry::Vacant(e) => { - let future = self.inner.request(id).map_err(Arc::new).boxed().shared(); - e.insert( + let inner_future = match &request { + SegmentRequest::Whole(id) => self.inner.request(*id), + SegmentRequest::Range(id, range) => { + self.inner.request_range(*id, range.clone()) + } + SegmentRequest::Ranges(..) => unreachable!(), + }; + let guard = InFlightGuard { + in_flight: Arc::clone(&self.in_flight), + request, + }; + let future = async move { + let _guard = guard; + inner_future.await.map_err(Arc::new) + } + .boxed() + .shared(); + e.insert(WeakInFlight::Segment( future .downgrade() .vortex_expect("just created, cannot be polled to completion"), - ); + )); return future.map_err(VortexError::from).boxed(); } } @@ -69,6 +204,7 @@ mod tests { use std::sync::atomic::Ordering; use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; use super::*; use crate::segments::SegmentSink; @@ -80,6 +216,8 @@ mod tests { struct CountingSegmentSource { segments: TestSegments, request_count: Arc, + range_request_count: Arc, + range_batch_count: Arc, } impl SegmentSource for CountingSegmentSource { @@ -87,6 +225,19 @@ mod tests { self.request_count.fetch_add(1, Ordering::SeqCst); self.segments.request(id) } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + self.range_request_count.fetch_add(1, Ordering::SeqCst); + self.segments.request_range(id, range) + } + + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + self.range_batch_count.fetch_add(1, Ordering::SeqCst); + ranges + .into_iter() + .map(|range| self.request_range(id, range)) + .collect() + } } #[tokio::test] @@ -116,6 +267,7 @@ mod tests { // The inner source should have been called only once assert_eq!(source.request_count.load(Ordering::Relaxed), 1); + assert!(shared_source.in_flight.is_empty()); } #[tokio::test] @@ -139,6 +291,7 @@ mod tests { let _future = shared_source.request(id); // Future is dropped here } + assert!(shared_source.in_flight.is_empty()); // A new request should still work correctly let result = shared_source.request(id).await; @@ -147,4 +300,42 @@ mod tests { // Should have made 2 requests since the first was dropped before completion assert_eq!(source.request_count.load(Ordering::Relaxed), 2); } + + #[tokio::test] + async fn test_shared_source_deduplicates_identical_ranges() -> VortexResult<()> { + let source = CountingSegmentSource::default(); + let data = ByteBuffer::from(vec![1, 2, 3, 4]); + let seq_id = SequenceId::root().downgrade(); + source.segments.write(seq_id, vec![data]).await?; + + let shared_source = SharedSegmentSource::new(source.clone()); + let id = SegmentId::from(0); + let (first, second) = futures::join!( + shared_source.request_range(id, 1..3), + shared_source.request_range(id, 1..3) + ); + assert_eq!(first?.unwrap_host(), ByteBuffer::from(vec![2, 3])); + assert_eq!(second?.unwrap_host(), ByteBuffer::from(vec![2, 3])); + assert_eq!(source.range_request_count.load(Ordering::Relaxed), 1); + assert!(shared_source.in_flight.is_empty()); + Ok(()) + } + + #[tokio::test] + async fn test_shared_source_forwards_missing_ranges_as_one_batch() -> VortexResult<()> { + let source = CountingSegmentSource::default(); + let data = ByteBuffer::from(vec![1, 2, 3, 4]); + let seq_id = SequenceId::root().downgrade(); + source.segments.write(seq_id, vec![data]).await?; + + let shared_source = SharedSegmentSource::new(source.clone()); + let reads = shared_source.request_ranges(SegmentId::from(0), vec![0..1, 2..4]); + let mut results = join_all(reads).await.into_iter(); + assert_eq!(results.next().vortex_expect("first range")?.len(), 1); + assert_eq!(results.next().vortex_expect("second range")?.len(), 2); + assert_eq!(source.range_batch_count.load(Ordering::Relaxed), 1); + assert_eq!(source.range_request_count.load(Ordering::Relaxed), 2); + assert!(shared_source.in_flight.is_empty()); + Ok(()) + } } diff --git a/vortex-layout/src/segments/source.rs b/vortex-layout/src/segments/source.rs index 5c709f5a7ad..45ab59b5c32 100644 --- a/vortex-layout/src/segments/source.rs +++ b/vortex-layout/src/segments/source.rs @@ -1,9 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; + +use futures::FutureExt; use futures::future::BoxFuture; use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::segments::SegmentId; /// Static future resolving to a segment byte buffer. @@ -14,6 +18,55 @@ pub type SegmentFuture = BoxFuture<'static, VortexResult>; /// Implementations may issue asynchronous file reads, object-store requests, cache lookups, or /// in-memory buffer slices. Returned futures must be independent and safe to poll concurrently. pub trait SegmentSource: 'static + Send + Sync { + /// Preferred size of independently requested byte ranges for this source. + /// + /// Layout readers can use this hint to divide a logical segment into canonical read ranges. + /// Returning `None` asks readers to preserve whole-segment reads. + fn preferred_read_size(&self) -> Option { + None + } + + /// Return the serialized length of `id`, when it is known without issuing I/O. + fn segment_len(&self, _id: SegmentId) -> Option { + None + } + /// Request a segment, returning a future that will eventually resolve to the segment data. fn request(&self, id: SegmentId) -> SegmentFuture; + + /// Request a byte range relative to the start of a segment. + /// + /// Sources backed by random-access storage should override this method. The default keeps + /// custom sources compatible by reading the segment and slicing it after bounds checking. + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + let segment = self.request(id); + async move { + let segment = segment.await?; + let start = usize::try_from(range.start)?; + let end = usize::try_from(range.end)?; + if start > end || end > segment.len() { + vortex_bail!( + "Segment {} range {}..{} is out of bounds for a {}-byte segment", + id, + range.start, + range.end, + segment.len() + ); + } + Ok(segment.slice(start..end)) + } + .boxed() + } + + /// Register multiple ranges from one segment together. + /// + /// The returned futures correspond positionally to `ranges`. Sources can override this to + /// amortize registration while retaining independent canonical range futures for sharing and + /// coalescing. + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + ranges + .into_iter() + .map(|range| self.request_range(id, range)) + .collect() + } } diff --git a/vortex-layout/src/segments/test.rs b/vortex-layout/src/segments/test.rs index d880d15cc1a..b6a3d3007e2 100644 --- a/vortex-layout/src/segments/test.rs +++ b/vortex-layout/src/segments/test.rs @@ -26,6 +26,13 @@ pub struct TestSegments { } impl SegmentSource for TestSegments { + fn segment_len(&self, id: SegmentId) -> Option { + self.segments + .lock() + .get(*id as usize) + .and_then(|buffer| u64::try_from(buffer.len()).ok()) + } + fn request(&self, id: SegmentId) -> SegmentFuture { let buffer = self.segments.lock().get(*id as usize).cloned(); async move { diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 57392ed627d..385e57309f2 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -115,6 +115,11 @@ name = "take_union" harness = false test = false +[[bench]] +name = "random_access_reconstruct" +harness = false +test = false + [[bench]] name = "pipeline" harness = false diff --git a/vortex/benches/random_access_reconstruct.rs b/vortex/benches/random_access_reconstruct.rs new file mode 100644 index 00000000000..a86adae3a05 --- /dev/null +++ b/vortex/benches/random_access_reconstruct.rs @@ -0,0 +1,364 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! In-memory reconstruction benchmarks for random-access feature-vector results. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use vortex::VortexSessionDefault; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::ExecutionCtx; +use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; +use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::FixedSizeListArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::StructArray; +use vortex::array::buffer::BufferHandle; +use vortex::array::dtype::FieldNames; +use vortex::array::dtype::PType; +use vortex::array::patches::Patches; +use vortex::array::validity::Validity; +use vortex::arrow::ArrowSessionExt; +use vortex::encodings::alp::ALPRD; +use vortex::encodings::alp::ALPRDArrayExt; +use vortex::encodings::alp::ALPRDArrayOwnedExt; +use vortex::encodings::alp::RDEncoder; +use vortex::encodings::alp::RDEncoderExt; +use vortex::encodings::fastlanes::BitPacked; +use vortex::encodings::fastlanes::BitPackedArrayExt; +use vortex::session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(VortexSession::default); + +const NUM_CHUNKS: usize = 100; +const LIST_SIZE: usize = 1024; +const SEGMENT_ROWS: usize = 64; +const SELECTED_ROW: usize = 31; + +#[derive(Clone)] +struct ResidentALPRDPage { + left: BufferHandle, + left_ptype: PType, + left_bit_width: u8, + right: BufferHandle, + right_ptype: PType, + right_bit_width: u8, + dictionary: vortex::buffer::Buffer, + dtype: vortex::dtype::DType, + patch_indices: BufferHandle, + patch_indices_ptype: PType, + patch_values: BufferHandle, + patch_values_ptype: PType, + patch_count: usize, + patch_offset: usize, +} + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +fn feature_values(chunk: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..LIST_SIZE).map(|index| { + let mixed = (chunk * LIST_SIZE + index).wrapping_mul(2_654_435_761); + (mixed % 100_003) as f32 / 100_003.0 + })) +} + +fn segment_values() -> PrimitiveArray { + PrimitiveArray::from_iter((0..SEGMENT_ROWS * LIST_SIZE).map(|index| { + let mixed = index.wrapping_mul(2_654_435_761); + (mixed % 100_003) as f32 / 100_003.0 + })) +} + +fn resident_alprd_page() -> ResidentALPRDPage { + let values = segment_values(); + let encoded = RDEncoder::new(values.as_slice::()).encode(values.as_view()); + let patches = encoded.left_parts_patches().unwrap(); + let dictionary = encoded.left_parts_dictionary().clone(); + let right_bit_width = encoded.right_bit_width(); + let dtype = encoded.dtype().clone(); + let patch_offset = patches.offset(); + let mut ctx = SESSION.create_execution_ctx(); + let patch_indices = patches + .indices() + .clone() + .execute::(&mut ctx) + .unwrap(); + let patch_values = patches + .values() + .clone() + .execute::(&mut ctx) + .unwrap(); + let patch_count = patch_indices.len(); + let patch_indices_ptype = patch_indices.ptype(); + let patch_values_ptype = patch_values.ptype(); + let patch_indices = patch_indices.buffer_handle().clone(); + let patch_values = patch_values.buffer_handle().clone(); + + let parts = encoded.into_data_parts(); + let element_range = SELECTED_ROW * LIST_SIZE..(SELECTED_ROW + 1) * LIST_SIZE; + let left_page = parts.left_parts.slice(element_range.clone()).unwrap(); + let right_page = parts.right_parts.slice(element_range).unwrap(); + let left_primitive = left_page + .clone() + .execute::(&mut ctx) + .unwrap(); + let right_primitive = right_page + .clone() + .execute::(&mut ctx) + .unwrap(); + let left = left_page.as_opt::().unwrap(); + let right = right_page.as_opt::().unwrap(); + assert_eq!(left.offset(), 0); + assert_eq!(right.offset(), 0); + + ResidentALPRDPage { + left: left.packed().clone(), + left_ptype: left_primitive.ptype(), + left_bit_width: left.bit_width(), + right: right.packed().clone(), + right_ptype: right_primitive.ptype(), + right_bit_width, + dictionary, + dtype, + patch_indices, + patch_indices_ptype, + patch_values, + patch_values_ptype, + patch_count, + patch_offset, + } +} + +fn resident_pages() -> Vec { + vec![resident_alprd_page(); NUM_CHUNKS] +} + +fn reconstruct_resident_pages( + pages: &[ResidentALPRDPage], + ctx: &mut ExecutionCtx, + eager_extract: bool, +) -> ArrayRef { + let names: FieldNames = ["id", "embedding"].into_iter().collect(); + let chunks = pages + .iter() + .enumerate() + .map(|(index, page)| { + let patch_indices = PrimitiveArray::from_buffer_handle( + page.patch_indices.clone(), + page.patch_indices_ptype, + Validity::NonNullable, + ) + .into_array() + .execute::(ctx) + .unwrap() + .into_array(); + let patch_values = PrimitiveArray::from_buffer_handle( + page.patch_values.clone(), + page.patch_values_ptype, + Validity::NonNullable, + ) + .into_array(); + let full_patches = Patches::new( + SEGMENT_ROWS * LIST_SIZE, + page.patch_offset, + patch_indices, + patch_values, + None, + ) + .unwrap(); + let element_start = SELECTED_ROW * LIST_SIZE; + let patches = full_patches + .slice(element_start..element_start + LIST_SIZE) + .unwrap(); + let left = BitPacked::try_new( + page.left.clone(), + page.left_ptype, + Validity::NonNullable, + None, + page.left_bit_width, + LIST_SIZE, + 0, + ) + .unwrap() + .into_array(); + let right = BitPacked::try_new( + page.right.clone(), + page.right_ptype, + Validity::NonNullable, + None, + page.right_bit_width, + LIST_SIZE, + 0, + ) + .unwrap() + .into_array(); + let elements = ALPRD::try_new( + page.dtype.clone(), + left, + page.dictionary.clone(), + right, + page.right_bit_width, + patches, + ) + .unwrap() + .into_array(); + let elements = if eager_extract { + elements + .execute::(ctx) + .unwrap() + .into_array() + } else { + elements + }; + let embedding = FixedSizeListArray::try_new( + elements, + u32::try_from(LIST_SIZE).unwrap(), + Validity::NonNullable, + 1, + ) + .unwrap() + .into_array(); + for child in embedding.depth_first_traversal() { + child.statistics().clear_all(); + } + let id = PrimitiveArray::from_iter([i64::try_from(index).unwrap()]).into_array(); + StructArray::try_new(names.clone(), [id, embedding], 1, Validity::NonNullable) + .unwrap() + .into_array() + }) + .collect::>(); + let dtype = chunks[0].dtype().clone(); + ChunkedArray::try_new(chunks, dtype) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn canonical_chunks() -> ChunkedArray { + let list_size = u32::try_from(LIST_SIZE).unwrap(); + let chunks = (0..NUM_CHUNKS).map(|chunk| { + FixedSizeListArray::new( + feature_values(chunk).into_array(), + list_size, + Validity::NonNullable, + 1, + ) + .into_array() + }); + let chunks = chunks.collect::>(); + let dtype = chunks[0].dtype().clone(); + ChunkedArray::try_new(chunks, dtype).unwrap() +} + +fn alprd_chunks() -> ChunkedArray { + let list_size = u32::try_from(LIST_SIZE).unwrap(); + let chunks = (0..NUM_CHUNKS).map(|chunk| { + let values = feature_values(chunk); + let encoder = RDEncoder::new(values.as_slice::()); + let encoded = encoder.encode(values.as_view()).into_array(); + FixedSizeListArray::new(encoded, list_size, Validity::NonNullable, 1).into_array() + }); + let chunks = chunks.collect::>(); + let dtype = chunks[0].dtype().clone(); + ChunkedArray::try_new(chunks, dtype).unwrap() +} + +#[divan::bench] +fn concat_100_canonical_one_row_arrays(bencher: Bencher) { + let chunked = canonical_chunks().into_array(); + bencher + .with_inputs(|| (&chunked, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx).unwrap()); +} + +#[divan::bench] +fn container_concat_100_alprd_one_row_arrays(bencher: Bencher) { + let chunked = alprd_chunks().into_array(); + bencher + .with_inputs(|| (&chunked, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx).unwrap()); +} + +#[divan::bench] +fn rebuild_100_flat_arrays_from_resident_buffers(bencher: Bencher) { + let pages = resident_pages(); + assert!(pages[0].patch_count > 0); + bencher + .with_inputs(|| (&pages, SESSION.create_execution_ctx())) + .bench_refs(|(pages, ctx)| reconstruct_resident_pages(pages, ctx, false)); +} + +#[divan::bench] +fn extract_100_prebuilt_arrays_to_arrow(bencher: Bencher) { + let pages = resident_pages(); + let array = reconstruct_resident_pages(&pages, &mut SESSION.create_execution_ctx(), false); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + SESSION + .arrow() + .execute_arrow((*array).clone(), None, ctx) + .unwrap() + }); +} + +#[divan::bench] +fn rebuild_and_extract_100_arrays_to_arrow(bencher: Bencher) { + let pages = resident_pages(); + assert!(pages[0].patch_count > 0); + bencher + .with_inputs(|| (&pages, SESSION.create_execution_ctx())) + .bench_refs(|(pages, ctx)| { + let array = reconstruct_resident_pages(pages, ctx, false); + SESSION.arrow().execute_arrow(array, None, ctx).unwrap() + }); +} + +#[divan::bench] +fn rebuild_with_eager_extract_100_arrays_to_arrow(bencher: Bencher) { + let pages = resident_pages(); + assert!(pages[0].patch_count > 0); + bencher + .with_inputs(|| (&pages, SESSION.create_execution_ctx())) + .bench_refs(|(pages, ctx)| { + let array = reconstruct_resident_pages(pages, ctx, true); + SESSION.arrow().execute_arrow(array, None, ctx).unwrap() + }); +} + +#[divan::bench] +fn rebuild_100_flat_arrays_with_eager_leaf_decode(bencher: Bencher) { + let pages = resident_pages(); + assert!(pages[0].patch_count > 0); + bencher + .with_inputs(|| (&pages, SESSION.create_execution_ctx())) + .bench_refs(|(pages, ctx)| reconstruct_resident_pages(pages, ctx, true)); +} + +#[divan::bench] +fn arrow_export_100_predecoded_arrays(bencher: Bencher) { + let pages = resident_pages(); + let array = reconstruct_resident_pages(&pages, &mut SESSION.create_execution_ctx(), true); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + SESSION + .arrow() + .execute_arrow((*array).clone(), None, ctx) + .unwrap() + }); +}