diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 7123b58512f..9b6e9aaf7c4 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -258,6 +258,7 @@ mod tests { use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; use vortex_fastlanes::bitpacked_v1_id; use vortex_fastlanes::bitpacked_v2_id; + use vortex_fastlanes::for_v1_id; use vortex_fsst::FSST; use super::*; @@ -288,6 +289,12 @@ mod tests { vec![bitpacked_v2_id(), Patched.id()], false )] + #[case::for_without_bitpacked(&integer::FOR_V1, vec![for_v1_id()], false)] + #[case::for_with_bitpacked( + &integer::FOR_V1, + vec![for_v1_id(), bitpacked_v1_id(), Patched.id()], + true + )] fn test_allowed_schemes( #[case] scheme: &'static dyn Scheme, #[case] allowed: Vec, diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 00d8aeb858a..c885bd97509 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -23,10 +23,11 @@ use vortex_fastlanes::FoR; use vortex_fastlanes::FoRArrayExt; use vortex_fastlanes::FoRArraySlotsExt; use vortex_fastlanes::FoRReferences; +use vortex_fastlanes::bitpacked_v2_id; use vortex_fastlanes::for_v1_id; use vortex_fastlanes::for_v2_id; -use super::BITPACKING_V1; +use super::BitPackingScheme; use crate::ArrayAndStats; use crate::CascadingCompressor; use crate::CompressorContext; @@ -45,6 +46,18 @@ pub(crate) static FOR_V1: FoRScheme = FoRScheme::v1(); /// The FoR scheme that may produce one reference per 1024-element chunk. pub(crate) static FOR_V2: FoRScheme = FoRScheme::v2(); +/// [`FOR_V1`] with a child that may give each 1024-element block its own bit width. +static FOR_V1_BLOCKED: FoRScheme = FoRScheme { + mode: FoRSchemeMode::V1, + bitpacking: BitPackingScheme::v2(), +}; + +/// [`FOR_V2`] with a child that may give each 1024-element block its own bit width. +static FOR_V2_BLOCKED: FoRScheme = FoRScheme { + mode: FoRSchemeMode::V2, + bitpacking: BitPackingScheme::v2(), +}; + /// Frame of Reference encoding. /// /// The v1 mode subtracts one reference, the array's minimum. The v2 mode subtracts one reference @@ -52,11 +65,16 @@ pub(crate) static FOR_V2: FoRScheme = FoRScheme::v2(); /// arrays whose references compress to a constant, serialize as `fastlanes.for`, while per-chunk /// references serialize as `fastlanes.for.v2`. /// -/// The default uses v1. [`refine`](Scheme::refine) picks v2 when the v2 ID is allowed and v1 -/// otherwise. +/// The encoded child is always bit-packed by a [`BitPackingScheme`] called directly, not chosen by +/// the compressor, so this scheme declares the child's serialized IDs as its own. Excluding the +/// BitPacking scheme does not stop FoR from bit-packing; disallowing `fastlanes.bitpacked` does. +/// +/// The default uses v1 for both. [`refine`](Scheme::refine) picks v2 references when +/// `fastlanes.for.v2` is allowed, and a v2 child when `fastlanes.bitpacked.v2` is allowed. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct FoRScheme { mode: FoRSchemeMode, + bitpacking: BitPackingScheme, } impl FoRScheme { @@ -64,6 +82,7 @@ impl FoRScheme { pub const fn v1() -> Self { Self { mode: FoRSchemeMode::V1, + bitpacking: BitPackingScheme::v1(), } } @@ -71,6 +90,7 @@ impl FoRScheme { pub const fn v2() -> Self { Self { mode: FoRSchemeMode::V2, + bitpacking: BitPackingScheme::v1(), } } } @@ -91,17 +111,20 @@ impl Scheme for FoRScheme { } fn produced_encodings(&self) -> Vec { - match self.mode { + let mut encodings = match self.mode { FoRSchemeMode::V1 => vec![for_v1_id()], FoRSchemeMode::V2 => vec![for_v1_id(), for_v2_id()], - } + }; + encodings.extend(self.bitpacking.produced_encodings()); + encodings } fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme { - if allowed(&for_v2_id()) { - &FOR_V2 - } else { - &FOR_V1 + match (allowed(&for_v2_id()), allowed(&bitpacked_v2_id())) { + (false, false) => &FOR_V1, + (false, true) => &FOR_V1_BLOCKED, + (true, false) => &FOR_V2, + (true, true) => &FOR_V2_BLOCKED, } } @@ -214,10 +237,12 @@ impl Scheme for FoRScheme { // of bitpacking. // NOTE: we could delegate in the future if we had another downstream codec that performs // as well. - let leaf_ctx = compress_ctx.clone().as_leaf(); let biased_data = ArrayAndStats::new(biased.into_array(), compress_ctx.merged_stats_options()); - let compressed = BITPACKING_V1.compress(compressor, &biased_data, leaf_ctx, exec_ctx)?; + // Our own context, not a leaf one, so a v2 child's block offsets can cascade. + let compressed = + self.bitpacking + .compress(compressor, &biased_data, compress_ctx.clone(), exec_ctx)?; // TODO(connor): This should really be `new_unchecked`. let for_compressed = match for_array.references() { @@ -241,9 +266,11 @@ impl Scheme for FoRScheme { mod tests { use rstest::rstest; use vortex_array::ArrayId; + use vortex_fastlanes::bitpacked_v2_id; use vortex_fastlanes::for_v1_id; use vortex_fastlanes::for_v2_id; + use super::BitPackingScheme; use super::FOR_V1; use super::FOR_V2; use crate::Scheme; @@ -265,11 +292,26 @@ mod tests { |id: &ArrayId| (allow_v1 && *id == for_v1_id()) || (allow_v2 && *id == for_v2_id()); let refined = scheme.refine(&allowed); assert_eq!(refined.id(), scheme.id()); - let expected = if expect_v2 { + let mut expected = if expect_v2 { vec![for_v1_id(), for_v2_id()] } else { vec![for_v1_id()] }; + expected.extend(BitPackingScheme::v1().produced_encodings()); assert_eq!(refined.produced_encodings(), expected); } + + /// Both variants give the child per-block bit widths exactly when `fastlanes.bitpacked.v2` is + /// allowed. + #[rstest] + fn refine_picks_blocked_child_when_bitpacked_v2_is_allowed( + #[values(false, true)] allow_v2: bool, + #[values(&FOR_V1, &FOR_V2)] scheme: &'static dyn Scheme, + ) { + let refined = scheme.refine(&|id: &ArrayId| allow_v2 && *id == bitpacked_v2_id()); + assert_eq!( + refined.produced_encodings().contains(&bitpacked_v2_id()), + allow_v2 + ); + } } diff --git a/vortex-btrblocks/tests/for_config.rs b/vortex-btrblocks/tests/for_config.rs index 2c2bd76a063..59979457837 100644 --- a/vortex-btrblocks/tests/for_config.rs +++ b/vortex-btrblocks/tests/for_config.rs @@ -27,6 +27,7 @@ use vortex_edition::EditionSession; use vortex_edition::EditionSessionExt; use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexResult; +use vortex_fastlanes::bitpacked_v2_id; use vortex_fastlanes::for_v1_id; use vortex_fastlanes::for_v2_id; use vortex_session::VortexSession; @@ -79,6 +80,12 @@ fn clustered() -> ArrayRef { PrimitiveArray::from_iter((0..8192u32).map(|i| 1_000_000 + i % 100)).into_array() } +/// Values whose range doubles every chunk from the same base, so the child needs a different bit +/// width in each block. +fn widening() -> ArrayRef { + PrimitiveArray::from_iter((0..8192u32).map(|i| 1_000_000 + i % (4 << (i / 1024)))).into_array() +} + /// Compresses `array`, round trips it through serialization, and returns the serialized IDs. fn compress_roundtrip( builder: BtrBlocksCompressorBuilder, @@ -167,3 +174,17 @@ fn cuda_preset_keeps_single_reference() -> VortexResult<()> { assert!(!ids.contains(&for_v2_id())); Ok(()) } + +/// FoR bit-packs its child with per-block widths exactly when `fastlanes.bitpacked.v2` is allowed. +#[rstest] +#[case::all_allowed(for_only(), true)] +#[case::cuda(for_only().only_cuda_compatible(), false)] +fn child_widths_follow_bitpacked_v2( + #[case] builder: BtrBlocksCompressorBuilder, + #[case] expect_v2: bool, +) -> VortexResult<()> { + let ids = compress_roundtrip(builder, &widening())?; + assert!(ids.contains(&for_v1_id())); + assert_eq!(ids.contains(&bitpacked_v2_id()), expect_v2); + Ok(()) +} diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 13746255c03..f95e02288ad 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1598,12 +1598,11 @@ async fn test_into_tokio_array_stream() -> VortexResult<()> { #[tokio::test] async fn test_array_stream_no_double_dict_encode() -> VortexResult<()> { let num_vals = 2048; - let mut values = Vec::::with_capacity(num_vals); - // Far apart, so per-block bit widths can't pack them tighter than a dictionary. - values.extend(iter::repeat_n(1_000_000_000, num_vals / 2)); - values.extend(iter::repeat_n(2_000_000_000, num_vals / 2)); - - let array = PrimitiveArray::from_iter(values).into_array(); + // Alternating and far apart, so neither runs nor bit widths pack them tighter than a + // dictionary. + let array = + PrimitiveArray::from_iter((0..num_vals).map(|i| [1_000_000_000i64, 2_000_000_000][i % 2])) + .into_array(); let mut buf = Vec::new(); SESSION .write_options()