diff --git a/benchmarks/vector-search-bench/src/compression.rs b/benchmarks/vector-search-bench/src/compression.rs index 6e2733cc1cd..65bf080db85 100644 --- a/benchmarks/vector-search-bench/src/compression.rs +++ b/benchmarks/vector-search-bench/src/compression.rs @@ -22,6 +22,7 @@ use vortex::session::VortexSession; use vortex::utils::aliases::hash_set::HashSet; use vortex_bench::Format; use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_tensor::encodings::l2_denorm::L2DenormScheme; use vortex_tensor::scalar_fns::l2_denorm::L2Denorm; use vortex_tensor::scalar_fns::sorf_transform::SorfTransform; @@ -76,10 +77,18 @@ impl VectorFlavor { pub fn create_write_options(&self, session: &VortexSession) -> VortexWriteOptions { let strategy = match self { VectorFlavor::Uncompressed => { - let compressor = BtrBlocksCompressorBuilder::empty().build(); + // Even though this is uncompressed, we still want to denormalize the data first so + // that the results are fair. + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&L2DenormScheme) + .build(); + + let mut allowed: HashSet = ALLOWED_ENCODINGS.clone(); + allowed.insert(L2Denorm.id()); WriteStrategyBuilder::default() .with_compressor(compressor) + .with_allow_encodings(allowed) .build() } VectorFlavor::TurboQuant => { diff --git a/vortex-tensor/public-api.lock b/vortex-tensor/public-api.lock index 99dc8f0edbd..96b95e1e91e 100644 --- a/vortex-tensor/public-api.lock +++ b/vortex-tensor/public-api.lock @@ -2,6 +2,24 @@ pub mod vortex_tensor pub mod vortex_tensor::encodings +pub mod vortex_tensor::encodings::l2_denorm + +pub struct vortex_tensor::encodings::l2_denorm::L2DenormScheme + +impl core::fmt::Debug for vortex_tensor::encodings::l2_denorm::L2DenormScheme + +pub fn vortex_tensor::encodings::l2_denorm::L2DenormScheme::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result + +impl vortex_compressor::scheme::Scheme for vortex_tensor::encodings::l2_denorm::L2DenormScheme + +pub fn vortex_tensor::encodings::l2_denorm::L2DenormScheme::compress(&self, compressor: &vortex_compressor::compressor::CascadingCompressor, data: &mut vortex_compressor::stats::cache::ArrayAndStats, _ctx: vortex_compressor::ctx::CompressorContext) -> vortex_error::VortexResult + +pub fn vortex_tensor::encodings::l2_denorm::L2DenormScheme::expected_compression_ratio(&self, _data: &mut vortex_compressor::stats::cache::ArrayAndStats, _ctx: vortex_compressor::ctx::CompressorContext) -> vortex_compressor::estimate::CompressionEstimate + +pub fn vortex_tensor::encodings::l2_denorm::L2DenormScheme::matches(&self, canonical: &vortex_array::canonical::Canonical) -> bool + +pub fn vortex_tensor::encodings::l2_denorm::L2DenormScheme::scheme_name(&self) -> &'static str + pub mod vortex_tensor::encodings::turboquant pub struct vortex_tensor::encodings::turboquant::TurboQuantConfig diff --git a/vortex-tensor/src/encodings/l2_denorm.rs b/vortex-tensor/src/encodings/l2_denorm.rs new file mode 100644 index 00000000000..172191abf6e --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm.rs @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_compressor::CascadingCompressor; +use vortex_compressor::ctx::CompressorContext; +use vortex_compressor::estimate::CompressionEstimate; +use vortex_compressor::estimate::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; +use vortex_error::VortexResult; + +use crate::matcher::AnyTensor; +use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm; + +#[derive(Debug)] +pub struct L2DenormScheme; + +impl Scheme for L2DenormScheme { + // TODO(connor): FIX THIS!!! + fn scheme_name(&self) -> &'static str { + "vortex.tensor.UNSTABLE.l2_denorm" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches!( + canonical, + Canonical::Extension(ext) if ext.ext_dtype().is::() + ) + } + + fn expected_compression_ratio( + &self, + _data: &mut ArrayAndStats, + _ctx: CompressorContext, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &mut ArrayAndStats, + _ctx: CompressorContext, + ) -> VortexResult { + let l2_denorm = + normalize_as_l2_denorm(data.array().clone(), &mut compressor.execution_ctx())?; + Ok(l2_denorm.into_array()) + } +} diff --git a/vortex-tensor/src/encodings/mod.rs b/vortex-tensor/src/encodings/mod.rs index 084baf97e57..22e57763171 100644 --- a/vortex-tensor/src/encodings/mod.rs +++ b/vortex-tensor/src/encodings/mod.rs @@ -6,4 +6,5 @@ // TODO(connor): // pub mod spherical; // Spherical transform on unit-normalized vectors. +pub mod l2_denorm; pub mod turboquant;