use crate::array::{ArrayBytes, ArraySubset, ChunkShape, DataType};
#[cfg(feature = "async")]
use zarrs_codec::AsyncArrayPartialDecoderTraits;
use zarrs_codec::{ArrayPartialDecoderTraits, CodecError, CodecOptions};
use zarrs_storage::StorageError;
pub(crate) struct ArrayPartialDecoderCache {
shape: ChunkShape,
data_type: DataType,
cache: ArrayBytes<'static>,
}
impl ArrayPartialDecoderCache {
pub(crate) fn new(
input_handle: &dyn ArrayPartialDecoderTraits,
shape: ChunkShape,
data_type: DataType,
options: &CodecOptions,
) -> Result<Self, CodecError> {
let bytes = input_handle
.partial_decode(
&ArraySubset::new_with_shape(bytemuck::must_cast_slice(&shape).to_vec()),
options,
)?
.into_owned();
Ok(Self {
shape,
data_type,
cache: bytes,
})
}
#[cfg(feature = "async")]
pub(crate) async fn async_new(
input_handle: &dyn AsyncArrayPartialDecoderTraits,
shape: ChunkShape,
data_type: DataType,
options: &CodecOptions,
) -> Result<ArrayPartialDecoderCache, CodecError> {
let bytes = input_handle
.partial_decode(
&ArraySubset::new_with_shape(bytemuck::must_cast_slice(&shape).to_vec()),
options,
)
.await?
.into_owned();
Ok(Self {
shape,
data_type,
cache: bytes,
})
}
}
impl ArrayPartialDecoderTraits for ArrayPartialDecoderCache {
fn exists(&self) -> Result<bool, StorageError> {
Ok(true)
}
fn size_held(&self) -> usize {
self.cache.size()
}
fn data_type(&self) -> &DataType {
&self.data_type
}
fn partial_decode(
&self,
indexer: &dyn crate::array::Indexer,
_options: &CodecOptions,
) -> Result<ArrayBytes<'_>, CodecError> {
let array_shape = bytemuck::must_cast_slice(&self.shape);
self.cache
.extract_array_subset(indexer, array_shape, &self.data_type)
}
fn supports_partial_decode(&self) -> bool {
true
}
}
#[cfg(feature = "async")]
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
impl AsyncArrayPartialDecoderTraits for ArrayPartialDecoderCache {
fn data_type(&self) -> &DataType {
&self.data_type
}
async fn exists(&self) -> Result<bool, StorageError> {
Ok(true)
}
fn size_held(&self) -> usize {
self.cache.size()
}
async fn partial_decode<'a>(
&'a self,
indexer: &dyn crate::array::Indexer,
options: &CodecOptions,
) -> Result<ArrayBytes<'a>, CodecError> {
ArrayPartialDecoderTraits::partial_decode(self, indexer, options)
}
fn supports_partial_decode(&self) -> bool {
true
}
}