Use pcodec bounded size representation

LDeakin · Mar 21, 2024 · 89e20c9 · 89e20c9
1 parent 23a9b51
commit 89e20c9
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git a/src/array/codec/array_to_bytes/pcodec.rs b/src/array/codec/array_to_bytes/pcodec.rs
@@ -227,11 +227,13 @@ mod tests {
             ChunkRepresentation::new(chunk_shape, data_type, fill_value).unwrap();
         let bytes: Vec<u8> = (0..chunk_representation.size()).map(|s| s as u8).collect();
 
+        let max_encoded_size = codec.compute_encoded_size(&chunk_representation)?;
         let encoded = codec.encode(
             bytes.clone(),
             &chunk_representation,
             &CodecOptions::default(),
         )?;
+        assert!((encoded.len() as u64) <= max_encoded_size.size().unwrap());
         let decoded = codec
             .decode(encoded, &chunk_representation, &CodecOptions::default())
             .unwrap();

diff --git a/src/array/codec/array_to_bytes/pcodec/pcodec_codec.rs b/src/array/codec/array_to_bytes/pcodec/pcodec_codec.rs
@@ -1,4 +1,4 @@
-use pco::{ChunkConfig, FloatMultSpec, IntMultSpec, PagingSpec};
+use pco::{standalone::guarantee::file_size, ChunkConfig, FloatMultSpec, IntMultSpec, PagingSpec};
 
 use crate::{
     array::{
@@ -211,9 +211,28 @@ impl ArrayToBytesCodecTraits for PcodecCodec {
 
     fn compute_encoded_size(
         &self,
-        _decoded_representation: &ChunkRepresentation,
+        decoded_representation: &ChunkRepresentation,
     ) -> Result<BytesRepresentation, CodecError> {
-        // FIXME: pcodec is likely bounded, but it doesn't have a nice API to figure out what the bounded size is
-        Ok(BytesRepresentation::UnboundedSize)
+        let data_type = decoded_representation.data_type();
+        let mut num_elements = decoded_representation.num_elements_usize();
+        if data_type == &DataType::Complex64 || data_type == &DataType::Complex128 {
+            num_elements *= 2;
+        }
+
+        let size = match data_type {
+            DataType::UInt32 | DataType::Int32 | DataType::Float32 | DataType::Complex64 => Ok(
+                file_size::<u32>(num_elements, &self.chunk_config.paging_spec)
+                    .map_err(|err| CodecError::from(err.to_string()))?,
+            ),
+            DataType::UInt64 | DataType::Int64 | DataType::Float64 | DataType::Complex128 => Ok(
+                file_size::<u64>(num_elements, &self.chunk_config.paging_spec)
+                    .map_err(|err| CodecError::from(err.to_string()))?,
+            ),
+            _ => Err(CodecError::UnsupportedDataType(
+                data_type.clone(),
+                IDENTIFIER.to_string(),
+            )),
+        }?;
+        Ok(BytesRepresentation::BoundedSize(size.try_into().unwrap()))
     }
 }