1#![allow(clippy::multiple_crate_versions)] use std::{borrow::Cow, fmt};
32
33use ndarray::{Array, Array1, ArrayView, Dimension, Zip};
34use numcodecs::{
35 AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
36 Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
37};
38use schemars::JsonSchema;
39use serde::{Deserialize, Serialize};
40use thiserror::Error;
41
42#[cfg(test)]
43use ::serde_json as _;
44
45mod ffi;
46
47type ZfpClassicCodecVersion = StaticCodecVersion<0, 2, 0>;
48
49#[derive(Clone, Serialize, Deserialize, JsonSchema)]
50#[schemars(deny_unknown_fields)]
52pub struct ZfpClassicCodec {
54 #[serde(flatten)]
56 pub mode: ZfpCompressionMode,
57 #[serde(default)]
59 pub non_finite: ZfpNonFiniteValuesMode,
60 #[serde(default, rename = "_version")]
62 pub version: ZfpClassicCodecVersion,
63}
64
65#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]
66#[serde(tag = "mode")]
67#[serde(deny_unknown_fields)]
68pub enum ZfpCompressionMode {
70 #[serde(rename = "expert")]
71 Expert {
73 min_bits: u32,
75 max_bits: u32,
77 max_prec: u32,
79 min_exp: i32,
84 },
85 #[serde(rename = "fixed-rate")]
90 FixedRate {
91 rate: f64,
93 },
94 #[serde(rename = "fixed-precision")]
98 FixedPrecision {
99 precision: u32,
101 },
102 #[serde(rename = "fixed-accuracy")]
107 FixedAccuracy {
108 tolerance: f64,
110 },
111 #[serde(rename = "reversible")]
114 Reversible,
115}
116
117#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, JsonSchema)]
118pub enum ZfpNonFiniteValuesMode {
120 #[default]
122 #[serde(rename = "deny")]
123 Deny,
124 #[serde(rename = "allow-unsafe")]
128 AllowUnsafe,
129}
130
131impl Codec for ZfpClassicCodec {
132 type Error = ZfpClassicCodecError;
133
134 fn encode(&self, data: AnyCowArray) -> Result<AnyArray, Self::Error> {
135 if matches!(data.dtype(), AnyArrayDType::I32 | AnyArrayDType::I64)
136 && matches!(
137 self.mode,
138 ZfpCompressionMode::FixedAccuracy { tolerance: _ }
139 )
140 {
141 return Err(ZfpClassicCodecError::FixedAccuracyModeIntegerData);
142 }
143
144 match data {
145 AnyCowArray::I32(data) => Ok(AnyArray::U8(
146 Array1::from(compress(data.view(), &self.mode, self.non_finite)?).into_dyn(),
147 )),
148 AnyCowArray::I64(data) => Ok(AnyArray::U8(
149 Array1::from(compress(data.view(), &self.mode, self.non_finite)?).into_dyn(),
150 )),
151 AnyCowArray::F32(data) => Ok(AnyArray::U8(
152 Array1::from(compress(data.view(), &self.mode, self.non_finite)?).into_dyn(),
153 )),
154 AnyCowArray::F64(data) => Ok(AnyArray::U8(
155 Array1::from(compress(data.view(), &self.mode, self.non_finite)?).into_dyn(),
156 )),
157 encoded => Err(ZfpClassicCodecError::UnsupportedDtype(encoded.dtype())),
158 }
159 }
160
161 fn decode(&self, encoded: AnyCowArray) -> Result<AnyArray, Self::Error> {
162 let AnyCowArray::U8(encoded) = encoded else {
163 return Err(ZfpClassicCodecError::EncodedDataNotBytes {
164 dtype: encoded.dtype(),
165 });
166 };
167
168 if !matches!(encoded.shape(), [_]) {
169 return Err(ZfpClassicCodecError::EncodedDataNotOneDimensional {
170 shape: encoded.shape().to_vec(),
171 });
172 }
173
174 decompress(&AnyCowArray::U8(encoded).as_bytes())
175 }
176
177 fn decode_into(
178 &self,
179 encoded: AnyArrayView,
180 decoded: AnyArrayViewMut,
181 ) -> Result<(), Self::Error> {
182 let AnyArrayView::U8(encoded) = encoded else {
183 return Err(ZfpClassicCodecError::EncodedDataNotBytes {
184 dtype: encoded.dtype(),
185 });
186 };
187
188 if !matches!(encoded.shape(), [_]) {
189 return Err(ZfpClassicCodecError::EncodedDataNotOneDimensional {
190 shape: encoded.shape().to_vec(),
191 });
192 }
193
194 decompress_into(&AnyArrayView::U8(encoded).as_bytes(), decoded)
195 }
196}
197
198impl StaticCodec for ZfpClassicCodec {
199 const CODEC_ID: &'static str = "zfp-classic.rs";
200
201 type Config<'de> = Self;
202
203 fn from_config(config: Self::Config<'_>) -> Self {
204 config
205 }
206
207 fn get_config(&self) -> StaticCodecConfig<'_, Self> {
208 StaticCodecConfig::from(self)
209 }
210}
211
212#[derive(Debug, Error)]
213pub enum ZfpClassicCodecError {
215 #[error("ZfpClassic does not support the dtype {0}")]
217 UnsupportedDtype(AnyArrayDType),
218 #[error("ZfpClassic does not support the fixed accuracy mode for integer data")]
221 FixedAccuracyModeIntegerData,
222 #[error("ZfpClassic only supports 1-4 dimensional data but found shape {shape:?}")]
224 ExcessiveDimensionality {
225 shape: Vec<usize>,
227 },
228 #[error("ZfpClassic was configured with an invalid expert mode {mode:?}")]
230 InvalidExpertMode {
231 mode: ZfpCompressionMode,
233 },
234 #[error(
237 "Zfp does not support non-finite (infinite or NaN) floating point data in non-reversible lossy compression"
238 )]
239 NonFiniteData,
240 #[error("ZfpClassic failed to encode the header")]
242 HeaderEncodeFailed,
243 #[error("ZfpClassic failed to encode the array metadata header")]
245 MetaHeaderEncodeFailed {
246 source: ZfpHeaderError,
248 },
249 #[error("ZfpClassic failed to encode the data")]
251 ZfpEncodeFailed,
252 #[error(
255 "ZfpClassic can only decode one-dimensional byte arrays but received an array of dtype {dtype}"
256 )]
257 EncodedDataNotBytes {
258 dtype: AnyArrayDType,
260 },
261 #[error(
264 "ZfpClassic can only decode one-dimensional byte arrays but received a byte array of shape {shape:?}"
265 )]
266 EncodedDataNotOneDimensional {
267 shape: Vec<usize>,
269 },
270 #[error("ZfpClassic failed to decode the header")]
272 HeaderDecodeFailed,
273 #[error("ZfpClassic failed to decode the array metadata header")]
275 MetaHeaderDecodeFailed {
276 source: ZfpHeaderError,
278 },
279 #[error("ZfpClassicCodec cannot decode into the provided array")]
281 MismatchedDecodeIntoArray {
282 #[from]
284 source: AnyArrayAssignError,
285 },
286 #[error("ZfpClassic failed to decode the data")]
288 ZfpDecodeFailed,
289}
290
291#[derive(Debug, Error)]
292#[error(transparent)]
293pub struct ZfpHeaderError(postcard::Error);
295
296pub fn compress<T: ffi::ZfpCompressible, D: Dimension>(
315 data: ArrayView<T, D>,
316 mode: &ZfpCompressionMode,
317 non_finite: ZfpNonFiniteValuesMode,
318) -> Result<Vec<u8>, ZfpClassicCodecError> {
319 if !matches!(mode, ZfpCompressionMode::Reversible)
320 && !matches!(non_finite, ZfpNonFiniteValuesMode::AllowUnsafe)
321 && !Zip::from(&data).all(|x| x.is_finite())
322 {
323 return Err(ZfpClassicCodecError::NonFiniteData);
324 }
325
326 let mut encoded = postcard::to_extend(
327 &CompressionHeader {
328 dtype: <T as ffi::ZfpCompressible>::D_TYPE,
329 shape: Cow::Borrowed(data.shape()),
330 version: StaticCodecVersion,
331 },
332 Vec::new(),
333 )
334 .map_err(|err| ZfpClassicCodecError::MetaHeaderEncodeFailed {
335 source: ZfpHeaderError(err),
336 })?;
337
338 if data.is_empty() {
340 return Ok(encoded);
341 }
342
343 let field = ffi::ZfpField::new(data.into_dyn().squeeze())?;
346 let stream = ffi::ZfpCompressionStream::new(&field, mode)?;
347
348 let stream = stream.with_bitstream(field, &mut encoded);
351
352 let stream = stream.write_header()?;
354
355 stream.compress()?;
357
358 Ok(encoded)
359}
360
361pub fn decompress(encoded: &[u8]) -> Result<AnyArray, ZfpClassicCodecError> {
373 let (header, encoded) =
374 postcard::take_from_bytes::<CompressionHeader>(encoded).map_err(|err| {
375 ZfpClassicCodecError::MetaHeaderDecodeFailed {
376 source: ZfpHeaderError(err),
377 }
378 })?;
379
380 if header.shape.iter().copied().product::<usize>() == 0 {
382 let decoded = match header.dtype {
383 ZfpDType::I32 => AnyArray::I32(Array::zeros(&*header.shape)),
384 ZfpDType::I64 => AnyArray::I64(Array::zeros(&*header.shape)),
385 ZfpDType::F32 => AnyArray::F32(Array::zeros(&*header.shape)),
386 ZfpDType::F64 => AnyArray::F64(Array::zeros(&*header.shape)),
387 };
388 return Ok(decoded);
389 }
390
391 let stream = ffi::ZfpDecompressionStream::new(encoded);
393
394 let stream = stream.read_header()?;
396
397 match header.dtype {
399 ZfpDType::I32 => {
400 let mut decompressed = Array::zeros(&*header.shape);
401 stream.decompress_into(decompressed.view_mut().squeeze())?;
402 Ok(AnyArray::I32(decompressed))
403 }
404 ZfpDType::I64 => {
405 let mut decompressed = Array::zeros(&*header.shape);
406 stream.decompress_into(decompressed.view_mut().squeeze())?;
407 Ok(AnyArray::I64(decompressed))
408 }
409 ZfpDType::F32 => {
410 let mut decompressed = Array::zeros(&*header.shape);
411 stream.decompress_into(decompressed.view_mut().squeeze())?;
412 Ok(AnyArray::F32(decompressed))
413 }
414 ZfpDType::F64 => {
415 let mut decompressed = Array::zeros(&*header.shape);
416 stream.decompress_into(decompressed.view_mut().squeeze())?;
417 Ok(AnyArray::F64(decompressed))
418 }
419 }
420}
421
422pub fn decompress_into(
436 encoded: &[u8],
437 decoded: AnyArrayViewMut,
438) -> Result<(), ZfpClassicCodecError> {
439 let (header, encoded) =
440 postcard::take_from_bytes::<CompressionHeader>(encoded).map_err(|err| {
441 ZfpClassicCodecError::MetaHeaderDecodeFailed {
442 source: ZfpHeaderError(err),
443 }
444 })?;
445
446 if decoded.shape() != &*header.shape {
447 return Err(ZfpClassicCodecError::MismatchedDecodeIntoArray {
448 source: AnyArrayAssignError::ShapeMismatch {
449 src: header.shape.into_owned(),
450 dst: decoded.shape().to_vec(),
451 },
452 });
453 }
454
455 if decoded.is_empty() {
457 return Ok(());
458 }
459
460 let stream = ffi::ZfpDecompressionStream::new(encoded);
462
463 let stream = stream.read_header()?;
465
466 match (decoded, header.dtype) {
468 (AnyArrayViewMut::I32(decoded), ZfpDType::I32) => stream.decompress_into(decoded.squeeze()),
469 (AnyArrayViewMut::I64(decoded), ZfpDType::I64) => stream.decompress_into(decoded.squeeze()),
470 (AnyArrayViewMut::F32(decoded), ZfpDType::F32) => stream.decompress_into(decoded.squeeze()),
471 (AnyArrayViewMut::F64(decoded), ZfpDType::F64) => stream.decompress_into(decoded.squeeze()),
472 (decoded, dtype) => Err(ZfpClassicCodecError::MismatchedDecodeIntoArray {
473 source: AnyArrayAssignError::DTypeMismatch {
474 src: dtype.into_dtype(),
475 dst: decoded.dtype(),
476 },
477 }),
478 }
479}
480
481#[derive(Serialize, Deserialize)]
482struct CompressionHeader<'a> {
483 dtype: ZfpDType,
484 #[serde(borrow)]
485 shape: Cow<'a, [usize]>,
486 version: ZfpClassicCodecVersion,
487}
488
489#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
491#[expect(missing_docs)]
492pub enum ZfpDType {
493 #[serde(rename = "i32", alias = "int32")]
494 I32,
495 #[serde(rename = "i64", alias = "int64")]
496 I64,
497 #[serde(rename = "f32", alias = "float32")]
498 F32,
499 #[serde(rename = "f64", alias = "float64")]
500 F64,
501}
502
503impl ZfpDType {
504 #[must_use]
506 pub const fn into_dtype(self) -> AnyArrayDType {
507 match self {
508 Self::I32 => AnyArrayDType::I32,
509 Self::I64 => AnyArrayDType::I64,
510 Self::F32 => AnyArrayDType::F32,
511 Self::F64 => AnyArrayDType::F64,
512 }
513 }
514}
515
516impl fmt::Display for ZfpDType {
517 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
518 fmt.write_str(match self {
519 Self::I32 => "i32",
520 Self::I64 => "i64",
521 Self::F32 => "f32",
522 Self::F64 => "f64",
523 })
524 }
525}
526
527#[cfg(test)]
528#[allow(clippy::unwrap_used)]
529mod tests {
530 use ndarray::ArrayView1;
531
532 use super::*;
533
534 #[test]
535 fn zero_length() {
536 let encoded = compress(
537 Array::<f32, _>::from_shape_vec([1, 27, 0].as_slice(), vec![])
538 .unwrap()
539 .view(),
540 &ZfpCompressionMode::FixedPrecision { precision: 7 },
541 ZfpNonFiniteValuesMode::Deny,
542 )
543 .unwrap();
544 let decoded = decompress(&encoded).unwrap();
545
546 assert_eq!(decoded.dtype(), AnyArrayDType::F32);
547 assert!(decoded.is_empty());
548 assert_eq!(decoded.shape(), &[1, 27, 0]);
549 }
550
551 #[test]
552 fn one_dimension() {
553 let data = Array::from_shape_vec(
554 [2_usize, 1, 2, 1, 1, 1].as_slice(),
555 vec![1.0, 2.0, 3.0, 4.0],
556 )
557 .unwrap();
558
559 let encoded = compress(
560 data.view(),
561 &ZfpCompressionMode::FixedAccuracy { tolerance: 0.1 },
562 ZfpNonFiniteValuesMode::Deny,
563 )
564 .unwrap();
565 let decoded = decompress(&encoded).unwrap();
566
567 assert_eq!(decoded, AnyArray::F32(data));
568 }
569
570 #[test]
571 fn small_state() {
572 for data in [
573 &[][..],
574 &[0.0],
575 &[0.0, 1.0],
576 &[0.0, 1.0, 0.0],
577 &[0.0, 1.0, 0.0, 1.0],
578 ] {
579 let encoded = compress(
580 ArrayView1::from(data),
581 &ZfpCompressionMode::FixedAccuracy { tolerance: 0.1 },
582 ZfpNonFiniteValuesMode::Deny,
583 )
584 .unwrap();
585 let decoded = decompress(&encoded).unwrap();
586
587 assert_eq!(
588 decoded,
589 AnyArray::F64(Array1::from_vec(data.to_vec()).into_dyn())
590 );
591 }
592 }
593}