vortex-tensor 0.71.0

Vortex tensor extension type
Documentation
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use vortex_array::ArrayRef;
use vortex_array::Canonical;
use vortex_array::ExecutionCtx;
use vortex_array::IntoArray;
use vortex_array::arrays::extension::ExtensionArrayExt;
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 {
    fn scheme_name(&self) -> &'static str {
        "vortex.tensor.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: &ArrayAndStats,
        _compress_ctx: CompressorContext,
        _exec_ctx: &mut ExecutionCtx,
    ) -> CompressionEstimate {
        CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse)
    }

    fn compress(
        &self,
        _compressor: &CascadingCompressor,
        data: &ArrayAndStats,
        _compress_ctx: CompressorContext,
        exec_ctx: &mut ExecutionCtx,
    ) -> VortexResult<ArrayRef> {
        let l2_denorm = normalize_as_l2_denorm(data.array().clone(), exec_ctx)?;
        Ok(l2_denorm.into_array())
    }
}