L2 denorm scheme

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<ArrayId> = ALLOWED_ENCODINGS.clone();
+                allowed.insert(L2Denorm.id());
 
                 WriteStrategyBuilder::default()
                     .with_compressor(compressor)
+                    .with_allow_encodings(allowed)
                     .build()
             }
             VectorFlavor::TurboQuant => {

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<vortex_array::array::erased::ArrayRef>
+
+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

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::<AnyTensor>()
+        )
+    }
+
+    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<ArrayRef> {
+        let l2_denorm =
+            normalize_as_l2_denorm(data.array().clone(), &mut compressor.execution_ctx())?;
+        Ok(l2_denorm.into_array())
+    }
+}

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;