vortex_array/compute/
uncompressed_size.rs

1use vortex_error::{VortexExpect as _, VortexResult, vortex_panic};
2
3use crate::stats::{Precision, Stat};
4use crate::{Array, Encoding};
5
6pub trait UncompressedSizeFn<A> {
7    /// Compute the approximated uncompressed size of the array, in bytes.
8    fn uncompressed_size(&self, array: A) -> VortexResult<usize>;
9}
10
11impl<E: Encoding> UncompressedSizeFn<&dyn Array> for E
12where
13    E: for<'a> UncompressedSizeFn<&'a E::Array>,
14{
15    fn uncompressed_size(&self, array: &dyn Array) -> VortexResult<usize> {
16        let array_ref = array
17            .as_any()
18            .downcast_ref::<E::Array>()
19            .vortex_expect("Failed to downcast array");
20        UncompressedSizeFn::uncompressed_size(self, array_ref)
21    }
22}
23
24/// Computes the uncompressed size of the array, in bytes.
25pub fn uncompressed_size(array: &dyn Array) -> VortexResult<usize> {
26    if let Some(Precision::Exact(uncompressed_size)) = array
27        .statistics()
28        .get_as::<usize>(Stat::UncompressedSizeInBytes)
29    {
30        return Ok(uncompressed_size);
31    }
32
33    let uncompressed_size = match array.vtable().uncompressed_size_fn() {
34        Some(size_fn) => size_fn.uncompressed_size(array),
35        None => {
36            log::debug!(
37                "No uncompressed_size implementation found for {}",
38                array.encoding()
39            );
40            let array = array.to_canonical()?;
41            let array_ref = array.as_ref();
42            if let Some(size_fn) = array_ref.vtable().uncompressed_size_fn() {
43                size_fn.uncompressed_size(array_ref)
44            } else {
45                vortex_panic!(
46                    "No uncompressed_size function for canonical array: {}",
47                    array.as_ref().encoding(),
48                )
49            }
50        }
51    }?;
52
53    array.statistics().set(
54        Stat::UncompressedSizeInBytes,
55        Precision::Exact(uncompressed_size.into()),
56    );
57
58    Ok(uncompressed_size)
59}