From 96b2fe91566b5cd1590cefaf8ebc328b7b562f8c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 18 Aug 2026 12:38:24 +0100 Subject: [PATCH] Enable OnPair in GPU compression benchmark Signed-off-by: "Robert Kruszewski" --- vortex-btrblocks/src/builder.rs | 13 +++++++--- vortex-btrblocks/src/schemes/string/onpair.rs | 5 +++- vortex-compressor/src/compressor/mod.rs | 5 ++++ vortex-cuda/src/kernel/encodings/onpair.rs | 24 ++++++++++++++++++- 4 files changed, 42 insertions(+), 5 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6f38e29cd86..a5e4c714781 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -163,8 +163,8 @@ impl BtrBlocksCompressorBuilder { /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. pub fn only_cuda_compatible(self) -> Self { - // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other - // string fragmentation and dictionary schemes still require unsupported decode paths. + // Keep FSST and OnPair, which have CUDA decoders and direct Arrow export. String + // dictionary schemes still require unsupported decode paths. #[cfg_attr( not(any(feature = "pco", feature = "unstable_encodings")), allow(unused_mut) @@ -269,7 +269,7 @@ mod tests { } #[test] - fn cuda_compatible_uses_fsst_for_strings() { + fn cuda_compatible_uses_gpu_string_encodings() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); assert!( builder @@ -277,6 +277,13 @@ mod tests { .iter() .any(|scheme| scheme.id() == string::FSSTScheme.id()) ); + #[cfg(feature = "unstable_encodings")] + assert!( + builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::OnPairScheme.id()) + ); #[cfg(feature = "zstd")] assert!( !builder diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 75ec97d2191..2d1d6cccaa4 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -26,6 +26,7 @@ use crate::CascadingCompressor; use crate::CompressorContext; use crate::Scheme; use crate::SchemeExt; +use crate::schemes::integer::DeltaScheme; use crate::schemes::integer::try_compress_delta; /// OnPair short-string compression (dict-12). @@ -174,7 +175,9 @@ fn compress_offsets_child( .into_array(); let plain = compressor.compress_child(&narrowed, compress_ctx, scheme_id, child_idx, exec_ctx)?; - if narrowed.len() < OFFSETS_DELTA_MIN_LEN { + if narrowed.len() < OFFSETS_DELTA_MIN_LEN + || !compressor.is_scheme_enabled(DeltaScheme::default().id()) + { return Ok(plain); } let delta = try_compress_delta( diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..01db348287a 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -65,6 +65,11 @@ impl CascadingCompressor { root_exclusions, } } + + /// Returns whether a compression scheme is enabled for this compressor. + pub fn is_scheme_enabled(&self, scheme_id: SchemeId) -> bool { + self.schemes.iter().any(|scheme| scheme.id() == scheme_id) + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index ee132e2bb8d..fe0f4d9de30 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -48,6 +48,7 @@ use num_traits::AsPrimitive; use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; +use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::primitive::PrimitiveDataParts; @@ -72,6 +73,7 @@ use vortex_onpair::OnPair; use vortex_onpair::OnPairArray; use vortex_onpair::OnPairArrayExt; use vortex_onpair::OnPairArraySlotsExt; +use vortex_onpair::OnPairData; use vortex_onpair::dict_view; use crate::CanonicalCudaExt; @@ -324,7 +326,27 @@ async fn stage_dict( onpair: ArrayView<'_, OnPair>, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let dict = dict_view(onpair, ctx.execution_ctx())?; + let dict_offsets = decode_primitive_child(onpair.dict_offsets().clone(), ctx).await?; + let dict_offsets = Canonical::Primitive(dict_offsets) + .into_host() + .await? + .into_array(); + let dict_bytes = onpair + .data() + .dict_bytes_handle() + .clone() + .try_to_host()? + .await?; + let host_onpair = OnPair::try_new_with_data( + onpair.dtype().clone(), + OnPairData::new(BufferHandle::new_host(dict_bytes)), + dict_offsets, + onpair.codes().clone(), + onpair.codes_offsets().clone(), + onpair.uncompressed_lengths().clone(), + onpair.array_validity(), + )?; + let dict = dict_view(host_onpair.as_view(), ctx.execution_ctx())?; let dict_size = dict.num_tokens(); let dict_size_u32 = u32::try_from(dict_size)?; let mut dict_padded = vec![0u8; dict_size * MAX_TOKEN_SIZE];