Skip to main content

zarrs/array/codec/array_to_array/bitround/
bitround_codec.rs

1use std::sync::Arc;
2
3use zarrs_plugin::{PluginCreateError, ZarrVersion};
4
5use super::{
6    BitroundCodecConfiguration, BitroundCodecConfigurationV1, BitroundDataTypeExt,
7    bitround_codec_partial, round_bytes,
8};
9use crate::array::{DataType, FillValue};
10use std::num::NonZeroU64;
11use zarrs_codec::{
12    ArrayBytes, ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayPartialEncoderTraits,
13    ArrayToArrayCodecTraits, CodecError, CodecMetadataOptions, CodecOptions, CodecTraits,
14    PartialDecoderCapability, PartialEncoderCapability, RecommendedConcurrency,
15};
16#[cfg(feature = "async")]
17use zarrs_codec::{AsyncArrayPartialDecoderTraits, AsyncArrayPartialEncoderTraits};
18use zarrs_metadata::Configuration;
19
20/// A `bitround` codec implementation.
21#[derive(Clone, Debug, Default)]
22pub struct BitroundCodec {
23    keepbits: u32,
24}
25
26impl BitroundCodec {
27    /// Create a new `bitround` codec.
28    ///
29    /// `keepbits` is the number of bits to round to in the floating point mantissa.
30    #[must_use]
31    pub const fn new(keepbits: u32) -> Self {
32        Self { keepbits }
33    }
34
35    /// Create a new `bitround` codec from a configuration.
36    ///
37    /// # Errors
38    /// Returns an error if the configuration is not supported.
39    pub fn new_with_configuration(
40        configuration: &BitroundCodecConfiguration,
41    ) -> Result<Self, PluginCreateError> {
42        match configuration {
43            BitroundCodecConfiguration::V1(configuration) => Ok(Self {
44                keepbits: configuration.keepbits,
45            }),
46            _ => Err(PluginCreateError::Other(
47                "this bitround codec configuration variant is unsupported".to_string(),
48            )),
49        }
50    }
51}
52
53impl CodecTraits for BitroundCodec {
54    fn as_any(&self) -> &dyn std::any::Any {
55        self
56    }
57
58    fn configuration(
59        &self,
60        _version: ZarrVersion,
61        options: &CodecMetadataOptions,
62    ) -> Option<Configuration> {
63        if options.codec_store_metadata_if_encode_only() {
64            let configuration = BitroundCodecConfiguration::V1(BitroundCodecConfigurationV1 {
65                keepbits: self.keepbits,
66            });
67            Some(configuration.into())
68        } else {
69            None
70        }
71    }
72
73    fn partial_decoder_capability(&self) -> PartialDecoderCapability {
74        PartialDecoderCapability {
75            partial_read: true,
76            partial_decode: true,
77        }
78    }
79
80    fn partial_encoder_capability(&self) -> PartialEncoderCapability {
81        PartialEncoderCapability {
82            partial_encode: true,
83        }
84    }
85}
86
87impl ArrayCodecTraits for BitroundCodec {
88    fn recommended_concurrency(
89        &self,
90        _shape: &[NonZeroU64],
91        _data_type: &DataType,
92    ) -> Result<RecommendedConcurrency, CodecError> {
93        // TODO: bitround is well suited to multithread, when is it optimal to kick in?
94        Ok(RecommendedConcurrency::new_maximum(1))
95    }
96}
97
98#[cfg_attr(
99    all(feature = "async", not(target_arch = "wasm32")),
100    async_trait::async_trait
101)]
102#[cfg_attr(all(feature = "async", target_arch = "wasm32"), async_trait::async_trait(?Send))]
103impl ArrayToArrayCodecTraits for BitroundCodec {
104    fn into_dyn(self: Arc<Self>) -> Arc<dyn ArrayToArrayCodecTraits> {
105        self as Arc<dyn ArrayToArrayCodecTraits>
106    }
107
108    fn encode<'a>(
109        &self,
110        bytes: ArrayBytes<'a>,
111        _shape: &[NonZeroU64],
112        data_type: &DataType,
113        _fill_value: &FillValue,
114        _options: &CodecOptions,
115    ) -> Result<ArrayBytes<'a>, CodecError> {
116        let mut bytes = bytes.into_fixed()?;
117        round_bytes(bytes.to_mut(), data_type, self.keepbits)?;
118        Ok(ArrayBytes::from(bytes))
119    }
120
121    fn decode<'a>(
122        &self,
123        bytes: ArrayBytes<'a>,
124        _shape: &[NonZeroU64],
125        _data_type: &DataType,
126        _fill_value: &FillValue,
127        _options: &CodecOptions,
128    ) -> Result<ArrayBytes<'a>, CodecError> {
129        Ok(bytes)
130    }
131
132    fn partial_decoder(
133        self: Arc<Self>,
134        input_handle: Arc<dyn ArrayPartialDecoderTraits>,
135        _shape: &[NonZeroU64],
136        data_type: &DataType,
137        _fill_value: &FillValue,
138        _options: &CodecOptions,
139    ) -> Result<Arc<dyn ArrayPartialDecoderTraits>, CodecError> {
140        Ok(Arc::new(bitround_codec_partial::BitroundCodecPartial::new(
141            input_handle,
142            data_type,
143            self.keepbits,
144        )?))
145    }
146
147    fn partial_encoder(
148        self: Arc<Self>,
149        input_output_handle: Arc<dyn ArrayPartialEncoderTraits>,
150        _shape: &[NonZeroU64],
151        data_type: &DataType,
152        _fill_value: &FillValue,
153        _options: &CodecOptions,
154    ) -> Result<Arc<dyn ArrayPartialEncoderTraits>, CodecError> {
155        Ok(Arc::new(bitround_codec_partial::BitroundCodecPartial::new(
156            input_output_handle,
157            data_type,
158            self.keepbits,
159        )?))
160    }
161
162    #[cfg(feature = "async")]
163    async fn async_partial_decoder(
164        self: Arc<Self>,
165        input_handle: Arc<dyn AsyncArrayPartialDecoderTraits>,
166        _shape: &[NonZeroU64],
167        data_type: &DataType,
168        _fill_value: &FillValue,
169        _options: &CodecOptions,
170    ) -> Result<Arc<dyn AsyncArrayPartialDecoderTraits>, CodecError> {
171        Ok(Arc::new(bitround_codec_partial::BitroundCodecPartial::new(
172            input_handle,
173            data_type,
174            self.keepbits,
175        )?))
176    }
177
178    #[cfg(feature = "async")]
179    async fn async_partial_encoder(
180        self: Arc<Self>,
181        input_output_handle: Arc<dyn AsyncArrayPartialEncoderTraits>,
182        _shape: &[NonZeroU64],
183        data_type: &DataType,
184        _fill_value: &FillValue,
185        _options: &CodecOptions,
186    ) -> Result<Arc<dyn AsyncArrayPartialEncoderTraits>, CodecError> {
187        Ok(Arc::new(bitround_codec_partial::BitroundCodecPartial::new(
188            input_output_handle,
189            data_type,
190            self.keepbits,
191        )?))
192    }
193
194    fn encoded_data_type(&self, decoded_data_type: &DataType) -> Result<DataType, CodecError> {
195        decoded_data_type.codec_bitround()?;
196        Ok(decoded_data_type.clone())
197    }
198}