Add a new AggregateFn for UncompressedSize stat

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-array/public-api.lock

@@ -728,6 +728,56 @@ pub struct vortex_array::aggregate_fn::fns::sum::SumPartial
 
 pub fn vortex_array::aggregate_fn::fns::sum::sum(array: &vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
 
+pub mod vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes
+
+pub struct vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
+
+impl core::clone::Clone for vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::clone(&self) -> vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
+
+impl core::fmt::Debug for vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
+
+pub type vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::Options = vortex_array::aggregate_fn::EmptyOptions
+
+pub type vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::Partial = u64
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::accumulate(&self, partial: &mut Self::Partial, batch: &vortex_array::Columnar, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::coerce_args(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::combine_partials(&self, partial: &mut Self::Partial, other: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::empty_partial(&self, _options: &Self::Options, _input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::finalize(&self, partials: vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::finalize_scalar(&self, partial: &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::is_saturated(&self, _partial: &Self::Partial) -> bool
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::partial_dtype(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::reset(&self, partial: &mut Self::Partial)
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::return_dtype(&self, _options: &Self::Options, _input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::serialize(&self, _options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::to_scalar(&self, partial: &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::try_accumulate(&self, partial: &mut Self::Partial, batch: &vortex_array::ArrayRef, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes(array: &vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<usize>
+
 pub mod vortex_array::aggregate_fn::kernels
 
 pub trait vortex_array::aggregate_fn::kernels::DynAggregateKernel: 'static + core::marker::Send + core::marker::Sync + core::fmt::Debug
@@ -1264,6 +1314,42 @@ pub fn vortex_array::aggregate_fn::fns::sum::Sum::to_scalar(&self, partial: &Sel
 
 pub fn vortex_array::aggregate_fn::fns::sum::Sum::try_accumulate(&self, _state: &mut Self::Partial, _batch: &vortex_array::ArrayRef, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
 
+impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
+
+pub type vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::Options = vortex_array::aggregate_fn::EmptyOptions
+
+pub type vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::Partial = u64
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::accumulate(&self, partial: &mut Self::Partial, batch: &vortex_array::Columnar, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::coerce_args(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::combine_partials(&self, partial: &mut Self::Partial, other: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::empty_partial(&self, _options: &Self::Options, _input_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::finalize(&self, partials: vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::finalize_scalar(&self, partial: &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::is_saturated(&self, _partial: &Self::Partial) -> bool
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::partial_dtype(&self, options: &Self::Options, input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::reset(&self, partial: &mut Self::Partial)
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::return_dtype(&self, _options: &Self::Options, _input_dtype: &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::serialize(&self, _options: &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::to_scalar(&self, partial: &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+pub fn vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes::try_accumulate(&self, partial: &mut Self::Partial, batch: &vortex_array::ArrayRef, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
+
 impl<T: vortex_array::aggregate_fn::combined::BinaryCombined> vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::combined::Combined<T>
 
 pub type vortex_array::aggregate_fn::combined::Combined<T>::Options = vortex_array::aggregate_fn::combined::PairOptions<<<T as vortex_array::aggregate_fn::combined::BinaryCombined>::Left as vortex_array::aggregate_fn::AggregateFnVTable>::Options, <<T as vortex_array::aggregate_fn::combined::BinaryCombined>::Right as vortex_array::aggregate_fn::AggregateFnVTable>::Options>
@@ -17604,7 +17690,7 @@ pub fn vortex_array::scalar_fn::internal::row_count::RowCount::child_name(&self,
 
 pub fn vortex_array::scalar_fn::internal::row_count::RowCount::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_array::dtype::DType>>
 
-pub fn vortex_array::scalar_fn::internal::row_count::RowCount::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::deserialize(&self, metadata: &[u8], session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
 
 pub fn vortex_array::scalar_fn::internal::row_count::RowCount::execute(&self, _options: &Self::Options, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
@@ -18800,7 +18886,7 @@ pub fn vortex_array::scalar_fn::internal::row_count::RowCount::child_name(&self,
 
 pub fn vortex_array::scalar_fn::internal::row_count::RowCount::coerce_args(&self, options: &Self::Options, args: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<alloc::vec::Vec<vortex_array::dtype::DType>>
 
-pub fn vortex_array::scalar_fn::internal::row_count::RowCount::deserialize(&self, _metadata: &[u8], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
+pub fn vortex_array::scalar_fn::internal::row_count::RowCount::deserialize(&self, metadata: &[u8], session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
 
 pub fn vortex_array::scalar_fn::internal::row_count::RowCount::execute(&self, _options: &Self::Options, _args: &dyn vortex_array::scalar_fn::ExecutionArgs, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 

vortex-array/src/aggregate_fn/fns/mod.rs

@@ -10,3 +10,4 @@ pub mod mean;
 pub mod min_max;
 pub mod nan_count;
 pub mod sum;
+pub mod uncompressed_size_in_bytes;

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/bool.rs

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::ExecutionCtx;
+use crate::arrays::BoolArray;
+
+pub(super) fn bool_uncompressed_size_in_bytes(
+    array: &BoolArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    let value_size = super::packed_bit_buffer_size_in_bytes(array.len())?;
+    let validity_size = super::validity_uncompressed_size_in_bytes(
+        array
+            .as_ref()
+            .validity()?
+            .execute_mask(array.as_ref().len(), ctx)?,
+    )?;
+
+    value_size
+        .checked_add(validity_size)
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))
+}

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/decimal.rs

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::ExecutionCtx;
+use crate::arrays::DecimalArray;
+use crate::arrays::decimal::DecimalArrayExt;
+use crate::dtype::DecimalType;
+
+pub(super) fn decimal_uncompressed_size_in_bytes(
+    array: &DecimalArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    let value_size = u64::try_from(array.len())
+        .map_err(|e| vortex_err!("Failed to convert decimal array length to u64: {e}"))?
+        .checked_mul(
+            u64::try_from(
+                DecimalType::smallest_decimal_value_type(&array.decimal_dtype()).byte_width(),
+            )
+            .map_err(|e| vortex_err!("Failed to convert decimal byte width to u64: {e}"))?,
+        )
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
+    let validity_size = super::validity_uncompressed_size_in_bytes(
+        array
+            .as_ref()
+            .validity()?
+            .execute_mask(array.as_ref().len(), ctx)?,
+    )?;
+
+    value_size
+        .checked_add(validity_size)
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))
+}

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/extension.rs

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::ExecutionCtx;
+use crate::arrays::ExtensionArray;
+use crate::arrays::extension::ExtensionArrayExt;
+
+pub(super) fn extension_uncompressed_size_in_bytes(
+    array: &ExtensionArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    super::uncompressed_size_in_bytes_u64(array.storage_array(), ctx)
+}

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/fixed_size_list.rs

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::ExecutionCtx;
+use crate::arrays::FixedSizeListArray;
+use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
+
+pub(super) fn fixed_size_list_uncompressed_size_in_bytes(
+    array: &FixedSizeListArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    let elements_size = super::uncompressed_size_in_bytes_u64(array.elements(), ctx)?;
+    let validity_size = super::validity_uncompressed_size_in_bytes(
+        array
+            .as_ref()
+            .validity()?
+            .execute_mask(array.as_ref().len(), ctx)?,
+    )?;
+
+    elements_size
+        .checked_add(validity_size)
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))
+}

vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list.rs

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::ExecutionCtx;
+use crate::arrays::ListViewArray;
+use crate::arrays::listview::ListViewArrayExt;
+use crate::arrays::listview::ListViewRebuildMode;
+
+pub(super) fn list_uncompressed_size_in_bytes(
+    array: &ListViewArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<u64> {
+    let mut size = if array.is_empty() {
+        0
+    } else {
+        let rebuilt = array.rebuild(ListViewRebuildMode::MakeExact)?;
+        super::uncompressed_size_in_bytes_u64(rebuilt.elements(), ctx)?
+    };
+
+    let view_buffer_size = u64::try_from(array.len())
+        .map_err(|e| vortex_err!("Failed to convert list array length to u64: {e}"))?
+        .checked_mul(8)
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
+
+    size = size
+        .checked_add(view_buffer_size)
+        .and_then(|size| size.checked_add(view_buffer_size))
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
+    size = size
+        .checked_add(super::validity_uncompressed_size_in_bytes(
+            array
+                .as_ref()
+                .validity()?
+                .execute_mask(array.as_ref().len(), ctx)?,
+        )?)
+        .ok_or_else(|| vortex_err!("uncompressed size in bytes overflowed u64"))?;
+
+    Ok(size)
+}