GeezerVison is showing github.com — vortex/vortex-python/src/compress.rs at develop · massive-com/vortex

Expand file tree

compress.rs

Copy path

More file actions

Latest commit

History

62 lines (56 loc) · 1.71 KB

compress.rs

Copy path

File metadata and controls

62 lines (56 loc) · 1.71 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: Copyright the Vortex contributors

use pyo3::prelude::*;

use vortex::array::VortexSessionExecute;

use vortex::compressor::BtrBlocksCompressor;

use crate::arrays::PyArrayRef;

use crate::error::PyVortexResult;

use crate::install_module;

use crate::session::session;

pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {

let m = PyModule::new(py, "compress")?;

parent.add_submodule(&m)?;

install_module("vortex._lib.compress", &m)?;

m.add_function(wrap_pyfunction!(compress, &m)?)?;

Ok(())

}

/// Attempt to compress a vortex array.

///

/// Parameters

/// ----------

/// array : :class:`~vortex.Array`

/// The array.

///

/// Examples

/// --------

///

/// Compress a very sparse array of integers:

///

/// >>> import vortex as vx

/// >>> a = vx.array([42 for _ in range(1000)])

/// >>> str(vx.compress(a))

/// 'vortex.constant(i64, len=1000)'

///

/// Compress an array of increasing integers:

///

/// >>> a = vx.array(list(range(1000)))

/// >>> str(vx.compress(a))

/// 'vortex.sequence(i64, len=1000)'

///

/// Compress an array of increasing floating-point numbers and a few nulls:

///

/// >>> a = vx.array([

/// ... float(x) if x % 20 != 0 else None

/// ... for x in range(1000)

/// ... ])

/// >>> str(vx.compress(a))

/// 'vortex.alp(f64?, len=1000)'

#[pyfunction]

pub fn compress(py: Python, array: PyArrayRef) -> PyVortexResult<PyArrayRef> {

let session = session();

let array = array.into_inner();

let compressed = py.detach(move || {

BtrBlocksCompressor::default().compress(&array, &mut session.create_execution_ctx())

})?;

Ok(PyArrayRef::from(compressed))

}