Skip to main content

zarrs_codec/
byte_interval_partial_decoder.rs

1use std::sync::Arc;
2
3use zarrs_storage::StorageError;
4use zarrs_storage::byte_range::{ByteLength, ByteOffset, ByteRange, ByteRangeIterator};
5
6use crate::{ArrayBytesRaw, CodecError, CodecOptions};
7
8#[cfg(feature = "async")]
9use crate::AsyncBytesPartialDecoderTraits;
10use crate::BytesPartialDecoderTraits;
11
12/// A partial decoder for a byte interval of a [`BytesPartialDecoderTraits`] partial decoder.
13///
14/// Modifies byte range requests to a specific byte interval in an inner bytes partial decoder.
15pub struct ByteIntervalPartialDecoder {
16    input_handle: Arc<dyn BytesPartialDecoderTraits>,
17    byte_offset: ByteOffset,
18    byte_length: ByteLength,
19}
20
21impl ByteIntervalPartialDecoder {
22    /// Create a new byte interval partial decoder.
23    pub fn new(
24        input_handle: Arc<dyn BytesPartialDecoderTraits>,
25        byte_offset: ByteOffset,
26        byte_length: ByteLength,
27    ) -> Self {
28        Self {
29            input_handle,
30            byte_offset,
31            byte_length,
32        }
33    }
34}
35
36impl BytesPartialDecoderTraits for ByteIntervalPartialDecoder {
37    fn exists(&self) -> Result<bool, StorageError> {
38        self.input_handle.exists()
39    }
40
41    fn size_held(&self) -> usize {
42        self.input_handle.size_held()
43    }
44
45    fn partial_decode_many(
46        &self,
47        byte_ranges: ByteRangeIterator,
48        options: &CodecOptions,
49    ) -> Result<Option<Vec<ArrayBytesRaw<'_>>>, CodecError> {
50        let byte_ranges = byte_ranges.map(|byte_range| match byte_range {
51            ByteRange::FromStart(offset, None) => {
52                ByteRange::FromStart(self.byte_offset + offset, Some(self.byte_length))
53            }
54            ByteRange::FromStart(offset, Some(length)) => {
55                ByteRange::FromStart(self.byte_offset + offset, Some(length))
56            }
57            ByteRange::Suffix(length) => {
58                ByteRange::FromStart(self.byte_offset + self.byte_length - length, Some(length))
59            }
60        });
61        self.input_handle
62            .partial_decode_many(Box::new(byte_ranges), options)
63    }
64
65    fn supports_partial_decode(&self) -> bool {
66        self.input_handle.supports_partial_decode()
67    }
68}
69
70#[cfg(feature = "async")]
71/// A partial decoder for a byte interval of a [`AsyncBytesPartialDecoderTraits`] partial decoder.
72///
73/// Modifies byte range requests to a specific byte interval in an inner bytes partial decoder.
74pub struct AsyncByteIntervalPartialDecoder {
75    input_handle: Arc<dyn AsyncBytesPartialDecoderTraits>,
76    byte_offset: ByteOffset,
77    byte_length: ByteLength,
78}
79
80#[cfg(feature = "async")]
81impl AsyncByteIntervalPartialDecoder {
82    /// Create a new byte interval partial decoder.
83    pub fn new(
84        input_handle: Arc<dyn AsyncBytesPartialDecoderTraits>,
85        byte_offset: ByteOffset,
86        byte_length: ByteLength,
87    ) -> Self {
88        Self {
89            input_handle,
90            byte_offset,
91            byte_length,
92        }
93    }
94}
95
96#[cfg(feature = "async")]
97#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
98#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
99impl AsyncBytesPartialDecoderTraits for AsyncByteIntervalPartialDecoder {
100    async fn exists(&self) -> Result<bool, StorageError> {
101        self.input_handle.exists().await
102    }
103
104    fn size_held(&self) -> usize {
105        self.input_handle.size_held()
106    }
107
108    async fn partial_decode_many<'a>(
109        &'a self,
110        byte_ranges: ByteRangeIterator<'a>,
111        options: &CodecOptions,
112    ) -> Result<Option<Vec<ArrayBytesRaw<'a>>>, CodecError> {
113        let byte_ranges = byte_ranges.map(|byte_range| match byte_range {
114            ByteRange::FromStart(offset, None) => {
115                ByteRange::FromStart(self.byte_offset + offset, Some(self.byte_length))
116            }
117            ByteRange::FromStart(offset, Some(length)) => {
118                ByteRange::FromStart(self.byte_offset + offset, Some(length))
119            }
120            ByteRange::Suffix(length) => {
121                ByteRange::FromStart(self.byte_offset + self.byte_length - length, Some(length))
122            }
123        });
124        self.input_handle
125            .partial_decode_many(Box::new(byte_ranges), options)
126            .await
127    }
128
129    fn supports_partial_decode(&self) -> bool {
130        self.input_handle.supports_partial_decode()
131    }
132}