diff --git a/CHANGELOG.md b/CHANGELOG.md index c41ed498..06101639 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Add store lock tests - Added `contiguous_elements` method to `ContiguousIndicesIterator` and `ContiguousLinearisedIndicesIterator` - Added `ChunkShape::num_elements` + - Added `codec::{Encode,Decode,PartialDecode,PartialDecoder}Options` + - Added new `Array::opt` methods which can use new encode/decode options + - **Breaking** Existing `Array` `_opt` use new encode/decode options insted of `parallel: bool` ### Changed - Dependency bumps @@ -32,7 +35,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `array_subset_iterators.rs` - **Major breaking**: storage transformers must be `Arc` wrapped as `StorageTransformerExtension` trait method now take `self: Arc` - Removed lifetimes from `{Async}{Readable,Writable,ReadableWritable,Listable,ReadableListable}Storage` - - **Breaking**: `Group` and `Array` methods generic on storage now require the storage with a `'static` lifetime + - **Breaking**: `Group` and `Array` methods generic on storage now require the storage have a `'static` lifetime + - **Breaking**: remove `Array::{set_}parallel_codecs` and `ArrayBuilder::parallel_codecs` + - **Breaking**: added `recommended_concurrency` to codec trait methods to facilitate improved parallelisation + - **Major breaking**: refactor codec traits: + - **Breaking**: remove `par_` variants, + - **Breaking**: `_opt` variants use new `codec::{Encode,Decode,PartialDecode,PartialDecoder}Options` instead of `parallel: bool` + - variants without prefix/suffix are no longer serial variants but parallel + - TODO: Remove these? ### Removed - **Breaking**: remove `InvalidArraySubsetError` and `ArrayExtractElementsError` diff --git a/Cargo.toml b/Cargo.toml index 5c1dfae6..9d24c4eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -71,6 +71,7 @@ walkdir = "2.3.2" zfp-sys = {version = "0.1.4", features = ["static"], optional = true } zip = { version = "0.6", optional = true } zstd = { version = "0.13", features = ["zstdmt"], optional = true } +rayon_iter_concurrent_limit = "0.1.0-alpha3" [dev-dependencies] chrono = "0.4" diff --git a/benches/array_blosc.rs b/benches/array_blosc.rs index ce7191e9..2455265f 100644 --- a/benches/array_blosc.rs +++ b/benches/array_blosc.rs @@ -32,9 +32,7 @@ fn array_blosc_write_all(c: &mut Criterion) { .unwrap(); let data = vec![1u8; num_elements.try_into().unwrap()]; let subset = zarrs::array_subset::ArraySubset::new_with_shape(vec![size; 3]); - array - .par_store_array_subset_elements(&subset, data) - .unwrap(); + array.store_array_subset_elements(&subset, data).unwrap(); }); }); } @@ -69,13 +67,11 @@ fn array_blosc_read_all(c: &mut Criterion) { .unwrap(); let data = vec![1u8; num_elements.try_into().unwrap()]; let subset = zarrs::array_subset::ArraySubset::new_with_shape(vec![size; 3]); - array - .par_store_array_subset_elements(&subset, data) - .unwrap(); + array.store_array_subset_elements(&subset, data).unwrap(); // Benchmark reading the data b.iter(|| { - let _bytes = array.par_retrieve_array_subset(&subset).unwrap(); + let _bytes = array.retrieve_array_subset(&subset).unwrap(); }); }); } diff --git a/benches/array_uncompressed.rs b/benches/array_uncompressed.rs index 05dc8122..9d8a69c0 100644 --- a/benches/array_uncompressed.rs +++ b/benches/array_uncompressed.rs @@ -19,9 +19,7 @@ fn array_write_all(c: &mut Criterion) { .unwrap(); let data = vec![1u8; num_elements.try_into().unwrap()]; let subset = zarrs::array_subset::ArraySubset::new_with_shape(vec![size; 3]); - array - .par_store_array_subset_elements(&subset, data) - .unwrap(); + array.store_array_subset_elements(&subset, data).unwrap(); }); }); } @@ -49,9 +47,7 @@ fn array_write_all_sharded(c: &mut Criterion) { .unwrap(); let data = vec![1u16; num_elements.try_into().unwrap()]; let subset = zarrs::array_subset::ArraySubset::new_with_shape(vec![size; 3]); - array - .par_store_array_subset_elements(&subset, data) - .unwrap(); + array.store_array_subset_elements(&subset, data).unwrap(); }); }); } @@ -76,13 +72,11 @@ fn array_read_all(c: &mut Criterion) { .unwrap(); let data = vec![1u16; num_elements.try_into().unwrap()]; let subset = zarrs::array_subset::ArraySubset::new_with_shape(vec![size; 3]); - array - .par_store_array_subset_elements(&subset, data) - .unwrap(); + array.store_array_subset_elements(&subset, data).unwrap(); // Benchmark reading the data b.iter(|| { - let _bytes = array.par_retrieve_array_subset(&subset).unwrap(); + let _bytes = array.retrieve_array_subset(&subset).unwrap(); }); }); } @@ -110,13 +104,11 @@ fn array_read_all_sharded(c: &mut Criterion) { .unwrap(); let data = vec![0u8; num_elements.try_into().unwrap()]; let subset = zarrs::array_subset::ArraySubset::new_with_shape(vec![size; 3]); - array - .par_store_array_subset_elements(&subset, data) - .unwrap(); + array.store_array_subset_elements(&subset, data).unwrap(); // Benchmark reading the data b.iter(|| { - let _bytes = array.par_retrieve_array_subset(&subset).unwrap(); + let _bytes = array.retrieve_array_subset(&subset).unwrap(); }); }); } diff --git a/benches/codecs.rs b/benches/codecs.rs index b450cabf..45421968 100644 --- a/benches/codecs.rs +++ b/benches/codecs.rs @@ -68,12 +68,6 @@ fn codec_blosc(c: &mut Criterion) { group.bench_function(BenchmarkId::new("decode", size3), |b| { b.iter(|| codec.decode(data_encoded.clone(), &rep).unwrap()); }); - group.bench_function(BenchmarkId::new("par_encode", size3), |b| { - b.iter(|| codec.par_encode(data_decoded.clone()).unwrap()); - }); - group.bench_function(BenchmarkId::new("par_decode", size3), |b| { - b.iter(|| codec.par_decode(data_encoded.clone(), &rep).unwrap()); - }); } } diff --git a/examples/array_write_read.rs b/examples/array_write_read.rs index ea117ab3..5dd7d68a 100644 --- a/examples/array_write_read.rs +++ b/examples/array_write_read.rs @@ -75,7 +75,7 @@ fn array_write_read() -> Result<(), Box> { println!("store_chunk [0, 0] and [0, 1]:\n{data_all:+4.1}\n"); // Store multiple chunks - array.store_chunks_elements_opt::( + array.store_chunks_elements::( &ArraySubset::new_with_ranges(&[1..2, 0..2]), vec![ // @@ -83,7 +83,6 @@ fn array_write_read() -> Result<(), Box> { // 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, 1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1, ], - true, )?; let data_all = array.retrieve_array_subset_ndarray::(&subset_all)?; println!("store_chunks [1..2, 0..2]:\n{data_all:+4.1}\n"); diff --git a/examples/sharded_array_write_read.rs b/examples/sharded_array_write_read.rs index 3647e529..fb966189 100644 --- a/examples/sharded_array_write_read.rs +++ b/examples/sharded_array_write_read.rs @@ -131,7 +131,7 @@ fn sharded_array_write_read() -> Result<(), Box> { ArraySubset::new_with_start_shape(vec![0, 0], inner_chunk_shape.clone())?, ArraySubset::new_with_start_shape(vec![0, 4], inner_chunk_shape.clone())?, ]; - let decoded_inner_chunks = partial_decoder.par_partial_decode(&inner_chunks_to_decode)?; + let decoded_inner_chunks = partial_decoder.partial_decode(&inner_chunks_to_decode)?; let decoded_inner_chunks = decoded_inner_chunks .into_iter() .map(|bytes| { diff --git a/src/array.rs b/src/array.rs index f6614067..bfab462f 100644 --- a/src/array.rs +++ b/src/array.rs @@ -215,8 +215,6 @@ pub struct Array { dimension_names: Option>, /// Additional fields annotated with `"must_understand": false`. additional_fields: AdditionalFields, - /// If true, codecs run with multithreading (where supported) - parallel_codecs: bool, /// Zarrs metadata. include_zarrs_metadata: bool, } @@ -289,7 +287,6 @@ impl Array { additional_fields: metadata.additional_fields, storage_transformers, dimension_names: metadata.dimension_names, - parallel_codecs: true, include_zarrs_metadata: true, }) } @@ -371,19 +368,6 @@ impl Array { &self.additional_fields } - /// Returns true if codecs can use multiple threads for encoding and decoding (where supported). - #[must_use] - pub const fn parallel_codecs(&self) -> bool { - self.parallel_codecs - } - - /// Enable or disable multithreaded codec encoding/decoding. Enabled by default. - /// - /// It may be advantageous to turn this off if parallelisation is external to avoid thrashing. - pub fn set_parallel_codecs(&mut self, parallel_codecs: bool) { - self.parallel_codecs = parallel_codecs; - } - /// Enable or disable the inclusion of zarrs metadata in the array attributes. Enabled by default. /// /// Zarrs metadata includes the zarrs version and some parameters. diff --git a/src/array/array_async_readable.rs b/src/array/array_async_readable.rs index 8eb52b66..8326fac1 100644 --- a/src/array/array_async_readable.rs +++ b/src/array/array_async_readable.rs @@ -12,7 +12,7 @@ use crate::{ use super::{ codec::{ ArrayCodecTraits, ArrayToBytesCodecTraits, AsyncArrayPartialDecoderTraits, - AsyncStoragePartialDecoder, + AsyncStoragePartialDecoder, DecodeOptions, PartialDecoderOptions, }, transmute_from_bytes_vec, unsafe_cell_slice::UnsafeCellSlice, @@ -50,9 +50,10 @@ impl Array { /// /// # Panics /// Panics if the number of elements in the chunk exceeds `usize::MAX`. - pub async fn async_retrieve_chunk_if_exists( + pub async fn async_retrieve_chunk_if_exists_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result>, ArrayError> { let storage_handle = Arc::new(StorageHandle::new(self.storage.clone())); let storage_transformer = self @@ -70,7 +71,7 @@ impl Array { let chunk_representation = self.chunk_array_representation(chunk_indices)?; let chunk_decoded = self .codecs() - .async_decode_opt(chunk_encoded, &chunk_representation, self.parallel_codecs()) + .async_decode_opt(chunk_encoded, &chunk_representation, options) .await .map_err(ArrayError::CodecError)?; let chunk_decoded_size = @@ -88,6 +89,16 @@ impl Array { } } + /// Read and decode the chunk at `chunk_indices` into its bytes if it exists (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_if_exists( + &self, + chunk_indices: &[u64], + ) -> Result>, ArrayError> { + self.async_retrieve_chunk_if_exists_opt(chunk_indices, &DecodeOptions::default()) + .await + } + /// Read and decode the chunk at `chunk_indices` into its bytes or the fill value if it does not exist. /// /// # Errors @@ -98,8 +109,14 @@ impl Array { /// /// # Panics /// Panics if the number of elements in the chunk exceeds `usize::MAX`. - pub async fn async_retrieve_chunk(&self, chunk_indices: &[u64]) -> Result, ArrayError> { - let chunk = self.async_retrieve_chunk_if_exists(chunk_indices).await?; + pub async fn async_retrieve_chunk_opt( + &self, + chunk_indices: &[u64], + options: &DecodeOptions, + ) -> Result, ArrayError> { + let chunk = self + .async_retrieve_chunk_if_exists_opt(chunk_indices, options) + .await?; if let Some(chunk) = chunk { Ok(chunk) } else { @@ -109,6 +126,13 @@ impl Array { } } + /// Read and decode the chunk at `chunk_indices` into its bytes or the fill value if it does not exist (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk(&self, chunk_indices: &[u64]) -> Result, ArrayError> { + self.async_retrieve_chunk_opt(chunk_indices, &DecodeOptions::default()) + .await + } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements if it exists. /// /// # Errors @@ -118,15 +142,28 @@ impl Array { /// - `chunk_indices` are invalid, /// - there is a codec decoding error, or /// - an underlying store error. - pub async fn async_retrieve_chunk_elements_if_exists( + pub async fn async_retrieve_chunk_elements_if_exists_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result>, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.async_retrieve_chunk_if_exists(chunk_indices).await?; + let bytes = self + .async_retrieve_chunk_if_exists_opt(chunk_indices, options) + .await?; Ok(bytes.map(|bytes| transmute_from_bytes_vec::(bytes))) } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements if it exists (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_elements_if_exists( + &self, + chunk_indices: &[u64], + ) -> Result>, ArrayError> { + self.async_retrieve_chunk_elements_if_exists_opt(chunk_indices, &DecodeOptions::default()) + .await + } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements or the fill value if it does not exist. /// /// # Errors @@ -136,15 +173,28 @@ impl Array { /// - `chunk_indices` are invalid, /// - there is a codec decoding error, or /// - an underlying store error. - pub async fn async_retrieve_chunk_elements( + pub async fn async_retrieve_chunk_elements_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.async_retrieve_chunk(chunk_indices).await?; + let bytes = self + .async_retrieve_chunk_opt(chunk_indices, options) + .await?; Ok(transmute_from_bytes_vec::(bytes)) } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements or the fill value if it does not exist (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_elements( + &self, + chunk_indices: &[u64], + ) -> Result, ArrayError> { + self.async_retrieve_chunk_elements_opt(chunk_indices, &DecodeOptions::default()) + .await + } + #[cfg(feature = "ndarray")] /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`] if it exists. /// @@ -158,9 +208,10 @@ impl Array { /// /// # Panics /// Will panic if a chunk dimension is larger than `usize::MAX`. - pub async fn async_retrieve_chunk_ndarray_if_exists( + pub async fn async_retrieve_chunk_ndarray_if_exists_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result>, ArrayError> { // validate_element_size::(self.data_type())?; in // async_retrieve_chunk_elements_if_exists let shape = self @@ -168,7 +219,7 @@ impl Array { .chunk_shape_u64(chunk_indices, self.shape())? .ok_or_else(|| ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec()))?; let elements = self - .async_retrieve_chunk_elements_if_exists(chunk_indices) + .async_retrieve_chunk_elements_if_exists_opt(chunk_indices, options) .await?; if let Some(elements) = elements { Ok(Some(elements_to_ndarray(&shape, elements)?)) @@ -177,6 +228,17 @@ impl Array { } } + #[cfg(feature = "ndarray")] + /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`] if it exists (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_ndarray_if_exists( + &self, + chunk_indices: &[u64], + ) -> Result>, ArrayError> { + self.async_retrieve_chunk_ndarray_if_exists_opt(chunk_indices, &DecodeOptions::default()) + .await + } + #[cfg(feature = "ndarray")] /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. It is filled with the fill value if it does not exist. /// @@ -190,19 +252,33 @@ impl Array { /// /// # Panics /// Will panic if a chunk dimension is larger than `usize::MAX`. - pub async fn async_retrieve_chunk_ndarray( + pub async fn async_retrieve_chunk_ndarray_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in async_retrieve_chunk_elements let shape = self .chunk_grid() .chunk_shape_u64(chunk_indices, self.shape())? .ok_or_else(|| ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec()))?; - let elements = self.async_retrieve_chunk_elements(chunk_indices).await?; + let elements = self + .async_retrieve_chunk_elements_opt(chunk_indices, options) + .await?; elements_to_ndarray(&shape, elements) } + #[cfg(feature = "ndarray")] + /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. It is filled with the fill value if it does not exist (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_ndarray( + &self, + chunk_indices: &[u64], + ) -> Result, ArrayError> { + self.async_retrieve_chunk_ndarray_opt(chunk_indices, &DecodeOptions::default()) + .await + } + /// Read and decode the chunk at `chunk_indices` into its bytes. /// /// # Errors @@ -213,7 +289,11 @@ impl Array { /// /// # Panics /// Panics if the number of elements in the chunk exceeds `usize::MAX`. - pub async fn async_retrieve_chunks(&self, chunks: &ArraySubset) -> Result, ArrayError> { + pub async fn async_retrieve_chunks_opt( + &self, + chunks: &ArraySubset, + options: &DecodeOptions, + ) -> Result, ArrayError> { if chunks.dimensionality() != self.chunk_grid().dimensionality() { return Err(ArrayError::InvalidArraySubset( chunks.clone(), @@ -229,7 +309,7 @@ impl Array { 0 => Ok(vec![]), 1 => { let chunk_indices = chunks.start(); - self.async_retrieve_chunk(chunk_indices).await + self.async_retrieve_chunk_opt(chunk_indices, options).await } _ => { // Decode chunks and copy to output @@ -249,6 +329,7 @@ impl Array { chunk_indices, &array_subset, unsafe { output_slice.get() }, + options, ) }) .collect::>(); @@ -262,38 +343,74 @@ impl Array { } } + /// Read and decode the chunk at `chunk_indices` into its bytes (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunks(&self, chunks: &ArraySubset) -> Result, ArrayError> { + self.async_retrieve_chunks_opt(chunks, &DecodeOptions::default()) + .await + } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements. /// /// # Errors /// Returns an [`ArrayError`] if the size of `T` does not match the data type size or a [`Array::async_retrieve_chunks`] error condition is met. - pub async fn async_retrieve_chunks_elements( + pub async fn async_retrieve_chunks_elements_opt( &self, chunks: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.async_retrieve_chunks(chunks).await?; + let bytes = self.async_retrieve_chunks_opt(chunks, options).await?; Ok(transmute_from_bytes_vec::(bytes)) } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements (default options). + /// + /// # Errors + /// Returns an [`ArrayError`] if the size of `T` does not match the data type size or a [`Array::async_retrieve_chunks`] error condition is met. + pub async fn async_retrieve_chunks_elements( + &self, + chunks: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_chunks_elements_opt(chunks, &DecodeOptions::default()) + .await + } + + #[cfg(feature = "ndarray")] /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. /// /// # Errors /// Returns an [`ArrayError`] if the size of `T` does not match the data type size or a [`Array::async_retrieve_chunks`] error condition is met. - pub async fn async_retrieve_chunks_ndarray( + pub async fn async_retrieve_chunks_ndarray_opt( &self, chunks: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; let array_subset = self.chunks_subset(chunks)?; - let elements = self.async_retrieve_chunks_elements(chunks).await?; + let elements = self + .async_retrieve_chunks_elements_opt(chunks, options) + .await?; elements_to_ndarray(array_subset.shape(), elements) } + #[cfg(feature = "ndarray")] + /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`] (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunks_ndarray( + &self, + chunks: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_chunks_ndarray_opt(chunks, &DecodeOptions::default()) + .await + } + async fn _async_decode_chunk_into_array_subset( &self, chunk_indices: &[u64], array_subset: &ArraySubset, output: &mut [u8], + options: &DecodeOptions, ) -> Result<(), ArrayError> { // Get the subset of the array corresponding to the chunk let chunk_subset_in_array = unsafe { @@ -312,7 +429,7 @@ impl Array { let array_subset_in_chunk_subset = unsafe { overlap.relative_to_unchecked(chunk_subset_in_array.start()) }; let decoded_bytes = self - .async_retrieve_chunk_subset(chunk_indices, &array_subset_in_chunk_subset) + .async_retrieve_chunk_subset_opt(chunk_indices, &array_subset_in_chunk_subset, options) .await?; // Copy decoded bytes to the output @@ -349,9 +466,10 @@ impl Array { /// # Panics /// Panics if attempting to reference a byte beyond `usize::MAX`. #[allow(clippy::too_many_lines)] - pub async fn async_retrieve_array_subset( + pub async fn async_retrieve_array_subset_opt( &self, array_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { if array_subset.dimensionality() != self.chunk_grid().dimensionality() { return Err(ArrayError::InvalidArraySubset( @@ -378,7 +496,7 @@ impl Array { let chunk_subset = self.chunk_subset(chunk_indices).unwrap(); if &chunk_subset == array_subset { // Single chunk fast path if the array subset domain matches the chunk domain - self.async_retrieve_chunk(chunk_indices).await + self.async_retrieve_chunk_opt(chunk_indices, options).await } else { let size_output = usize::try_from( array_subset.num_elements() * self.data_type().size() as u64, @@ -395,6 +513,7 @@ impl Array { chunk_indices, array_subset, output_slice, + options, ) .await?; #[allow(clippy::transmute_undefined_repr)] @@ -452,11 +571,15 @@ impl Array { self.chunk_array_representation(&chunk_indices)?; let partial_decoder = self .codecs() - .async_partial_decoder(input_handle, &chunk_representation) + .async_partial_decoder_opt( + input_handle, + &chunk_representation, + options, // FIXME: Adjust internal decode options + ) .await?; partial_decoder - .par_partial_decode(&[array_subset_in_chunk_subset]) + .partial_decode_opt(&[array_subset_in_chunk_subset], options) // FIXME: Adjust internal decode options .await? .remove(0) }; @@ -501,6 +624,16 @@ impl Array { } } + /// Read and decode the `array_subset` of array into its bytes (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_array_subset( + &self, + array_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_array_subset_opt(array_subset, &DecodeOptions::default()) + .await + } + /// Read and decode the `array_subset` of array into a vector of its elements. /// /// # Errors @@ -510,15 +643,28 @@ impl Array { /// - an array subset is invalid or out of bounds of the array, /// - there is a codec decoding error, or /// - an underlying store error. - pub async fn async_retrieve_array_subset_elements( + pub async fn async_retrieve_array_subset_elements_opt( &self, array_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.async_retrieve_array_subset(array_subset).await?; + let bytes = self + .async_retrieve_array_subset_opt(array_subset, options) + .await?; Ok(transmute_from_bytes_vec::(bytes)) } + /// Read and decode the `array_subset` of array into a vector of its elements (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_array_subset_elements( + &self, + array_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_array_subset_elements_opt(array_subset, &DecodeOptions::default()) + .await + } + #[cfg(feature = "ndarray")] /// Read and decode the `array_subset` of array into an [`ndarray::ArrayD`]. /// @@ -530,17 +676,29 @@ impl Array { /// /// # Panics /// Will panic if any dimension in `chunk_subset` is `usize::MAX` or larger. - pub async fn async_retrieve_array_subset_ndarray( + pub async fn async_retrieve_array_subset_ndarray_opt( &self, array_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in async_retrieve_array_subset_elements let elements = self - .async_retrieve_array_subset_elements(array_subset) + .async_retrieve_array_subset_elements_opt(array_subset, options) .await?; elements_to_ndarray(array_subset.shape(), elements) } + #[cfg(feature = "ndarray")] + /// Read and decode the `array_subset` of array into an [`ndarray::ArrayD`] (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_array_subset_ndarray( + &self, + array_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_array_subset_ndarray_opt(array_subset, &DecodeOptions::default()) + .await + } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its bytes. /// /// # Errors @@ -552,10 +710,11 @@ impl Array { /// /// # Panics /// Will panic if the number of elements in `chunk_subset` is `usize::MAX` or larger. - pub async fn async_retrieve_chunk_subset( + pub async fn async_retrieve_chunk_subset_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { let chunk_representation = self.chunk_array_representation(chunk_indices)?; if !chunk_subset.inbounds(&chunk_representation.shape_u64()) { @@ -576,9 +735,9 @@ impl Array { let decoded_bytes = self .codecs() - .async_partial_decoder_opt(input_handle, &chunk_representation, self.parallel_codecs()) + .async_partial_decoder_opt(input_handle, &chunk_representation, options) .await? - .partial_decode_opt(&[chunk_subset.clone()], self.parallel_codecs()) + .partial_decode_opt(&[chunk_subset.clone()], options) .await?; let total_size = decoded_bytes.iter().map(Vec::len).sum::(); @@ -593,6 +752,17 @@ impl Array { } } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its bytes (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_subset( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_chunk_subset_opt(chunk_indices, chunk_subset, &DecodeOptions::default()) + .await + } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its elements. /// /// # Errors @@ -601,18 +771,34 @@ impl Array { /// - the chunk subset is invalid, /// - there is a codec decoding error, or /// - an underlying store error. - pub async fn async_retrieve_chunk_subset_elements( + pub async fn async_retrieve_chunk_subset_elements_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; let bytes = self - .async_retrieve_chunk_subset(chunk_indices, chunk_subset) + .async_retrieve_chunk_subset_opt(chunk_indices, chunk_subset, options) .await?; Ok(transmute_from_bytes_vec::(bytes)) } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its elements (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_subset_elements( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_chunk_subset_elements_opt( + chunk_indices, + chunk_subset, + &DecodeOptions::default(), + ) + .await + } + #[cfg(feature = "ndarray")] /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. /// @@ -625,26 +811,43 @@ impl Array { /// /// # Panics /// Will panic if the number of elements in `chunk_subset` is `usize::MAX` or larger. - pub async fn async_retrieve_chunk_subset_ndarray( + pub async fn async_retrieve_chunk_subset_ndarray_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in async_retrieve_chunk_subset_elements let elements = self - .async_retrieve_chunk_subset_elements(chunk_indices, chunk_subset) + .async_retrieve_chunk_subset_elements_opt(chunk_indices, chunk_subset, options) .await?; elements_to_ndarray(chunk_subset.shape(), elements) } - /// Initialises a partial decoder for the chunk at `chunk_indices` with optional parallelism. + #[cfg(feature = "ndarray")] + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into an [`ndarray::ArrayD`] (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_retrieve_chunk_subset_ndarray( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.async_retrieve_chunk_subset_ndarray_opt( + chunk_indices, + chunk_subset, + &DecodeOptions::default(), + ) + .await + } + + /// Initialises a partial decoder for the chunk at `chunk_indices`. /// /// # Errors /// Returns an [`ArrayError`] if initialisation of the partial decoder fails. pub async fn async_partial_decoder_opt<'a>( &'a self, chunk_indices: &[u64], - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, ArrayError> { let storage_handle = Arc::new(StorageHandle::new(self.storage.clone())); let storage_transformer = self @@ -657,11 +860,11 @@ impl Array { let chunk_representation = self.chunk_array_representation(chunk_indices)?; Ok(self .codecs() - .async_partial_decoder_opt(input_handle, &chunk_representation, parallel) + .async_partial_decoder_opt(input_handle, &chunk_representation, options) .await?) } - /// Initialises a partial decoder for the chunk at `chunk_indices`. + /// Initialises a partial decoder for the chunk at `chunk_indices` (default options). /// /// # Errors /// Returns an [`ArrayError`] if initialisation of the partial decoder fails. @@ -669,17 +872,7 @@ impl Array { &'a self, chunk_indices: &[u64], ) -> Result, ArrayError> { - self.async_partial_decoder_opt(chunk_indices, false).await - } - - /// Initialises a partial decoder for the chunk at `chunk_indices` using multithreading if applicable. - /// - /// # Errors - /// Returns an [`ArrayError`] if initialisation of the partial decoder fails. - pub async fn async_par_partial_decoder<'a>( - &'a self, - chunk_indices: &[u64], - ) -> Result, ArrayError> { - self.async_partial_decoder_opt(chunk_indices, true).await + self.async_partial_decoder_opt(chunk_indices, &PartialDecoderOptions::default()) + .await } } diff --git a/src/array/array_async_readable_writable.rs b/src/array/array_async_readable_writable.rs index 917d6951..d715ac2b 100644 --- a/src/array/array_async_readable_writable.rs +++ b/src/array/array_async_readable_writable.rs @@ -5,7 +5,10 @@ use crate::{ storage::{data_key, AsyncReadableWritableStorageTraits}, }; -use super::{Array, ArrayError}; +use super::{ + codec::{DecodeOptions, EncodeOptions}, + Array, ArrayError, +}; impl Array { /// Encode `subset_bytes` and store in `array_subset`. @@ -20,10 +23,12 @@ impl Array, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { // Validation if array_subset.dimensionality() != self.shape().len() { @@ -60,7 +65,8 @@ impl Array Array Array Array, + ) -> Result<(), ArrayError> { + self.async_store_array_subset_opt( + array_subset, + subset_bytes, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) + .await + } + /// Encode `subset_elements` and store in `array_subset`. /// /// Prefer to use [`store_chunk`](Array::store_chunk) since this will decode and encode each chunk intersecting `array_subset`. @@ -142,28 +168,53 @@ impl Array( + pub async fn async_store_array_subset_elements_opt( &self, array_subset: &ArraySubset, subset_elements: Vec, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { array_async_store_elements!( self, subset_elements, - async_store_array_subset(array_subset, subset_elements) + async_store_array_subset_opt( + array_subset, + subset_elements, + encode_options, + decode_options + ) ) } + /// Encode `subset_elements` and store in `array_subset` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_array_subset_elements( + &self, + array_subset: &ArraySubset, + subset_elements: Vec, + ) -> Result<(), ArrayError> { + self.async_store_array_subset_elements_opt( + array_subset, + subset_elements, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) + .await + } + #[cfg(feature = "ndarray")] /// Encode `subset_array` and store in the array subset starting at `subset_start`. /// /// # Errors /// Returns an [`ArrayError`] if a [`store_array_subset_elements`](Array::store_array_subset_elements) error condition is met. #[allow(clippy::missing_panics_doc)] - pub async fn async_store_array_subset_ndarray( + pub async fn async_store_array_subset_ndarray_opt( &self, subset_start: &[u64], subset_array: &ndarray::ArrayViewD<'_, T>, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { let subset = ArraySubset::new_with_start_shape( subset_start.to_vec(), @@ -172,8 +223,30 @@ impl Array( + &self, + subset_start: &[u64], + subset_array: &ndarray::ArrayViewD<'_, T>, + ) -> Result<(), ArrayError> { + self.async_store_array_subset_ndarray_opt( + subset_start, + subset_array, + &EncodeOptions::default(), + &DecodeOptions::default(), ) + .await } /// Encode `chunk_subset_bytes` and store in `chunk_subset` of the chunk at `chunk_indices`. @@ -188,11 +261,13 @@ impl Array, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { // Validation if let Some(chunk_shape) = self @@ -219,7 +294,7 @@ impl Array Array Array Array, + ) -> Result<(), ArrayError> { + self.async_store_chunk_subset_opt( + chunk_indices, + chunk_subset, + chunk_subset_bytes, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) + .await + } + /// Encode `chunk_subset_elements` and store in `chunk_subset` of the chunk at `chunk_indices`. /// /// Prefer to use [`store_chunk`](Array::store_chunk) since this will decode the chunk before updating it and reencoding it. @@ -263,17 +359,43 @@ impl Array( + pub async fn async_store_chunk_subset_elements_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, chunk_subset_elements: Vec, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { array_async_store_elements!( self, chunk_subset_elements, - async_store_chunk_subset(chunk_indices, chunk_subset, chunk_subset_elements) + async_store_chunk_subset_opt( + chunk_indices, + chunk_subset, + chunk_subset_elements, + encode_options, + decode_options + ) + ) + } + + /// Encode `chunk_subset_elements` and store in `chunk_subset` of the chunk at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunk_subset_elements( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + chunk_subset_elements: Vec, + ) -> Result<(), ArrayError> { + self.async_store_chunk_subset_elements_opt( + chunk_indices, + chunk_subset, + chunk_subset_elements, + &EncodeOptions::default(), + &DecodeOptions::default(), ) + .await } #[cfg(feature = "ndarray")] @@ -284,11 +406,13 @@ impl Array( + pub async fn async_store_chunk_subset_ndarray_opt( &self, chunk_indices: &[u64], chunk_subset_start: &[u64], chunk_subset_array: &ndarray::ArrayViewD<'_, T>, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { let subset = ArraySubset::new_with_start_shape( chunk_subset_start.to_vec(), @@ -301,7 +425,32 @@ impl Array( + &self, + chunk_indices: &[u64], + chunk_subset_start: &[u64], + chunk_subset_array: &ndarray::ArrayViewD<'_, T>, + ) -> Result<(), ArrayError> { + self.async_store_chunk_subset_ndarray_opt( + chunk_indices, + chunk_subset_start, + chunk_subset_array, + &EncodeOptions::default(), + &DecodeOptions::default(), ) + .await } } diff --git a/src/array/array_async_writable.rs b/src/array/array_async_writable.rs index e4a976cc..41f88ee9 100644 --- a/src/array/array_async_writable.rs +++ b/src/array/array_async_writable.rs @@ -7,7 +7,10 @@ use crate::{ storage::{AsyncWritableStorageTraits, StorageError, StorageHandle}, }; -use super::{codec::ArrayCodecTraits, Array, ArrayError}; +use super::{ + codec::{ArrayCodecTraits, EncodeOptions}, + Array, ArrayError, +}; impl Array { /// Store metadata. @@ -33,10 +36,11 @@ impl Array { /// - the length of `chunk_bytes` is not equal to the expected length (the product of the number of elements in the chunk and the data type size in bytes), /// - there is a codec encoding error, or /// - an underlying store error. - pub async fn async_store_chunk( + pub async fn async_store_chunk_opt( &self, chunk_indices: &[u64], chunk_bytes: Vec, + options: &EncodeOptions, ) -> Result<(), ArrayError> { // Validation let chunk_array_representation = self.chunk_array_representation(chunk_indices)?; @@ -58,11 +62,7 @@ impl Array { .create_async_writable_transformer(storage_handle); let chunk_encoded: Vec = self .codecs() - .async_encode_opt( - chunk_bytes, - &chunk_array_representation, - self.parallel_codecs(), - ) + .async_encode_opt(chunk_bytes, &chunk_array_representation, options) .await .map_err(ArrayError::CodecError)?; crate::storage::async_store_chunk( @@ -77,6 +77,17 @@ impl Array { } } + /// Encode `chunk_bytes` and store at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunk( + &self, + chunk_indices: &[u64], + chunk_bytes: Vec, + ) -> Result<(), ArrayError> { + self.async_store_chunk_opt(chunk_indices, chunk_bytes, &EncodeOptions::default()) + .await + } + /// Encode `chunk_elements` and store at `chunk_indices`. /// /// A chunk composed entirely of the fill value will not be written to the store. @@ -85,18 +96,34 @@ impl Array { /// Returns an [`ArrayError`] if /// - the size of `T` does not match the data type size, or /// - a [`store_chunk`](Array::store_chunk) error condition is met. - pub async fn async_store_chunk_elements( + pub async fn async_store_chunk_elements_opt( &self, chunk_indices: &[u64], chunk_elements: Vec, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_async_store_elements!( self, chunk_elements, - async_store_chunk(chunk_indices, chunk_elements) + async_store_chunk_opt(chunk_indices, chunk_elements, options) ) } + /// Encode `chunk_elements` and store at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunk_elements( + &self, + chunk_indices: &[u64], + chunk_elements: Vec, + ) -> Result<(), ArrayError> { + self.async_store_chunk_elements_opt( + chunk_indices, + chunk_elements, + &EncodeOptions::default(), + ) + .await + } + #[cfg(feature = "ndarray")] /// Encode `chunk_array` and store at `chunk_indices`. /// @@ -105,18 +132,31 @@ impl Array { /// - the size of `T` does not match the size of the data type, /// - a [`store_chunk_elements`](Array::store_chunk_elements) error condition is met. #[allow(clippy::missing_panics_doc)] - pub async fn async_store_chunk_ndarray( + pub async fn async_store_chunk_ndarray_opt( &self, chunk_indices: &[u64], chunk_array: &ndarray::ArrayViewD<'_, T>, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_async_store_ndarray!( self, chunk_array, - async_store_chunk_elements(chunk_indices, chunk_array) + async_store_chunk_elements_opt(chunk_indices, chunk_array, options) ) } + #[cfg(feature = "ndarray")] + /// Encode `chunk_array` and store at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunk_ndarray( + &self, + chunk_indices: &[u64], + chunk_array: &ndarray::ArrayViewD<'_, T>, + ) -> Result<(), ArrayError> { + self.async_store_chunk_ndarray_opt(chunk_indices, chunk_array, &EncodeOptions::default()) + .await + } + /// Encode `chunks_bytes` and store at the chunks with indices represented by the `chunks` array subset. /// /// A chunk composed entirely of the fill value will not be written to the store. @@ -128,15 +168,17 @@ impl Array { /// - there is a codec encoding error, or /// - an underlying store error. #[allow(clippy::similar_names, clippy::missing_panics_doc)] - pub async fn async_store_chunks( + pub async fn async_store_chunks_opt( &self, chunks: &ArraySubset, chunks_bytes: Vec, + options: &EncodeOptions, ) -> Result<(), ArrayError> { let num_chunks = chunks.num_elements_usize(); if num_chunks == 1 { let chunk_indices = chunks.start(); - self.async_store_chunk(chunk_indices, chunks_bytes).await?; + self.async_store_chunk_opt(chunk_indices, chunks_bytes, options) + .await?; } else { let array_subset = self.chunks_subset(chunks)?; let element_size = self.data_type().size(); @@ -184,7 +226,7 @@ impl Array { chunk_subset_in_array_subset.num_elements() ); - self.async_store_chunk(chunk_indices, chunk_bytes) + self.async_store_chunk_opt(chunk_indices, chunk_bytes, options) }) .collect::>(); while let Some(item) = futures.next().await { @@ -194,39 +236,76 @@ impl Array { Ok(()) } + /// Encode `chunks_bytes` and store at the chunks with indices represented by the `chunks` array subset (default options). + #[allow(clippy::similar_names)] + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunks( + &self, + chunks: &ArraySubset, + chunks_bytes: Vec, + ) -> Result<(), ArrayError> { + self.async_store_chunks_opt(chunks, chunks_bytes, &EncodeOptions::default()) + .await + } + /// Variation of [`Array::async_store_chunks`] for elements with a known type. /// /// # Errors /// In addition to [`Array::async_store_chunks`] errors, returns an [`ArrayError`] if the size of `T` does not match the data type size. - pub async fn async_store_chunks_elements( + pub async fn async_store_chunks_elements_opt( &self, chunks: &ArraySubset, chunks_elements: Vec, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_async_store_elements!( self, chunks_elements, - async_store_chunks(chunks, chunks_elements) + async_store_chunks_opt(chunks, chunks_elements, options) ) } + /// Variation of [`Array::async_store_chunks`] for elements with a known type (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunks_elements( + &self, + chunks: &ArraySubset, + chunks_elements: Vec, + ) -> Result<(), ArrayError> { + self.async_store_chunks_elements_opt(chunks, chunks_elements, &EncodeOptions::default()) + .await + } + #[cfg(feature = "ndarray")] /// Variation of [`Array::async_store_chunks`] for an [`ndarray::ArrayViewD`]. /// /// # Errors /// In addition to [`Array::async_store_chunks`] errors, returns an [`ArrayError`] if the size of `T` does not match the data type size. - pub async fn async_store_chunks_ndarray( + pub async fn async_store_chunks_ndarray_opt( &self, chunks: &ArraySubset, chunks_array: &ndarray::ArrayViewD<'_, T>, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_async_store_ndarray!( self, chunks_array, - async_store_chunks_elements(chunks, chunks_array) + async_store_chunks_elements_opt(chunks, chunks_array, options) ) } + #[cfg(feature = "ndarray")] + /// Variation of [`Array::async_store_chunks`] for an [`ndarray::ArrayViewD`] (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub async fn async_store_chunks_ndarray( + &self, + chunks: &ArraySubset, + chunks_array: &ndarray::ArrayViewD<'_, T>, + ) -> Result<(), ArrayError> { + self.async_store_chunks_ndarray_opt(chunks, chunks_array, &EncodeOptions::default()) + .await + } + /// Erase the chunk at `chunk_indices`. /// /// Succeeds if the key does not exist. diff --git a/src/array/array_builder.rs b/src/array/array_builder.rs index 14523bab..360d216a 100644 --- a/src/array/array_builder.rs +++ b/src/array/array_builder.rs @@ -78,8 +78,6 @@ pub struct ArrayBuilder { pub dimension_names: Option>, /// Additional fields. pub additional_fields: AdditionalFields, - /// Parallel codecs. - pub parallel_codecs: bool, } impl ArrayBuilder { @@ -107,7 +105,6 @@ impl ArrayBuilder { storage_transformers: StorageTransformerChain::default(), dimension_names: None, additional_fields: AdditionalFields::default(), - parallel_codecs: true, } } @@ -125,7 +122,6 @@ impl ArrayBuilder { .attributes(array.attributes().clone()) .chunk_key_encoding(array.chunk_key_encoding().clone()) .dimension_names(array.dimension_names().clone()) - .parallel_codecs(array.parallel_codecs()) .array_to_array_codecs(array.codecs().array_to_array_codecs().to_vec()) .array_to_bytes_codec(array.codecs().array_to_bytes_codec().clone()) .bytes_to_bytes_codecs(array.codecs().bytes_to_bytes_codecs().to_vec()) @@ -251,14 +247,6 @@ impl ArrayBuilder { self } - /// Set whether or not to use multithreaded codec encoding and decoding. - /// - /// If parallel codecs is not set, it defaults to true. - pub fn parallel_codecs(&mut self, parallel_codecs: bool) -> &mut Self { - self.parallel_codecs = parallel_codecs; - self - } - /// Build into an [`Array`]. /// /// # Errors @@ -312,7 +300,6 @@ impl ArrayBuilder { attributes: self.attributes.clone(), dimension_names: self.dimension_names.clone(), additional_fields: self.additional_fields.clone(), - parallel_codecs: self.parallel_codecs, include_zarrs_metadata: true, }) } @@ -346,7 +333,6 @@ mod tests { builder.fill_value(FillValue::from(0i8)); builder.dimension_names(Some(vec!["y".into(), "x".into()])); - builder.parallel_codecs(true); let mut attributes = serde_json::Map::new(); attributes.insert("key".to_string(), "value".into()); @@ -376,7 +362,6 @@ mod tests { assert_eq!(array.chunk_grid_shape(), Some(vec![4, 4])); assert_eq!(array.fill_value(), &FillValue::from(0i8)); assert_eq!(array.dimension_names(), &Some(vec!["y".into(), "x".into()])); - assert!(array.parallel_codecs()); assert_eq!(array.attributes(), &attributes); assert_eq!(array.additional_fields(), &additional_fields); @@ -387,7 +372,6 @@ mod tests { assert_eq!(builder.attributes, builder2.attributes); assert_eq!(builder.dimension_names, builder2.dimension_names); assert_eq!(builder.additional_fields, builder2.additional_fields); - assert_eq!(builder.parallel_codecs, builder2.parallel_codecs); } #[test] diff --git a/src/array/array_sync_readable.rs b/src/array/array_sync_readable.rs index 2f12b3e5..584d3d6a 100644 --- a/src/array/array_sync_readable.rs +++ b/src/array/array_sync_readable.rs @@ -10,7 +10,8 @@ use crate::{ use super::{ codec::{ - ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, StoragePartialDecoder, + ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, DecodeOptions, + PartialDecoderOptions, StoragePartialDecoder, }, transmute_from_bytes_vec, unravel_index, unsafe_cell_slice::UnsafeCellSlice, @@ -47,9 +48,10 @@ impl Array { /// /// # Panics /// Panics if the number of elements in the chunk exceeds `usize::MAX`. - pub fn retrieve_chunk_if_exists( + pub fn retrieve_chunk_if_exists_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result>, ArrayError> { let storage_handle = Arc::new(StorageHandle::new(self.storage.clone())); let storage_transformer = self @@ -66,7 +68,7 @@ impl Array { let chunk_representation = self.chunk_array_representation(chunk_indices)?; let chunk_decoded = self .codecs() - .decode_opt(chunk_encoded, &chunk_representation, self.parallel_codecs()) + .decode_opt(chunk_encoded, &chunk_representation, options) .map_err(ArrayError::CodecError)?; let chunk_decoded_size = chunk_representation.num_elements_usize() * chunk_representation.data_type().size(); @@ -83,6 +85,15 @@ impl Array { } } + /// Read and decode the chunk at `chunk_indices` into its bytes if it exists (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_if_exists( + &self, + chunk_indices: &[u64], + ) -> Result>, ArrayError> { + self.retrieve_chunk_if_exists_opt(chunk_indices, &DecodeOptions::default()) + } + /// Read and decode the chunk at `chunk_indices` into its bytes or the fill value if it does not exist. /// /// # Errors @@ -93,8 +104,12 @@ impl Array { /// /// # Panics /// Panics if the number of elements in the chunk exceeds `usize::MAX`. - pub fn retrieve_chunk(&self, chunk_indices: &[u64]) -> Result, ArrayError> { - let chunk = self.retrieve_chunk_if_exists(chunk_indices)?; + pub fn retrieve_chunk_opt( + &self, + chunk_indices: &[u64], + options: &DecodeOptions, + ) -> Result, ArrayError> { + let chunk = self.retrieve_chunk_if_exists_opt(chunk_indices, options)?; if let Some(chunk) = chunk { Ok(chunk) } else { @@ -104,6 +119,12 @@ impl Array { } } + /// Read and decode the chunk at `chunk_indices` into its bytes or the fill value if it does not exist (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk(&self, chunk_indices: &[u64]) -> Result, ArrayError> { + self.retrieve_chunk_opt(chunk_indices, &DecodeOptions::default()) + } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements if it exists. /// /// # Errors @@ -113,15 +134,25 @@ impl Array { /// - `chunk_indices` are invalid, /// - there is a codec decoding error, or /// - an underlying store error. - pub fn retrieve_chunk_elements_if_exists( + pub fn retrieve_chunk_elements_if_exists_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result>, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.retrieve_chunk_if_exists(chunk_indices)?; + let bytes = self.retrieve_chunk_if_exists_opt(chunk_indices, options)?; Ok(bytes.map(|bytes| transmute_from_bytes_vec::(bytes))) } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements if it exists (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_elements_if_exists( + &self, + chunk_indices: &[u64], + ) -> Result>, ArrayError> { + self.retrieve_chunk_elements_if_exists_opt(chunk_indices, &DecodeOptions::default()) + } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements or the fill value if it does not exist. /// /// # Errors @@ -131,15 +162,25 @@ impl Array { /// - `chunk_indices` are invalid, /// - there is a codec decoding error, or /// - an underlying store error. - pub fn retrieve_chunk_elements( + pub fn retrieve_chunk_elements_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.retrieve_chunk(chunk_indices)?; + let bytes = self.retrieve_chunk_opt(chunk_indices, options)?; Ok(transmute_from_bytes_vec::(bytes)) } + /// Read and decode the chunk at `chunk_indices` into a vector of its elements or the fill value if it does not exist (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_elements( + &self, + chunk_indices: &[u64], + ) -> Result, ArrayError> { + self.retrieve_chunk_elements_opt(chunk_indices, &DecodeOptions::default()) + } + #[cfg(feature = "ndarray")] /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`] if it exists. /// @@ -153,16 +194,17 @@ impl Array { /// /// # Panics /// Will panic if a chunk dimension is larger than `usize::MAX`. - pub fn retrieve_chunk_ndarray_if_exists( + pub fn retrieve_chunk_ndarray_if_exists_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result>, ArrayError> { // validate_element_size::(self.data_type())?; // in retrieve_chunk_elements_if_exists let shape = self .chunk_grid() .chunk_shape_u64(chunk_indices, self.shape())? .ok_or_else(|| ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec()))?; - let elements = self.retrieve_chunk_elements_if_exists::(chunk_indices)?; + let elements = self.retrieve_chunk_elements_if_exists_opt::(chunk_indices, options)?; if let Some(elements) = elements { Ok(Some(elements_to_ndarray(&shape, elements)?)) } else { @@ -170,6 +212,16 @@ impl Array { } } + #[cfg(feature = "ndarray")] + /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`] if it exists (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_ndarray_if_exists( + &self, + chunk_indices: &[u64], + ) -> Result>, ArrayError> { + self.retrieve_chunk_ndarray_if_exists_opt(chunk_indices, &DecodeOptions::default()) + } + #[cfg(feature = "ndarray")] /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. It is filled with the fill value if it does not exist. /// @@ -183,16 +235,30 @@ impl Array { /// /// # Panics /// Will panic if a chunk dimension is larger than `usize::MAX`. - pub fn retrieve_chunk_ndarray( + pub fn retrieve_chunk_ndarray_opt( &self, chunk_indices: &[u64], + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in retrieve_chunk_elements let shape = self .chunk_grid() .chunk_shape_u64(chunk_indices, self.shape())? .ok_or_else(|| ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec()))?; - elements_to_ndarray(&shape, self.retrieve_chunk_elements::(chunk_indices)?) + elements_to_ndarray( + &shape, + self.retrieve_chunk_elements_opt::(chunk_indices, options)?, + ) + } + + #[cfg(feature = "ndarray")] + /// Read and decode the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. It is filled with the fill value if it does not exist (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_ndarray( + &self, + chunk_indices: &[u64], + ) -> Result, ArrayError> { + self.retrieve_chunk_ndarray_opt(chunk_indices, &DecodeOptions::default()) } /// Read and decode the chunks at `chunks` into their bytes. @@ -208,7 +274,7 @@ impl Array { pub fn retrieve_chunks_opt( &self, chunks: &ArraySubset, - parallel: bool, + options: &DecodeOptions, ) -> Result, ArrayError> { if chunks.dimensionality() != self.chunk_grid().dimensionality() { return Err(ArrayError::InvalidArraySubset( @@ -225,7 +291,7 @@ impl Array { 0 => Ok(vec![]), 1 => { let chunk_indices = chunks.start(); - self.retrieve_chunk(chunk_indices) + self.retrieve_chunk_opt(chunk_indices, options) } _ => { // Decode chunks and copy to output @@ -239,7 +305,7 @@ impl Array { let output_slice = unsafe { std::slice::from_raw_parts_mut(output.as_mut_ptr().cast::(), size_output) }; - if parallel { + if options.is_parallel() { let output = UnsafeCellSlice::new(output_slice); (0..chunks.shape().iter().product()) .into_par_iter() @@ -259,6 +325,7 @@ impl Array { &chunk_indices, &array_subset, unsafe { output.get() }, + options, ) })?; } else { @@ -267,6 +334,7 @@ impl Array { &chunk_indices, &array_subset, output_slice, + options, )?; } } @@ -277,16 +345,14 @@ impl Array { } } - /// Serial version of [`Array::retrieve_chunks_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + /// Read and decode the chunks at `chunks` into their bytes (default options). + /// + /// # Errors + /// See [`Array::retrieve_chunks_opt`]. + /// # Panics + /// See [`Array::retrieve_chunks_opt`]. pub fn retrieve_chunks(&self, chunks: &ArraySubset) -> Result, ArrayError> { - self.retrieve_chunks_opt(chunks, false) - } - - /// Parallel version of [`Array::retrieve_chunks_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_retrieve_chunks(&self, chunks: &ArraySubset) -> Result, ArrayError> { - self.retrieve_chunks_opt(chunks, true) + self.retrieve_chunks_opt(chunks, &DecodeOptions::default()) } /// Read and decode the chunks at `chunks` into a vector of its elements. @@ -299,29 +365,24 @@ impl Array { pub fn retrieve_chunks_elements_opt( &self, chunks: &ArraySubset, - parallel: bool, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.retrieve_chunks_opt(chunks, parallel)?; + let bytes = self.retrieve_chunks_opt(chunks, options)?; Ok(transmute_from_bytes_vec::(bytes)) } - /// Serial version of [`Array::retrieve_chunks_elements_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + /// Read and decode the chunks at `chunks` into a vector of its elements with default options + /// + /// # Errors + /// See [`Array::retrieve_chunks_elements_opt`]. + /// # Panics + /// See [`Array::retrieve_chunks_elements_opt`]. pub fn retrieve_chunks_elements( &self, chunks: &ArraySubset, ) -> Result, ArrayError> { - self.retrieve_chunks_elements_opt(chunks, false) - } - - /// Parallel version of [`Array::retrieve_chunks_elements_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_retrieve_chunks_elements( - &self, - chunks: &ArraySubset, - ) -> Result, ArrayError> { - self.retrieve_chunks_elements_opt(chunks, true) + self.retrieve_chunks_elements_opt(chunks, &DecodeOptions::default()) } #[cfg(feature = "ndarray")] @@ -335,32 +396,26 @@ impl Array { pub fn retrieve_chunks_ndarray_opt( &self, chunks: &ArraySubset, - parallel: bool, + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in retrieve_chunks_elements_opt let array_subset = self.chunks_subset(chunks)?; - let elements = self.retrieve_chunks_elements_opt::(chunks, parallel)?; + let elements = self.retrieve_chunks_elements_opt::(chunks, options)?; elements_to_ndarray(array_subset.shape(), elements) } #[cfg(feature = "ndarray")] - /// Serial version of [`Array::retrieve_chunks_ndarray_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + /// Read and decode the chunks at `chunks` into an [`ndarray::ArrayD`] (default options). + /// + /// # Errors + /// See [`Array::retrieve_chunks_ndarray_opt`]. + /// # Panics + /// See [`Array::retrieve_chunks_ndarray_opt`]. pub fn retrieve_chunks_ndarray( &self, chunks: &ArraySubset, ) -> Result, ArrayError> { - self.retrieve_chunks_ndarray_opt(chunks, false) - } - - #[cfg(feature = "ndarray")] - /// Parallel version of [`Array::retrieve_chunks_ndarray_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_retrieve_chunks_ndarray( - &self, - chunks: &ArraySubset, - ) -> Result, ArrayError> { - self.retrieve_chunks_ndarray_opt(chunks, true) + self.retrieve_chunks_ndarray_opt(chunks, &DecodeOptions::default()) } fn _decode_chunk_into_array_subset( @@ -368,6 +423,7 @@ impl Array { chunk_indices: &[u64], array_subset: &ArraySubset, output: &mut [u8], + options: &DecodeOptions, ) -> Result<(), ArrayError> { // Get the subset of the array corresponding to the chunk let chunk_subset_in_array = unsafe { @@ -386,7 +442,7 @@ impl Array { let array_subset_in_chunk_subset = unsafe { overlap.relative_to_unchecked(chunk_subset_in_array.start()) }; let decoded_bytes = - self.retrieve_chunk_subset(chunk_indices, &array_subset_in_chunk_subset)?; + self.retrieve_chunk_subset_opt(chunk_indices, &array_subset_in_chunk_subset, options)?; // Copy decoded bytes to the output let element_size = self.data_type().size() as u64; @@ -424,7 +480,7 @@ impl Array { pub fn retrieve_array_subset_opt( &self, array_subset: &ArraySubset, - parallel: bool, + options: &DecodeOptions, ) -> Result, ArrayError> { if array_subset.dimensionality() != self.chunk_grid().dimensionality() { return Err(ArrayError::InvalidArraySubset( @@ -451,7 +507,7 @@ impl Array { let chunk_subset = self.chunk_subset(chunk_indices).unwrap(); if &chunk_subset == array_subset { // Single chunk fast path if the array subset domain matches the chunk domain - self.retrieve_chunk(chunk_indices) + self.retrieve_chunk_opt(chunk_indices, options) } else { let size_output = usize::try_from( array_subset.num_elements() * self.data_type().size() as u64, @@ -468,6 +524,7 @@ impl Array { chunk_indices, array_subset, output_slice, + options, )?; #[allow(clippy::transmute_undefined_repr)] let output: Vec = unsafe { core::mem::transmute(output) }; @@ -486,7 +543,9 @@ impl Array { let output_slice = unsafe { std::slice::from_raw_parts_mut(output.as_mut_ptr().cast::(), size_output) }; - if parallel { + if options.is_parallel() { + // FIXME: Constrain concurrency here based on parallelism internally vs externally + let output = UnsafeCellSlice::new(output_slice); (0..chunks.shape().iter().product()) .into_par_iter() @@ -498,14 +557,12 @@ impl Array { .map(|(chunk_indices, chunks_start)| chunk_indices + chunks_start) .collect::>() }) - // chunks - // .iter_indices() - // .par_bridge() .try_for_each(|chunk_indices| { self._decode_chunk_into_array_subset( &chunk_indices, array_subset, unsafe { output.get() }, + options, ) })?; } else { @@ -514,6 +571,7 @@ impl Array { &chunk_indices, array_subset, output_slice, + options, )?; } } @@ -524,19 +582,14 @@ impl Array { } } - /// Serial version of [`Array::retrieve_array_subset_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + /// Read and decode the `array_subset` of array into its bytes. + /// + /// # Errors + /// See [`Array::retrieve_array_subset_opt`]. + /// # Panics + /// See [`Array::retrieve_array_subset_opt`]. pub fn retrieve_array_subset(&self, array_subset: &ArraySubset) -> Result, ArrayError> { - self.retrieve_array_subset_opt(array_subset, false) - } - - /// Parallel version of [`Array::retrieve_array_subset_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_retrieve_array_subset( - &self, - array_subset: &ArraySubset, - ) -> Result, ArrayError> { - self.retrieve_array_subset_opt(array_subset, true) + self.retrieve_array_subset_opt(array_subset, &DecodeOptions::default()) } /// Read and decode the `array_subset` of array into a vector of its elements. @@ -551,29 +604,24 @@ impl Array { pub fn retrieve_array_subset_elements_opt( &self, array_subset: &ArraySubset, - parallel: bool, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.retrieve_array_subset_opt(array_subset, parallel)?; + let bytes = self.retrieve_array_subset_opt(array_subset, options)?; Ok(transmute_from_bytes_vec::(bytes)) } - /// Serial version of [`Array::retrieve_array_subset_elements_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + /// Read and decode the `array_subset` of array into a vector of its elements (default options). + /// + /// # Errors + /// See [`Array::retrieve_array_subset_elements_opt`]. + /// # Panics + /// See [`Array::retrieve_array_subset_elements_opt`]. pub fn retrieve_array_subset_elements( &self, array_subset: &ArraySubset, ) -> Result, ArrayError> { - self.retrieve_array_subset_elements_opt(array_subset, false) - } - - /// Parallel version of [`Array::retrieve_array_subset_elements_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_retrieve_array_subset_elements( - &self, - array_subset: &ArraySubset, - ) -> Result, ArrayError> { - self.retrieve_array_subset_elements_opt(array_subset, true) + self.retrieve_array_subset_elements_opt(array_subset, &DecodeOptions::default()) } #[cfg(feature = "ndarray")] @@ -590,31 +638,25 @@ impl Array { pub fn retrieve_array_subset_ndarray_opt( &self, array_subset: &ArraySubset, - parallel: bool, + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in retrieve_array_subset_elements_opt - let elements = self.retrieve_array_subset_elements_opt::(array_subset, parallel)?; + let elements = self.retrieve_array_subset_elements_opt::(array_subset, options)?; elements_to_ndarray(array_subset.shape(), elements) } #[cfg(feature = "ndarray")] - /// Serial version of [`Array::retrieve_array_subset_ndarray`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + /// Read and decode the `array_subset` of array into an [`ndarray::ArrayD`] (default options). + /// + /// # Errors + /// See [`Array::retrieve_array_subset_ndarray`]. + /// # Panics + /// See [`Array::retrieve_array_subset_ndarray`]. pub fn retrieve_array_subset_ndarray( &self, array_subset: &ArraySubset, ) -> Result, ArrayError> { - self.retrieve_array_subset_ndarray_opt(array_subset, false) - } - - #[cfg(feature = "ndarray")] - /// Parallel version of [`Array::retrieve_array_subset_ndarray`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_retrieve_array_subset_ndarray( - &self, - array_subset: &ArraySubset, - ) -> Result, ArrayError> { - self.retrieve_array_subset_ndarray_opt(array_subset, true) + self.retrieve_array_subset_ndarray_opt(array_subset, &DecodeOptions::default()) } /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its bytes. @@ -628,10 +670,11 @@ impl Array { /// /// # Panics /// Will panic if the number of elements in `chunk_subset` is `usize::MAX` or larger. - pub fn retrieve_chunk_subset( + pub fn retrieve_chunk_subset_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { let chunk_representation = self.chunk_array_representation(chunk_indices)?; if !chunk_subset.inbounds(&chunk_representation.shape_u64()) { @@ -652,8 +695,8 @@ impl Array { let decoded_bytes = self .codecs() - .partial_decoder_opt(input_handle, &chunk_representation, self.parallel_codecs())? - .partial_decode_opt(&[chunk_subset.clone()], self.parallel_codecs())?; + .partial_decoder_opt(input_handle, &chunk_representation, options)? + .partial_decode_opt(&[chunk_subset.clone()], options)?; let total_size = decoded_bytes.iter().map(Vec::len).sum::(); let expected_size = chunk_subset.num_elements_usize() * self.data_type().size(); @@ -667,6 +710,16 @@ impl Array { } } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its bytes (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_subset( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.retrieve_chunk_subset_opt(chunk_indices, chunk_subset, &DecodeOptions::default()) + } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its elements. /// /// # Errors @@ -675,16 +728,31 @@ impl Array { /// - the chunk subset is invalid, /// - there is a codec decoding error, or /// - an underlying store error. - pub fn retrieve_chunk_subset_elements( + pub fn retrieve_chunk_subset_elements_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { validate_element_size::(self.data_type())?; - let bytes = self.retrieve_chunk_subset(chunk_indices, chunk_subset)?; + let bytes = self.retrieve_chunk_subset_opt(chunk_indices, chunk_subset, options)?; Ok(transmute_from_bytes_vec::(bytes)) } + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into its elements (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_subset_elements( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.retrieve_chunk_subset_elements_opt( + chunk_indices, + chunk_subset, + &DecodeOptions::default(), + ) + } + #[cfg(feature = "ndarray")] /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into an [`ndarray::ArrayD`]. /// @@ -697,24 +765,41 @@ impl Array { /// /// # Panics /// Will panic if the number of elements in `chunk_subset` is `usize::MAX` or larger. - pub fn retrieve_chunk_subset_ndarray( + pub fn retrieve_chunk_subset_ndarray_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, + options: &DecodeOptions, ) -> Result, ArrayError> { // validate_element_size::(self.data_type())?; // in retrieve_chunk_subset_elements - let elements = self.retrieve_chunk_subset_elements::(chunk_indices, chunk_subset)?; + let elements = + self.retrieve_chunk_subset_elements_opt::(chunk_indices, chunk_subset, options)?; elements_to_ndarray(chunk_subset.shape(), elements) } - /// Initialises a partial decoder for the chunk at `chunk_indices` with optional parallelism. + #[cfg(feature = "ndarray")] + /// Read and decode the `chunk_subset` of the chunk at `chunk_indices` into an [`ndarray::ArrayD`] (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn retrieve_chunk_subset_ndarray( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + ) -> Result, ArrayError> { + self.retrieve_chunk_subset_ndarray_opt( + chunk_indices, + chunk_subset, + &DecodeOptions::default(), + ) + } + + /// Initialises a partial decoder for the chunk at `chunk_indices`. /// /// # Errors /// Returns an [`ArrayError`] if initialisation of the partial decoder fails. pub fn partial_decoder_opt<'a>( &'a self, chunk_indices: &[u64], - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, ArrayError> { let storage_handle = Arc::new(StorageHandle::new(self.storage.clone())); let storage_transformer = self @@ -727,10 +812,10 @@ impl Array { let chunk_representation = self.chunk_array_representation(chunk_indices)?; Ok(self .codecs() - .partial_decoder_opt(input_handle, &chunk_representation, parallel)?) + .partial_decoder_opt(input_handle, &chunk_representation, options)?) } - /// Initialises a partial decoder for the chunk at `chunk_indices`. + /// Initialises a partial decoder for the chunk at `chunk_indices` (default options). /// /// # Errors /// Returns an [`ArrayError`] if initialisation of the partial decoder fails. @@ -738,17 +823,6 @@ impl Array { &'a self, chunk_indices: &[u64], ) -> Result, ArrayError> { - self.partial_decoder_opt(chunk_indices, false) - } - - /// Initialises a partial decoder for the chunk at `chunk_indices` using multithreading if applicable. - /// - /// # Errors - /// Returns an [`ArrayError`] if initialisation of the partial decoder fails. - pub fn par_partial_decoder<'a>( - &'a self, - chunk_indices: &[u64], - ) -> Result, ArrayError> { - self.partial_decoder_opt(chunk_indices, true) + self.partial_decoder_opt(chunk_indices, &PartialDecoderOptions::default()) } } diff --git a/src/array/array_sync_readable_writable.rs b/src/array/array_sync_readable_writable.rs index 45b0012e..80a3a36c 100644 --- a/src/array/array_sync_readable_writable.rs +++ b/src/array/array_sync_readable_writable.rs @@ -5,7 +5,10 @@ use crate::{ storage::{data_key, ReadableWritableStorageTraits}, }; -use super::{unravel_index, Array, ArrayError}; +use super::{ + codec::{DecodeOptions, EncodeOptions}, + unravel_index, Array, ArrayError, +}; impl Array { /// Encode `subset_bytes` and store in `array_subset`. @@ -24,7 +27,8 @@ impl Array &self, array_subset: &ArraySubset, subset_bytes: Vec, - parallel: bool, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { // Validation if array_subset.dimensionality() != self.shape().len() { @@ -61,7 +65,7 @@ impl Array if array_subset == &chunk_subset_in_array { // A fast path if the array subset matches the chunk subset // This skips the internal decoding occurring in store_chunk_subset - self.store_chunk(chunk_indices, subset_bytes)?; + self.store_chunk_opt(chunk_indices, subset_bytes, encode_options)?; } else { let overlap = unsafe { array_subset.overlap_unchecked(&chunk_subset_in_array) }; let chunk_subset_in_array_subset = @@ -78,10 +82,12 @@ impl Array let array_subset_in_chunk_subset = unsafe { overlap.relative_to_unchecked(chunk_subset_in_array.start()) }; - self.store_chunk_subset( + self.store_chunk_subset_opt( chunk_indices, &array_subset_in_chunk_subset, chunk_subset_bytes, + encode_options, + decode_options, )?; } } else { @@ -106,15 +112,17 @@ impl Array let array_subset_in_chunk_subset = unsafe { overlap.relative_to_unchecked(chunk_subset_in_array.start()) }; - self.store_chunk_subset( + self.store_chunk_subset_opt( &chunk_indices, &array_subset_in_chunk_subset, chunk_subset_bytes, + encode_options, + decode_options, )?; Ok(()) }; - if parallel { + if encode_options.is_parallel() { (0..chunks.shape().iter().product()) .into_par_iter() .map(|chunk_index| { @@ -135,24 +143,21 @@ impl Array Ok(()) } - /// Serial version of [`Array::store_array_subset_opt`]. + /// Encode `subset_bytes` and store in `array_subset` (default options). + /// + /// See [`Array::store_array_subset_opt`]. #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] pub fn store_array_subset( &self, array_subset: &ArraySubset, subset_bytes: Vec, ) -> Result<(), ArrayError> { - self.store_array_subset_opt(array_subset, subset_bytes, false) - } - - /// Parallel version of [`Array::store_array_subset_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_store_array_subset( - &self, - array_subset: &ArraySubset, - subset_bytes: Vec, - ) -> Result<(), ArrayError> { - self.store_array_subset_opt(array_subset, subset_bytes, true) + self.store_array_subset_opt( + array_subset, + subset_bytes, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) } /// Encode `subset_elements` and store in `array_subset`. @@ -167,33 +172,34 @@ impl Array &self, array_subset: &ArraySubset, subset_elements: Vec, - parallel: bool, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { array_store_elements!( self, subset_elements, - store_array_subset_opt(array_subset, subset_elements, parallel) + store_array_subset_opt( + array_subset, + subset_elements, + encode_options, + decode_options + ) ) } - /// Serial version of [`Array::store_array_subset_elements_opt`]. + /// Encode `subset_elements` and store in `array_subset` (default options). #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] pub fn store_array_subset_elements( &self, array_subset: &ArraySubset, subset_elements: Vec, ) -> Result<(), ArrayError> { - self.store_array_subset_elements_opt(array_subset, subset_elements, false) - } - - /// Parallel version of [`Array::store_array_subset_elements_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_store_array_subset_elements( - &self, - array_subset: &ArraySubset, - subset_elements: Vec, - ) -> Result<(), ArrayError> { - self.store_array_subset_elements_opt(array_subset, subset_elements, true) + self.store_array_subset_elements_opt( + array_subset, + subset_elements, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) } #[cfg(feature = "ndarray")] @@ -206,7 +212,8 @@ impl Array &self, subset_start: &[u64], subset_array: &ndarray::ArrayViewD, - parallel: bool, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { let subset = ArraySubset::new_with_start_shape( subset_start.to_vec(), @@ -215,30 +222,24 @@ impl Array array_store_ndarray!( self, subset_array, - store_array_subset_elements_opt(&subset, subset_array, parallel) + store_array_subset_elements_opt(&subset, subset_array, encode_options, decode_options) ) } #[cfg(feature = "ndarray")] - /// Serial version of [`Array::store_array_subset_ndarray_opt`]. + /// Encode `subset_array` and store in the array subset starting at `subset_start` (default options). #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] pub fn store_array_subset_ndarray( &self, subset_start: &[u64], subset_array: &ndarray::ArrayViewD, ) -> Result<(), ArrayError> { - self.store_array_subset_ndarray_opt(subset_start, subset_array, false) - } - - #[cfg(feature = "ndarray")] - /// Parallel version of [`Array::store_array_subset_ndarray_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_store_array_subset_ndarray( - &self, - subset_start: &[u64], - subset_array: &ndarray::ArrayViewD, - ) -> Result<(), ArrayError> { - self.store_array_subset_ndarray_opt(subset_start, subset_array, true) + self.store_array_subset_ndarray_opt( + subset_start, + subset_array, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) } /// Encode `chunk_subset_bytes` and store in `chunk_subset` of the chunk at `chunk_indices`. @@ -253,11 +254,13 @@ impl Array /// /// # Panics /// Panics if attempting to reference a byte beyond `usize::MAX`. - pub fn store_chunk_subset( + pub fn store_chunk_subset_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, chunk_subset_bytes: Vec, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { let chunk_shape = self .chunk_grid() @@ -285,7 +288,7 @@ impl Array if chunk_subset.shape() == chunk_shape && chunk_subset.start().iter().all(|&x| x == 0) { // The subset spans the whole chunk, so store the bytes directly and skip decoding - self.store_chunk(chunk_indices, chunk_subset_bytes) + self.store_chunk_opt(chunk_indices, chunk_subset_bytes, encode_options) } else { // Lock the chunk let key = data_key(self.path(), chunk_indices, self.chunk_key_encoding()); @@ -293,7 +296,7 @@ impl Array let _lock = mutex.lock(); // Decode the entire chunk - let mut chunk_bytes = self.retrieve_chunk(chunk_indices)?; + let mut chunk_bytes = self.retrieve_chunk_opt(chunk_indices, decode_options)?; // Update the intersecting subset of the chunk let element_size = self.data_type().size() as u64; @@ -311,10 +314,27 @@ impl Array } // Store the updated chunk - self.store_chunk(chunk_indices, chunk_bytes) + self.store_chunk_opt(chunk_indices, chunk_bytes, encode_options) } } + /// Encode `chunk_subset_bytes` and store in `chunk_subset` of the chunk at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn store_chunk_subset( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + chunk_subset_bytes: Vec, + ) -> Result<(), ArrayError> { + self.store_chunk_subset_opt( + chunk_indices, + chunk_subset, + chunk_subset_bytes, + &EncodeOptions::default(), + &DecodeOptions::default(), + ) + } + /// Encode `chunk_subset_elements` and store in `chunk_subset` of the chunk at `chunk_indices`. /// /// Prefer to use [`store_chunk`](Array::store_chunk) since this will decode the chunk before updating it and reencoding it. @@ -323,16 +343,41 @@ impl Array /// Returns an [`ArrayError`] if /// - the size of `T` does not match the data type size, or /// - a [`store_chunk_subset`](Array::store_chunk_subset) error condition is met. - pub fn store_chunk_subset_elements( + pub fn store_chunk_subset_elements_opt( &self, chunk_indices: &[u64], chunk_subset: &ArraySubset, chunk_subset_elements: Vec, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { array_store_elements!( self, chunk_subset_elements, - store_chunk_subset(chunk_indices, chunk_subset, chunk_subset_elements) + store_chunk_subset_opt( + chunk_indices, + chunk_subset, + chunk_subset_elements, + encode_options, + decode_options + ) + ) + } + + /// Encode `chunk_subset_elements` and store in `chunk_subset` of the chunk at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn store_chunk_subset_elements( + &self, + chunk_indices: &[u64], + chunk_subset: &ArraySubset, + chunk_subset_elements: Vec, + ) -> Result<(), ArrayError> { + self.store_chunk_subset_elements_opt( + chunk_indices, + chunk_subset, + chunk_subset_elements, + &EncodeOptions::default(), + &DecodeOptions::default(), ) } @@ -344,11 +389,13 @@ impl Array /// # Errors /// Returns an [`ArrayError`] if a [`store_chunk_subset_elements`](Array::store_chunk_subset_elements) error condition is met. #[allow(clippy::missing_panics_doc)] - pub fn store_chunk_subset_ndarray( + pub fn store_chunk_subset_ndarray_opt( &self, chunk_indices: &[u64], chunk_subset_start: &[u64], chunk_subset_array: &ndarray::ArrayViewD, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, ) -> Result<(), ArrayError> { let subset = ArraySubset::new_with_start_shape( chunk_subset_start.to_vec(), @@ -361,7 +408,31 @@ impl Array array_store_ndarray!( self, chunk_subset_array, - store_chunk_subset_elements(chunk_indices, &subset, chunk_subset_array) + store_chunk_subset_elements_opt( + chunk_indices, + &subset, + chunk_subset_array, + encode_options, + decode_options + ) + ) + } + + #[cfg(feature = "ndarray")] + /// Encode `chunk_subset_array` and store in `chunk_subset` of the chunk in the subset starting at `chunk_subset_start` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn store_chunk_subset_ndarray( + &self, + chunk_indices: &[u64], + chunk_subset_start: &[u64], + chunk_subset_array: &ndarray::ArrayViewD, + ) -> Result<(), ArrayError> { + self.store_chunk_subset_ndarray_opt( + chunk_indices, + chunk_subset_start, + chunk_subset_array, + &EncodeOptions::default(), + &DecodeOptions::default(), ) } } diff --git a/src/array/array_sync_writable.rs b/src/array/array_sync_writable.rs index 8e382997..8feac4ea 100644 --- a/src/array/array_sync_writable.rs +++ b/src/array/array_sync_writable.rs @@ -7,7 +7,10 @@ use crate::{ storage::{StorageError, StorageHandle, WritableStorageTraits}, }; -use super::{codec::ArrayCodecTraits, unravel_index, Array, ArrayError}; +use super::{ + codec::{ArrayCodecTraits, EncodeOptions}, + unravel_index, Array, ArrayError, +}; impl Array { /// Store metadata. @@ -32,10 +35,11 @@ impl Array { /// - the length of `chunk_bytes` is not equal to the expected length (the product of the number of elements in the chunk and the data type size in bytes), /// - there is a codec encoding error, or /// - an underlying store error. - pub fn store_chunk( + pub fn store_chunk_opt( &self, chunk_indices: &[u64], chunk_bytes: Vec, + options: &EncodeOptions, ) -> Result<(), ArrayError> { // Validation let chunk_array_representation = self.chunk_array_representation(chunk_indices)?; @@ -57,11 +61,7 @@ impl Array { .create_writable_transformer(storage_handle); let chunk_encoded: Vec = self .codecs() - .encode_opt( - chunk_bytes, - &chunk_array_representation, - self.parallel_codecs(), - ) + .encode_opt(chunk_bytes, &chunk_array_representation, options) .map_err(ArrayError::CodecError)?; crate::storage::store_chunk( &*storage_transformer, @@ -74,6 +74,16 @@ impl Array { } } + /// Encode `chunk_bytes` and store at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn store_chunk( + &self, + chunk_indices: &[u64], + chunk_bytes: Vec, + ) -> Result<(), ArrayError> { + self.store_chunk_opt(chunk_indices, chunk_bytes, &EncodeOptions::default()) + } + /// Encode `chunk_elements` and store at `chunk_indices`. /// /// A chunk composed entirely of the fill value will not be written to the store. @@ -82,18 +92,29 @@ impl Array { /// Returns an [`ArrayError`] if /// - the size of `T` does not match the data type size, or /// - a [`store_chunk`](Array::store_chunk) error condition is met. - pub fn store_chunk_elements( + pub fn store_chunk_elements_opt( &self, chunk_indices: &[u64], chunk_elements: Vec, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_store_elements!( self, chunk_elements, - store_chunk(chunk_indices, chunk_elements) + store_chunk_opt(chunk_indices, chunk_elements, options) ) } + /// Encode `chunk_elements` and store at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn store_chunk_elements( + &self, + chunk_indices: &[u64], + chunk_elements: Vec, + ) -> Result<(), ArrayError> { + self.store_chunk_elements_opt(chunk_indices, chunk_elements, &EncodeOptions::default()) + } + #[cfg(feature = "ndarray")] /// Encode `chunk_array` and store at `chunk_indices`. /// @@ -102,18 +123,30 @@ impl Array { /// - the size of `T` does not match the size of the data type, /// - a [`store_chunk_elements`](Array::store_chunk_elements) error condition is met. #[allow(clippy::missing_panics_doc)] - pub fn store_chunk_ndarray( + pub fn store_chunk_ndarray_opt( &self, chunk_indices: &[u64], chunk_array: &ndarray::ArrayViewD, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_store_ndarray!( self, chunk_array, - store_chunk_elements(chunk_indices, chunk_array) + store_chunk_elements_opt(chunk_indices, chunk_array, options) ) } + #[cfg(feature = "ndarray")] + /// Encode `chunk_array` and store at `chunk_indices` (default options). + #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] + pub fn store_chunk_ndarray( + &self, + chunk_indices: &[u64], + chunk_array: &ndarray::ArrayViewD, + ) -> Result<(), ArrayError> { + self.store_chunk_ndarray_opt(chunk_indices, chunk_array, &EncodeOptions::default()) + } + /// Encode `chunks_bytes` and store at the chunks with indices represented by the `chunks` array subset. /// /// A chunk composed entirely of the fill value will not be written to the store. @@ -129,12 +162,12 @@ impl Array { &self, chunks: &ArraySubset, chunks_bytes: Vec, - parallel: bool, + options: &EncodeOptions, ) -> Result<(), ArrayError> { let num_chunks = chunks.num_elements_usize(); if num_chunks == 1 { let chunk_indices = chunks.start(); - self.store_chunk(chunk_indices, chunks_bytes)?; + self.store_chunk_opt(chunk_indices, chunks_bytes, options)?; } else { let array_subset = self.chunks_subset(chunks)?; let element_size = self.data_type().size(); @@ -172,11 +205,11 @@ impl Array { ); // Store the chunk - self.store_chunk(&chunk_indices, chunk_bytes)?; + self.store_chunk_opt(&chunk_indices, chunk_bytes, options)?; Ok(()) }; - if parallel { + if options.is_parallel() { (0..chunks.shape().iter().product()) .into_par_iter() .map(|chunk_index| { @@ -195,24 +228,15 @@ impl Array { Ok(()) } - /// Serial version of [`Array::store_chunks_opt`]. + /// Encode `chunks_bytes` and store at the chunks with indices represented by the `chunks` array subset (default options). + #[allow(clippy::similar_names)] #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] pub fn store_chunks( &self, chunks: &ArraySubset, chunks_bytes: Vec, ) -> Result<(), ArrayError> { - self.store_chunks_opt(chunks, chunks_bytes, false) - } - - /// Parallel version of [`Array::store_chunks_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_store_chunks( - &self, - chunks: &ArraySubset, - chunks_bytes: Vec, - ) -> Result<(), ArrayError> { - self.store_chunks_opt(chunks, chunks_bytes, true) + self.store_chunks_opt(chunks, chunks_bytes, &EncodeOptions::default()) } /// Variation of [`Array::store_chunks_opt`] for elements with a known type. @@ -223,41 +247,23 @@ impl Array { &self, chunks: &ArraySubset, chunks_elements: Vec, - parallel: bool, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_store_elements!( self, chunks_elements, - store_chunks_opt(chunks, chunks_elements, parallel) + store_chunks_opt(chunks, chunks_elements, options) ) } - /// Serial version of [`Array::store_chunks_elements_opt`]. + /// Variation of [`Array::store_chunks_opt`] for elements with a known type (default options). #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] pub fn store_chunks_elements( &self, chunks: &ArraySubset, chunks_elements: Vec, ) -> Result<(), ArrayError> { - array_store_elements!( - self, - chunks_elements, - store_chunks_opt(chunks, chunks_elements, false) - ) - } - - /// Parallel version of [`Array::store_chunks_elements_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_store_chunks_elements( - &self, - chunks: &ArraySubset, - chunks_elements: Vec, - ) -> Result<(), ArrayError> { - array_store_elements!( - self, - chunks_elements, - store_chunks_opt(chunks, chunks_elements, true) - ) + self.store_chunks_elements_opt(chunks, chunks_elements, &EncodeOptions::default()) } #[cfg(feature = "ndarray")] @@ -269,43 +275,24 @@ impl Array { &self, chunks: &ArraySubset, chunks_array: &ndarray::ArrayViewD<'_, T>, - parallel: bool, + options: &EncodeOptions, ) -> Result<(), ArrayError> { array_store_ndarray!( self, chunks_array, - store_chunks_elements_opt(chunks, chunks_array, parallel) + store_chunks_elements_opt(chunks, chunks_array, options) ) } #[cfg(feature = "ndarray")] - /// Serial version of [`Array::store_chunks_ndarray_opt`]. + /// Variation of [`Array::store_chunks_opt`] for an [`ndarray::ArrayViewD`] (default options). #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] pub fn store_chunks_ndarray( &self, chunks: &ArraySubset, chunks_array: &ndarray::ArrayViewD<'_, T>, ) -> Result<(), ArrayError> { - array_store_ndarray!( - self, - chunks_array, - store_chunks_elements_opt(chunks, chunks_array, false) - ) - } - - #[cfg(feature = "ndarray")] - /// Parallel version of [`Array::store_chunks_ndarray_opt`]. - #[allow(clippy::missing_panics_doc, clippy::missing_errors_doc)] - pub fn par_store_chunks_ndarray( - &self, - chunks: &ArraySubset, - chunks_array: &ndarray::ArrayViewD<'_, T>, - ) -> Result<(), ArrayError> { - array_store_ndarray!( - self, - chunks_array, - store_chunks_elements_opt(chunks, chunks_array, true) - ) + self.store_chunks_ndarray_opt(chunks, chunks_array, &EncodeOptions::default()) } /// Erase the chunk at `chunk_indices`. @@ -331,57 +318,29 @@ impl Array { /// /// # Errors /// Returns a [`StorageError`] if there is an underlying store error. - pub fn erase_chunks_opt( - &self, - chunks: &ArraySubset, - parallel: bool, - ) -> Result<(), StorageError> { + pub fn erase_chunks(&self, chunks: &ArraySubset) -> Result<(), StorageError> { let storage_handle = Arc::new(StorageHandle::new(self.storage.clone())); let storage_transformer = self .storage_transformers() .create_writable_transformer(storage_handle); - if parallel { - (0..chunks.shape().iter().product()) - .into_par_iter() - .map(|chunk_index| { - std::iter::zip(unravel_index(chunk_index, chunks.shape()), chunks.start()) - .map(|(chunk_indices, chunks_start)| chunk_indices + chunks_start) - .collect::>() - }) - // chunks - // .iter_indices() - // .par_bridge() - .try_for_each(|chunk_indices| { - crate::storage::erase_chunk( - &*storage_transformer, - self.path(), - &chunk_indices, - self.chunk_key_encoding(), - )?; - Ok::<_, StorageError>(()) - })?; - } else { - for chunk_indices in chunks.iter_indices() { + (0..chunks.shape().iter().product()) + .into_par_iter() + .map(|chunk_index| { + std::iter::zip(unravel_index(chunk_index, chunks.shape()), chunks.start()) + .map(|(chunk_indices, chunks_start)| chunk_indices + chunks_start) + .collect::>() + }) + // chunks + // .iter_indices() + // .par_bridge() + .try_for_each(|chunk_indices| { crate::storage::erase_chunk( &*storage_transformer, self.path(), &chunk_indices, self.chunk_key_encoding(), )?; - } - } - Ok(()) - } - - /// Serial version of [`Array::store_chunks_ndarray_opt`]. - #[allow(clippy::missing_errors_doc)] - pub fn erase_chunks(&self, chunks: &ArraySubset) -> Result<(), StorageError> { - self.erase_chunks_opt(chunks, false) - } - - /// Parallel version of [`Array::store_chunks_ndarray_opt`]. - #[allow(clippy::missing_errors_doc)] - pub fn par_erase_chunks(&self, chunks: &ArraySubset) -> Result<(), StorageError> { - self.erase_chunks_opt(chunks, true) + Ok::<_, StorageError>(()) + }) } } diff --git a/src/array/codec.rs b/src/array/codec.rs index 85995390..262e506b 100644 --- a/src/array/codec.rs +++ b/src/array/codec.rs @@ -13,6 +13,9 @@ pub mod array_to_array; pub mod array_to_bytes; pub mod bytes_to_bytes; +pub mod options; + +pub use options::{DecodeOptions, EncodeOptions, PartialDecodeOptions, PartialDecoderOptions}; // Array to array #[cfg(feature = "bitround")] @@ -74,6 +77,7 @@ use crate::storage::AsyncReadableStorage; use std::{ collections::{BTreeMap, BTreeSet}, io::{Read, Seek, SeekFrom}, + num::NonZeroU64, }; use super::{BytesRepresentation, ChunkRepresentation, DataType, MaybeBytes}; @@ -93,6 +97,53 @@ pub enum Codec { BytesToBytes(Box), } +/// The recommended concurrency of a codec includes the most efficient and maximum recommended concurrency. +/// +/// Consider a chain that does slow decoding first on a single thread, but subsequent codecs can run on multiple threads. +/// In this case, recommended concurrency is best expressed by two numbers: +/// - the efficient concurrency, equal to the minimum of codecs +/// - the maximum concurrency, equal to the maximum of codecs +// TODO: Compression codec example in docs? +#[derive(Debug, Copy, Clone)] +pub struct RecommendedConcurrency { + efficient: NonZeroU64, + maximum: NonZeroU64, +} + +impl RecommendedConcurrency { + /// Create a new recommended concurrency struct with an explicit efficient and maximum recommendation. + #[must_use] + pub fn new(efficient: NonZeroU64, maximum: NonZeroU64) -> Self { + Self { efficient, maximum } + } + + /// Create a new recommended concurrency struct with an efficient and maximum recommendation of one. + #[must_use] + pub fn one() -> Self { + Self::new(unsafe { NonZeroU64::new_unchecked(1) }, unsafe { + NonZeroU64::new_unchecked(1) + }) + } + + /// Return the recommended efficient concurrency. + #[must_use] + pub fn efficient(&self) -> u64 { + self.efficient.get() + } + + /// Return the recommended maximum concurrency. + #[must_use] + pub fn maximum(&self) -> u64 { + self.maximum.get() + } + + /// Merge another concurrency, reducing the minimum concurrency or increasing the maximum concurrency to match `other`. + pub fn merge(&mut self, other: &RecommendedConcurrency) { + self.efficient = std::cmp::min(self.efficient, other.efficient); + self.maximum = std::cmp::max(self.maximum, other.maximum); + } +} + impl Codec { /// Create a codec from metadata. /// @@ -181,7 +232,16 @@ pub trait CodecTraits: Send + Sync { /// Traits for both array to array and array to bytes codecs. #[cfg_attr(feature = "async", async_trait::async_trait)] pub trait ArrayCodecTraits: CodecTraits { - /// Encode array with optional parallelism. + /// Return the recommended concurrency for the requested decoded representation. + /// + /// # Errors + /// Returns [`CodecError`] if the decoded representation is not valid for the codec. + fn recommended_concurrency( + &self, + decoded_representation: &ChunkRepresentation, + ) -> Result; + + /// Encode a chunk. /// /// # Errors /// Returns [`CodecError`] if a codec fails or `decoded_value` is incompatible with `decoded_representation`. @@ -189,13 +249,13 @@ pub trait ArrayCodecTraits: CodecTraits { &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &EncodeOptions, ) -> Result, CodecError>; #[cfg(feature = "async")] - /// Asynchronously encode array with optional parallelism. + /// Asynchronously encode a chunk. /// - /// The default implementation calls [`encode_opt`](ArrayCodecTraits::encode_opt) with parallelism disabled. + /// The default implementation calls [`encode_opt`](ArrayCodecTraits::encode_opt). /// /// # Errors /// Returns [`CodecError`] if a codec fails or the decoded output is incompatible with `decoded_representation`. @@ -203,12 +263,12 @@ pub trait ArrayCodecTraits: CodecTraits { &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + options: &EncodeOptions, ) -> Result, CodecError> { - self.encode_opt(decoded_value, decoded_representation, false) + self.encode_opt(decoded_value, decoded_representation, options) } - /// Decode array with optional parallelism. + /// Decode an array. /// /// # Errors /// Returns [`CodecError`] if a codec fails or the decoded output is incompatible with `decoded_representation`. @@ -216,13 +276,13 @@ pub trait ArrayCodecTraits: CodecTraits { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError>; #[cfg(feature = "async")] - /// Asynchronously decode array with optional parallelism. + /// Asynchronously decode an array. /// - /// The default implementation calls [`decode_opt`](ArrayCodecTraits::decode_opt) with parallelism disabled. + /// The default implementation calls [`decode_opt`](ArrayCodecTraits::decode_opt). /// /// # Errors /// Returns [`CodecError`] if a codec fails or the decoded output is incompatible with `decoded_representation`. @@ -230,12 +290,12 @@ pub trait ArrayCodecTraits: CodecTraits { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { - self.decode_opt(encoded_value, decoded_representation, false) + self.decode_opt(encoded_value, decoded_representation, options) } - /// Encode array. + /// Encode an array (default options). /// /// # Errors /// Returns [`CodecError`] if a codec fails or `decoded_value` is incompatible with `decoded_representation`. @@ -244,22 +304,14 @@ pub trait ArrayCodecTraits: CodecTraits { decoded_value: Vec, decoded_representation: &ChunkRepresentation, ) -> Result, CodecError> { - self.encode_opt(decoded_value, decoded_representation, false) + self.encode_opt( + decoded_value, + decoded_representation, + &EncodeOptions::default(), + ) } - /// Encode array using multithreading (if supported). - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails or `decoded_value` is incompatible with `decoded_representation`. - fn par_encode( - &self, - decoded_value: Vec, - decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.encode_opt(decoded_value, decoded_representation, true) - } - - /// Decode array. + /// Decode an array (default options). /// /// # Errors /// Returns [`CodecError`] if a codec fails or the decoded output is incompatible with `decoded_representation`. @@ -268,25 +320,17 @@ pub trait ArrayCodecTraits: CodecTraits { encoded_value: Vec, decoded_representation: &ChunkRepresentation, ) -> Result, CodecError> { - self.decode_opt(encoded_value, decoded_representation, false) - } - - /// Decode array using multithreading (if supported). - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails or the decoded output is incompatible with `decoded_representation`. - fn par_decode( - &self, - encoded_value: Vec, - decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.decode_opt(encoded_value, decoded_representation, true) + self.decode_opt( + encoded_value, + decoded_representation, + &DecodeOptions::default(), + ) } } /// Partial bytes decoder traits. pub trait BytesPartialDecoderTraits: Send + Sync { - /// Partially decode bytes with optional parallelism. + /// Partially decode bytes. /// /// Returns [`None`] if partial decoding of the input handle returns [`None`]. /// @@ -295,22 +339,22 @@ pub trait BytesPartialDecoderTraits: Send + Sync { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError>; - /// Decode all bytes with optional parallelism. + /// Decode all bytes. /// /// Returns [`None`] if partial decoding of the input handle returns [`None`]. /// /// # Errors /// Returns [`CodecError`] if a codec fails. - fn decode_opt(&self, parallel: bool) -> Result { + fn decode_opt(&self, options: &PartialDecodeOptions) -> Result { Ok(self - .partial_decode_opt(&[ByteRange::FromStart(0, None)], parallel)? + .partial_decode_opt(&[ByteRange::FromStart(0, None)], options)? .map(|mut v| v.remove(0))) } - /// Partially decode bytes. + /// Partially decode bytes (default options). /// /// Returns [`None`] if partial decoding of the input handle returns [`None`]. /// @@ -320,40 +364,17 @@ pub trait BytesPartialDecoderTraits: Send + Sync { &self, byte_ranges: &[ByteRange], ) -> Result>>, CodecError> { - self.partial_decode_opt(byte_ranges, false) - } - - /// Partially decode bytes using multithreading if applicable. - /// - /// Returns [`None`] if partial decoding of the input handle returns [`None`]. - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails or a byte range is invalid. - fn par_partial_decode( - &self, - byte_ranges: &[ByteRange], - ) -> Result>>, CodecError> { - self.partial_decode_opt(byte_ranges, true) + self.partial_decode_opt(byte_ranges, &PartialDecodeOptions::default()) } - /// Decode all bytes. + /// Decode all bytes (default options). /// /// Returns [`None`] if decoding of the input handle returns [`None`]. /// /// # Errors /// Returns [`CodecError`] if a codec fails. fn decode(&self) -> Result { - self.decode_opt(false) - } - - /// Decode all bytes using multithreading (if supported). - /// - /// Returns [`None`] if decoding of the input handle returns [`None`]. - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails. - fn par_decode(&self) -> Result { - self.decode_opt(true) + self.decode_opt(&PartialDecodeOptions::default()) } } @@ -361,7 +382,7 @@ pub trait BytesPartialDecoderTraits: Send + Sync { /// Asynchronous partial bytes decoder traits. #[async_trait::async_trait] pub trait AsyncBytesPartialDecoderTraits: Send + Sync { - /// Partially decode bytes with optional parallelism. + /// Partially decode bytes. /// /// Returns [`None`] if partial decoding of the input handle returns [`None`]. /// @@ -370,23 +391,23 @@ pub trait AsyncBytesPartialDecoderTraits: Send + Sync { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError>; - /// Decode all bytes with optional parallelism. + /// Decode all bytes. /// /// Returns [`None`] if partial decoding of the input handle returns [`None`]. /// /// # Errors /// Returns [`CodecError`] if a codec fails. - async fn decode_opt(&self, parallel: bool) -> Result { + async fn decode_opt(&self, options: &PartialDecodeOptions) -> Result { Ok(self - .partial_decode_opt(&[ByteRange::FromStart(0, None)], parallel) + .partial_decode_opt(&[ByteRange::FromStart(0, None)], options) .await? .map(|mut v| v.remove(0))) } - /// Partially decode bytes. + /// Partially decode bytes with defaullt options. /// /// Returns [`None`] if partial decoding of the input handle returns [`None`]. /// @@ -396,40 +417,18 @@ pub trait AsyncBytesPartialDecoderTraits: Send + Sync { &self, byte_ranges: &[ByteRange], ) -> Result>>, CodecError> { - self.partial_decode_opt(byte_ranges, false).await - } - - /// Partially decode bytes using multithreading if applicable. - /// - /// Returns [`None`] if partial decoding of the input handle returns [`None`]. - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails or a byte range is invalid. - async fn par_partial_decode( - &self, - byte_ranges: &[ByteRange], - ) -> Result>>, CodecError> { - self.partial_decode_opt(byte_ranges, true).await + self.partial_decode_opt(byte_ranges, &PartialDecodeOptions::default()) + .await } - /// Decode all bytes. + /// Decode all bytes (default options). /// /// Returns [`None`] if decoding of the input handle returns [`None`]. /// /// # Errors /// Returns [`CodecError`] if a codec fails. async fn decode(&self) -> Result { - self.decode_opt(false).await - } - - /// Decode all bytes using multithreading (if supported). - /// - /// Returns [`None`] if decoding of the input handle returns [`None`]. - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails. - async fn par_decode(&self) -> Result { - self.decode_opt(true).await + self.decode_opt(&PartialDecodeOptions::default()).await } } @@ -444,10 +443,10 @@ pub trait ArrayPartialDecoderTraits: Send + Sync { fn partial_decode_opt( &self, array_subsets: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError>; - /// Partially decode an array. + /// Partially decode an array (default options). /// /// If the inner `input_handle` is a bytes decoder and partial decoding returns [`None`], then the array subsets have the fill value. /// @@ -455,20 +454,7 @@ pub trait ArrayPartialDecoderTraits: Send + Sync { /// /// Returns [`CodecError`] if a codec fails or an array subset is invalid. fn partial_decode(&self, array_subsets: &[ArraySubset]) -> Result>, CodecError> { - self.partial_decode_opt(array_subsets, false) - } - - /// Partially decode an array using multithreading (if supported). - /// - /// If the inner `input_handle` is a bytes decoder and partial decoding returns [`None`], then the array subsets have the fill value. - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails or an array subset is invalid. - fn par_partial_decode( - &self, - array_subsets: &[ArraySubset], - ) -> Result>, CodecError> { - self.partial_decode_opt(array_subsets, true) + self.partial_decode_opt(array_subsets, &PartialDecodeOptions::default()) } } @@ -485,10 +471,10 @@ pub trait AsyncArrayPartialDecoderTraits: Send + Sync { async fn partial_decode_opt( &self, array_subsets: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError>; - /// Partially decode an array. + /// Partially decode an array (default options). /// /// If the inner `input_handle` is a bytes decoder and partial decoding returns [`None`], then the array subsets have the fill value. /// @@ -499,20 +485,8 @@ pub trait AsyncArrayPartialDecoderTraits: Send + Sync { &self, array_subsets: &[ArraySubset], ) -> Result>, CodecError> { - self.partial_decode_opt(array_subsets, false).await - } - - /// Partially decode an array using multithreading (if supported). - /// - /// If the inner `input_handle` is a bytes decoder and partial decoding returns [`None`], then the array subsets have the fill value. - /// - /// # Errors - /// Returns [`CodecError`] if a codec fails or an array subset is invalid. - async fn par_partial_decode( - &self, - array_subsets: &[ArraySubset], - ) -> Result>, CodecError> { - self.partial_decode_opt(array_subsets, true).await + self.partial_decode_opt(array_subsets, &PartialDecodeOptions::default()) + .await } } @@ -533,7 +507,7 @@ impl BytesPartialDecoderTraits for StoragePartialDecoder { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _options: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(self .storage @@ -562,7 +536,7 @@ impl AsyncBytesPartialDecoderTraits for AsyncStoragePartialDecoder { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _options: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(self .storage @@ -576,35 +550,6 @@ impl AsyncBytesPartialDecoderTraits for AsyncStoragePartialDecoder { pub trait ArrayToArrayCodecTraits: ArrayCodecTraits + dyn_clone::DynClone + core::fmt::Debug { - /// Initialise a partial decoder with optional parallelism. - /// - /// `parallel` only affects parallelism on initialisation, which is irrelevant for most codecs. - /// Parallel partial decoding support is independent of how the partial decoder is initialised. - /// - /// # Errors - /// Returns a [`CodecError`] if initialisation fails. - fn partial_decoder_opt<'a>( - &'a self, - input_handle: Box, - decoded_representation: &ChunkRepresentation, - parallel: bool, - ) -> Result, CodecError>; - - #[cfg(feature = "async")] - /// Initialise an asynchronous partial decoder with optional parallelism. - /// - /// `parallel` only affects parallelism on initialisation, which is irrelevant for most codecs. - /// Parallel partial decoding support is independent of how the partial decoder is initialised. - /// - /// # Errors - /// Returns a [`CodecError`] if initialisation fails. - async fn async_partial_decoder_opt<'a>( - &'a self, - input_handle: Box, - decoded_representation: &ChunkRepresentation, - parallel: bool, - ) -> Result, CodecError>; - /// Returns the size of the encoded representation given a size of the decoded representation. /// /// # Errors @@ -617,54 +562,65 @@ pub trait ArrayToArrayCodecTraits: /// Initialise a partial decoder. /// + /// `parallel` only affects parallelism on initialisation, which is irrelevant for most codecs. + /// Parallel partial decoding support is independent of how the partial decoder is initialised. + /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - fn partial_decoder<'a>( + fn partial_decoder_opt<'a>( &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.partial_decoder_opt(input_handle, decoded_representation, false) - } + options: &PartialDecoderOptions, + ) -> Result, CodecError>; - /// Initialise a partial decoder with multithreading (where supported). + /// Initialise a partial decoder (default options). /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - fn par_partial_decoder<'a>( + fn partial_decoder<'a>( &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, ) -> Result, CodecError> { - self.partial_decoder_opt(input_handle, decoded_representation, true) + self.partial_decoder_opt( + input_handle, + decoded_representation, + &PartialDecoderOptions::default(), + ) } #[cfg(feature = "async")] /// Initialise an asynchronous partial decoder. /// + /// `parallel` only affects parallelism on initialisation, which is irrelevant for most codecs. + /// Parallel partial decoding support is independent of how the partial decoder is initialised. + /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_partial_decoder<'a>( + async fn async_partial_decoder_opt<'a>( &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.async_partial_decoder_opt(input_handle, decoded_representation, false) - .await - } + options: &PartialDecoderOptions, + ) -> Result, CodecError>; #[cfg(feature = "async")] - /// Initialise an asynchronous partial decoder with multithreading (where supported). + /// Initialise an asynchronous partial decoder (default options). /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_par_partial_decoder<'a>( + async fn async_partial_decoder<'a>( &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, ) -> Result, CodecError> { - self.async_partial_decoder_opt(input_handle, decoded_representation, true) - .await + self.async_partial_decoder_opt( + input_handle, + decoded_representation, + &PartialDecoderOptions::default(), + ) + .await } } @@ -675,33 +631,9 @@ dyn_clone::clone_trait_object!(ArrayToArrayCodecTraits); pub trait ArrayToBytesCodecTraits: ArrayCodecTraits + dyn_clone::DynClone + core::fmt::Debug { - /// Initialise a partial decoder with optional parallelism. - /// - /// # Errors - /// Returns a [`CodecError`] if initialisation fails. - fn partial_decoder_opt<'a>( - &'a self, - input_handle: Box, - decoded_representation: &ChunkRepresentation, - parallel: bool, - ) -> Result, CodecError>; - - #[cfg(feature = "async")] - /// Initialise an asynchronous partial decoder with optional parallelism. - /// - /// # Errors - /// Returns a [`CodecError`] if initialisation fails. - async fn async_partial_decoder_opt<'a>( - &'a self, - mut input_handle: Box, - decoded_representation: &ChunkRepresentation, - parallel: bool, - ) -> Result, CodecError>; - /// Returns the size of the encoded representation given a size of the decoded representation. /// /// # Errors - /// /// Returns a [`CodecError`] if the decoded representation is not supported by this codec. fn compute_encoded_size( &self, @@ -712,52 +644,57 @@ pub trait ArrayToBytesCodecTraits: /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - fn partial_decoder<'a>( + fn partial_decoder_opt<'a>( &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.partial_decoder_opt(input_handle, decoded_representation, false) - } + options: &PartialDecoderOptions, + ) -> Result, CodecError>; - /// Initialise a partial decoder with multithreading (where supported). + #[cfg(feature = "async")] + /// Initialise an asynchronous partial decoder. /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - fn par_partial_decoder<'a>( + async fn async_partial_decoder_opt<'a>( &'a self, - input_handle: Box, + mut input_handle: Box, decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.partial_decoder_opt(input_handle, decoded_representation, true) - } + options: &PartialDecoderOptions, + ) -> Result, CodecError>; - #[cfg(feature = "async")] - /// Initialise an asynchronous partial decoder. + /// Initialise a partial decoder (default options). /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_partial_decoder<'a>( + fn partial_decoder<'a>( &'a self, - input_handle: Box, + input_handle: Box, decoded_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - self.async_partial_decoder_opt(input_handle, decoded_representation, false) - .await + ) -> Result, CodecError> { + self.partial_decoder_opt( + input_handle, + decoded_representation, + &PartialDecoderOptions::default(), + ) } #[cfg(feature = "async")] - /// Initialise an asynchronous partial decoder with multithreading (where supported). + /// Initialise an asynchronous partial decoder (default options). /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_par_partial_decoder<'a>( + async fn async_partial_decoder<'a>( &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, ) -> Result, CodecError> { - self.async_partial_decoder_opt(input_handle, decoded_representation, true) - .await + self.async_partial_decoder_opt( + input_handle, + decoded_representation, + &PartialDecoderOptions::default(), + ) + .await } } @@ -766,28 +703,40 @@ dyn_clone::clone_trait_object!(ArrayToBytesCodecTraits); /// Traits for bytes to bytes codecs. #[cfg_attr(feature = "async", async_trait::async_trait)] pub trait BytesToBytesCodecTraits: CodecTraits + dyn_clone::DynClone + core::fmt::Debug { - /// Encode bytes with optional parallelism. + /// Return the maximum internal concurrency supported for the requested decoded representation. /// /// # Errors - /// Returns [`CodecError`] if a codec fails. - fn encode_opt(&self, decoded_value: Vec, parallel: bool) -> Result, CodecError>; + /// Returns [`CodecError`] if the decoded representation is not valid for the codec. + fn recommended_concurrency( + &self, + decoded_representation: &BytesRepresentation, + ) -> Result; - #[cfg(feature = "async")] - /// Asynchronously encode bytes with optional parallelism. - /// - /// The default implementation calls [`encode_opt`](BytesToBytesCodecTraits::encode_opt) with parallelism disabled. + /// Returns the size of the encoded representation given a size of the decoded representation. + fn compute_encoded_size( + &self, + decoded_representation: &BytesRepresentation, + ) -> BytesRepresentation; + + /// Encode chunk bytes. /// /// # Errors /// Returns [`CodecError`] if a codec fails. - async fn async_encode_opt( + fn encode_opt( &self, decoded_value: Vec, - _parallel: bool, - ) -> Result, CodecError> { - self.encode_opt(decoded_value, false) + options: &EncodeOptions, + ) -> Result, CodecError>; + + /// Encode chunk bytes (default options). + /// + /// # Errors + /// Returns [`CodecError`] if a codec fails. + fn encode(&self, decoded_value: Vec) -> Result, CodecError> { + self.encode_opt(decoded_value, &EncodeOptions::default()) } - /// Decode bytes with optional parallelism. + /// Decode chunk bytes. // /// # Errors /// Returns [`CodecError`] if a codec fails. @@ -795,26 +744,26 @@ pub trait BytesToBytesCodecTraits: CodecTraits + dyn_clone::DynClone + core::fmt &self, encoded_value: Vec, decoded_representation: &BytesRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError>; - #[cfg(feature = "async")] - /// Asynchronously decode bytes with optional parallelism. - /// - /// The default implementation calls [`decode_opt`](BytesToBytesCodecTraits::decode_opt) with parallelism disabled. + /// Decode bytes (default options). /// /// # Errors /// Returns [`CodecError`] if a codec fails. - async fn async_decode_opt( + fn decode( &self, encoded_value: Vec, decoded_representation: &BytesRepresentation, - _parallel: bool, ) -> Result, CodecError> { - self.decode_opt(encoded_value, decoded_representation, false) + self.decode_opt( + encoded_value, + decoded_representation, + &DecodeOptions::default(), + ) } - /// Initialises a partial decoder with optional parallelism. + /// Initialises a partial decoder. /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. @@ -822,89 +771,86 @@ pub trait BytesToBytesCodecTraits: CodecTraits + dyn_clone::DynClone + core::fmt &'a self, input_handle: Box, decoded_representation: &BytesRepresentation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, CodecError>; - #[cfg(feature = "async")] - /// Initialises an asynchronous partial decoder with optional parallelism. + /// Initialises a partial decoder (default options). /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_partial_decoder_opt<'a>( + fn partial_decoder<'a>( &'a self, - input_handle: Box, - decoded_representation: &BytesRepresentation, - parallel: bool, - ) -> Result, CodecError>; - - /// Returns the size of the encoded representation given a size of the decoded representation. - fn compute_encoded_size( - &self, + input_handle: Box, decoded_representation: &BytesRepresentation, - ) -> BytesRepresentation; + ) -> Result, CodecError> { + self.partial_decoder_opt( + input_handle, + decoded_representation, + &PartialDecodeOptions::default(), + ) + } - /// Encode bytes. + #[cfg(feature = "async")] + /// Asynchronously encode chunk bytes. + /// + /// The default implementation calls [`encode_opt`](BytesToBytesCodecTraits::encode_opt). /// /// # Errors /// Returns [`CodecError`] if a codec fails. - fn encode(&self, decoded_value: Vec) -> Result, CodecError> { - self.encode_opt(decoded_value, false) + async fn async_encode_opt( + &self, + decoded_value: Vec, + options: &EncodeOptions, + ) -> Result, CodecError> { + self.encode_opt(decoded_value, options) } - /// Encode bytes using using multithreading (if supported). + #[cfg(feature = "async")] + /// Asynchronously encode chunk bytes (default options). + /// + /// The default implementation calls [`async_encode_opt`](BytesToBytesCodecTraits::async_encode_opt). /// /// # Errors /// Returns [`CodecError`] if a codec fails. - fn par_encode(&self, decoded_value: Vec) -> Result, CodecError> { - self.encode_opt(decoded_value, true) + async fn async_encode(&self, decoded_value: Vec) -> Result, CodecError> { + self.async_encode_opt(decoded_value, &EncodeOptions::default()) + .await } - /// Decode bytes. + #[cfg(feature = "async")] + /// Asynchronously decode chunk bytes. + /// + /// The default implementation calls [`decode_opt`](BytesToBytesCodecTraits::decode_opt). /// /// # Errors /// Returns [`CodecError`] if a codec fails. - fn decode( + async fn async_decode_opt( &self, encoded_value: Vec, decoded_representation: &BytesRepresentation, + options: &DecodeOptions, ) -> Result, CodecError> { - self.decode_opt(encoded_value, decoded_representation, false) + self.decode_opt(encoded_value, decoded_representation, options) } - /// Decode bytes using using multithreading (if supported). + #[cfg(feature = "async")] + /// Asynchronously decode chunk bytes. + /// + /// The default implementation calls [`decode_opt`](BytesToBytesCodecTraits::decode_opt). /// /// # Errors /// Returns [`CodecError`] if a codec fails. - fn par_decode( + async fn async_decode( &self, encoded_value: Vec, decoded_representation: &BytesRepresentation, ) -> Result, CodecError> { - self.decode_opt(encoded_value, decoded_representation, true) - } - - /// Initialises a partial decoder. - /// - /// # Errors - /// Returns a [`CodecError`] if initialisation fails. - fn partial_decoder<'a>( - &'a self, - input_handle: Box, - decoded_representation: &BytesRepresentation, - ) -> Result, CodecError> { - self.partial_decoder_opt(input_handle, decoded_representation, false) - } - - /// Initialise a partial decoder with multithreading (where supported). - /// - /// # Errors - /// Returns a [`CodecError`] if initialisation fails. - fn par_partial_decoder<'a>( - &'a self, - input_handle: Box, - decoded_representation: &BytesRepresentation, - ) -> Result, CodecError> { - self.partial_decoder_opt(input_handle, decoded_representation, true) + self.async_decode_opt( + encoded_value, + decoded_representation, + &DecodeOptions::default(), + ) + .await } #[cfg(feature = "async")] @@ -912,27 +858,29 @@ pub trait BytesToBytesCodecTraits: CodecTraits + dyn_clone::DynClone + core::fmt /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_partial_decoder<'a>( + async fn async_partial_decoder_opt<'a>( &'a self, input_handle: Box, decoded_representation: &BytesRepresentation, - ) -> Result, CodecError> { - self.async_partial_decoder_opt(input_handle, decoded_representation, false) - .await - } + options: &PartialDecoderOptions, + ) -> Result, CodecError>; #[cfg(feature = "async")] - /// Initialise an asynchronous partial decoder with multithreading (where supported). + /// Initialises an asynchronous partial decoder (default options). /// /// # Errors /// Returns a [`CodecError`] if initialisation fails. - async fn async_par_partial_decoder<'a>( + async fn async_partial_decoder<'a>( &'a self, input_handle: Box, decoded_representation: &BytesRepresentation, ) -> Result, CodecError> { - self.async_partial_decoder_opt(input_handle, decoded_representation, true) - .await + self.async_partial_decoder_opt( + input_handle, + decoded_representation, + &PartialDecoderOptions::default(), + ) + .await } } @@ -942,7 +890,7 @@ impl BytesPartialDecoderTraits for std::io::Cursor<&[u8]> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _parallel: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(Some(extract_byte_ranges_read_seek( &mut self.clone(), @@ -955,7 +903,7 @@ impl BytesPartialDecoderTraits for std::io::Cursor> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _parallel: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(Some(extract_byte_ranges_read_seek( &mut self.clone(), @@ -970,7 +918,7 @@ impl AsyncBytesPartialDecoderTraits for std::io::Cursor<&[u8]> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _parallel: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(Some(extract_byte_ranges_read_seek( &mut self.clone(), @@ -985,7 +933,7 @@ impl AsyncBytesPartialDecoderTraits for std::io::Cursor> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _parallel: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(Some(extract_byte_ranges_read_seek( &mut self.clone(), diff --git a/src/array/codec/array_to_array/bitround/bitround_codec.rs b/src/array/codec/array_to_array/bitround/bitround_codec.rs index 057b0f03..bae1e8b5 100644 --- a/src/array/codec/array_to_array/bitround/bitround_codec.rs +++ b/src/array/codec/array_to_array/bitround/bitround_codec.rs @@ -2,7 +2,8 @@ use crate::{ array::{ codec::{ ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToArrayCodecTraits, CodecError, - CodecTraits, + CodecTraits, DecodeOptions, EncodeOptions, PartialDecoderOptions, + RecommendedConcurrency, }, ChunkRepresentation, DataType, }, @@ -60,11 +61,19 @@ impl CodecTraits for BitroundCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for BitroundCodec { + fn recommended_concurrency( + &self, + _decoded_representation: &ChunkRepresentation, + ) -> Result { + // TODO: bitround is well suited to multithread, when is it optimal to kick in? + Ok(RecommendedConcurrency::one()) + } + fn encode_opt( &self, mut decoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &EncodeOptions, ) -> Result, CodecError> { round_bytes( &mut decoded_value, @@ -78,7 +87,7 @@ impl ArrayCodecTraits for BitroundCodec { &self, encoded_value: Vec, _decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { Ok(encoded_value) } @@ -90,7 +99,7 @@ impl ArrayToArrayCodecTraits for BitroundCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( bitround_partial_decoder::BitroundPartialDecoder::new( @@ -106,7 +115,7 @@ impl ArrayToArrayCodecTraits for BitroundCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( bitround_partial_decoder::AsyncBitroundPartialDecoder::new( diff --git a/src/array/codec/array_to_array/bitround/bitround_partial_decoder.rs b/src/array/codec/array_to_array/bitround/bitround_partial_decoder.rs index c07fa290..17fe3063 100644 --- a/src/array/codec/array_to_array/bitround/bitround_partial_decoder.rs +++ b/src/array/codec/array_to_array/bitround/bitround_partial_decoder.rs @@ -1,6 +1,6 @@ use crate::{ array::{ - codec::{ArrayPartialDecoderTraits, CodecError}, + codec::{ArrayPartialDecoderTraits, CodecError, PartialDecodeOptions}, DataType, }, array_subset::ArraySubset, @@ -54,11 +54,11 @@ impl ArrayPartialDecoderTraits for BitroundPartialDecoder<'_> { fn partial_decode_opt( &self, array_subsets: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let mut bytes = self .input_handle - .partial_decode_opt(array_subsets, parallel)?; + .partial_decode_opt(array_subsets, options)?; for bytes in &mut bytes { round_bytes(bytes, &self.data_type, self.keepbits)?; @@ -115,11 +115,11 @@ impl AsyncArrayPartialDecoderTraits for AsyncBitroundPartialDecoder<'_> { async fn partial_decode_opt( &self, array_subsets: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let mut bytes = self .input_handle - .partial_decode_opt(array_subsets, parallel) + .partial_decode_opt(array_subsets, options) .await?; for bytes in &mut bytes { diff --git a/src/array/codec/array_to_array/transpose/transpose_codec.rs b/src/array/codec/array_to_array/transpose/transpose_codec.rs index d5394720..36fce5bf 100644 --- a/src/array/codec/array_to_array/transpose/transpose_codec.rs +++ b/src/array/codec/array_to_array/transpose/transpose_codec.rs @@ -5,7 +5,8 @@ use crate::{ array::{ codec::{ ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToArrayCodecTraits, CodecError, - CodecTraits, + CodecTraits, DecodeOptions, EncodeOptions, PartialDecoderOptions, + RecommendedConcurrency, }, ChunkRepresentation, }, @@ -76,7 +77,7 @@ impl ArrayToArrayCodecTraits for TransposeCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( super::transpose_partial_decoder::TransposePartialDecoder::new( @@ -92,7 +93,7 @@ impl ArrayToArrayCodecTraits for TransposeCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( super::transpose_partial_decoder::AsyncTransposePartialDecoder::new( @@ -120,11 +121,19 @@ impl ArrayToArrayCodecTraits for TransposeCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for TransposeCodec { + fn recommended_concurrency( + &self, + _decoded_representation: &ChunkRepresentation, + ) -> Result { + // TODO: This could be increased, need to implement `transpose_array` without ndarray + Ok(RecommendedConcurrency::one()) + } + fn encode_opt( &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &EncodeOptions, ) -> Result, CodecError> { if decoded_value.len() as u64 != decoded_representation.size() { return Err(CodecError::UnexpectedChunkDecodedSize( @@ -153,7 +162,7 @@ impl ArrayCodecTraits for TransposeCodec { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { let order_decode = calculate_order_decode(&self.order, decoded_representation.shape().len()); diff --git a/src/array/codec/array_to_array/transpose/transpose_partial_decoder.rs b/src/array/codec/array_to_array/transpose/transpose_partial_decoder.rs index 79eab557..0c825356 100644 --- a/src/array/codec/array_to_array/transpose/transpose_partial_decoder.rs +++ b/src/array/codec/array_to_array/transpose/transpose_partial_decoder.rs @@ -1,6 +1,6 @@ use super::{calculate_order_decode, permute, transpose_array, TransposeOrder}; use crate::array::{ - codec::{ArrayPartialDecoderTraits, ArraySubset, CodecError}, + codec::{ArrayPartialDecoderTraits, ArraySubset, CodecError, PartialDecodeOptions}, ChunkRepresentation, }; @@ -33,7 +33,7 @@ impl ArrayPartialDecoderTraits for TransposePartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { // Get transposed array subsets let mut decoded_regions_transposed = Vec::with_capacity(decoded_regions.len()); @@ -46,7 +46,7 @@ impl ArrayPartialDecoderTraits for TransposePartialDecoder<'_> { } let mut encoded_value = self .input_handle - .partial_decode_opt(&decoded_regions_transposed, parallel)?; + .partial_decode_opt(&decoded_regions_transposed, options)?; // Reverse the transpose on each subset let order_decode = @@ -99,7 +99,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncTransposePartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { // Get transposed array subsets let mut decoded_regions_transposed = Vec::with_capacity(decoded_regions.len()); @@ -112,7 +112,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncTransposePartialDecoder<'_> { } let mut encoded_value = self .input_handle - .partial_decode_opt(&decoded_regions_transposed, parallel) + .partial_decode_opt(&decoded_regions_transposed, options) .await?; // Reverse the transpose on each subset diff --git a/src/array/codec/array_to_bytes/bytes/bytes_codec.rs b/src/array/codec/array_to_bytes/bytes/bytes_codec.rs index 581bd192..d5e26030 100644 --- a/src/array/codec/array_to_bytes/bytes/bytes_codec.rs +++ b/src/array/codec/array_to_bytes/bytes/bytes_codec.rs @@ -4,7 +4,8 @@ use crate::{ array::{ codec::{ ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, - BytesPartialDecoderTraits, CodecError, CodecTraits, + BytesPartialDecoderTraits, CodecError, CodecTraits, DecodeOptions, EncodeOptions, + PartialDecoderOptions, RecommendedConcurrency, }, BytesRepresentation, ChunkRepresentation, }, @@ -104,11 +105,31 @@ impl CodecTraits for BytesCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for BytesCodec { + fn recommended_concurrency( + &self, + _decoded_representation: &ChunkRepresentation, + ) -> Result { + // TODO: Recomment > 1 if endianness needs changing and input is sufficiently large + // if let Some(endian) = &self.endian { + // if !endian.is_native() { + // FIXME: Support parallel + // let min_elements_per_thread = 32768; // 32^3 + // unsafe { + // NonZeroU64::new_unchecked( + // (decoded_representation.num_elements() + min_elements_per_thread - 1) + // / min_elements_per_thread, + // ) + // } + // } + // } + Ok(RecommendedConcurrency::one()) + } + fn encode_opt( &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &EncodeOptions, ) -> Result, CodecError> { self.do_encode_or_decode(decoded_value, decoded_representation) } @@ -117,7 +138,7 @@ impl ArrayCodecTraits for BytesCodec { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { self.do_encode_or_decode(encoded_value, decoded_representation) } @@ -129,7 +150,7 @@ impl ArrayToBytesCodecTraits for BytesCodec { &self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(bytes_partial_decoder::BytesPartialDecoder::new( input_handle, @@ -143,7 +164,7 @@ impl ArrayToBytesCodecTraits for BytesCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( bytes_partial_decoder::AsyncBytesPartialDecoder::new( diff --git a/src/array/codec/array_to_bytes/bytes/bytes_partial_decoder.rs b/src/array/codec/array_to_bytes/bytes/bytes_partial_decoder.rs index d7142ef9..188aa526 100644 --- a/src/array/codec/array_to_bytes/bytes/bytes_partial_decoder.rs +++ b/src/array/codec/array_to_bytes/bytes/bytes_partial_decoder.rs @@ -1,6 +1,9 @@ use crate::{ array::{ - codec::{ArrayPartialDecoderTraits, ArraySubset, BytesPartialDecoderTraits, CodecError}, + codec::{ + ArrayPartialDecoderTraits, ArraySubset, BytesPartialDecoderTraits, CodecError, + PartialDecodeOptions, + }, ChunkRepresentation, }, array_subset::IncompatibleArraySubsetAndShapeError, @@ -37,7 +40,7 @@ impl ArrayPartialDecoderTraits for BytesPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let mut bytes = Vec::with_capacity(decoded_regions.len()); let chunk_shape = self.decoded_representation.shape_u64(); @@ -55,7 +58,7 @@ impl ArrayPartialDecoderTraits for BytesPartialDecoder<'_> { // Decode let decoded = self .input_handle - .partial_decode_opt(&byte_ranges, parallel)?; + .partial_decode_opt(&byte_ranges, options)?; let bytes_subset = decoded.map_or_else( || { @@ -114,7 +117,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncBytesPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let mut bytes = Vec::with_capacity(decoded_regions.len()); let chunk_shape = self.decoded_representation.shape_u64(); @@ -132,7 +135,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncBytesPartialDecoder<'_> { // Decode let decoded = self .input_handle - .partial_decode_opt(&byte_ranges, parallel) + .partial_decode_opt(&byte_ranges, options) .await?; let bytes_subset = decoded.map_or_else( diff --git a/src/array/codec/array_to_bytes/codec_chain.rs b/src/array/codec/array_to_bytes/codec_chain.rs index b3b1c567..bc4c35c3 100644 --- a/src/array/codec/array_to_bytes/codec_chain.rs +++ b/src/array/codec/array_to_bytes/codec_chain.rs @@ -1,12 +1,15 @@ //! An array to bytes codec formed by joining an array to array sequence, array to bytes, and bytes to bytes sequence of codecs. +use std::num::NonZeroU64; + use crate::{ array::{ codec::{ partial_decoder_cache::{ArrayPartialDecoderCache, BytesPartialDecoderCache}, ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToArrayCodecTraits, ArrayToBytesCodecTraits, BytesPartialDecoderTraits, BytesToBytesCodecTraits, Codec, - CodecError, CodecTraits, + CodecError, CodecTraits, DecodeOptions, EncodeOptions, PartialDecoderOptions, + RecommendedConcurrency, }, BytesRepresentation, ChunkRepresentation, }, @@ -223,7 +226,7 @@ impl ArrayToBytesCodecTraits for CodecChain { &'a self, mut input_handle: Box, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, CodecError> { let array_representations = self.get_array_representations(decoded_representation.clone())?; @@ -236,22 +239,22 @@ impl ArrayToBytesCodecTraits for CodecChain { bytes_representations.iter().rev().skip(1), ) { if Some(codec_index) == self.cache_index { - input_handle = Box::new(BytesPartialDecoderCache::new(&*input_handle, parallel)?); + input_handle = Box::new(BytesPartialDecoderCache::new(&*input_handle, options)?); } codec_index += 1; input_handle = - codec.partial_decoder_opt(input_handle, bytes_representation, parallel)?; + codec.partial_decoder_opt(input_handle, bytes_representation, options)?; } if Some(codec_index) == self.cache_index { - input_handle = Box::new(BytesPartialDecoderCache::new(&*input_handle, parallel)?); + input_handle = Box::new(BytesPartialDecoderCache::new(&*input_handle, options)?); }; let mut input_handle = { let array_representation = array_representations.last().unwrap(); let codec = &self.array_to_bytes; codec_index += 1; - codec.partial_decoder_opt(input_handle, array_representation, parallel)? + codec.partial_decoder_opt(input_handle, array_representation, options)? }; for (codec, array_representation) in std::iter::zip( @@ -262,19 +265,19 @@ impl ArrayToBytesCodecTraits for CodecChain { input_handle = Box::new(ArrayPartialDecoderCache::new( &*input_handle, array_representation.clone(), - parallel, + options, )?); } codec_index += 1; input_handle = - codec.partial_decoder_opt(input_handle, array_representation, parallel)?; + codec.partial_decoder_opt(input_handle, array_representation, options)?; } if Some(codec_index) == self.cache_index { input_handle = Box::new(ArrayPartialDecoderCache::new( &*input_handle, array_representations.first().unwrap().clone(), - parallel, + options, )?); } @@ -286,7 +289,7 @@ impl ArrayToBytesCodecTraits for CodecChain { &'a self, mut input_handle: Box, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, CodecError> { let array_representations = self.get_array_representations(decoded_representation.clone())?; @@ -300,17 +303,17 @@ impl ArrayToBytesCodecTraits for CodecChain { ) { if Some(codec_index) == self.cache_index { input_handle = - Box::new(BytesPartialDecoderCache::async_new(&*input_handle, parallel).await?); + Box::new(BytesPartialDecoderCache::async_new(&*input_handle, options).await?); } codec_index += 1; input_handle = codec - .async_partial_decoder_opt(input_handle, bytes_representation, parallel) + .async_partial_decoder_opt(input_handle, bytes_representation, options) .await?; } if Some(codec_index) == self.cache_index { input_handle = - Box::new(BytesPartialDecoderCache::async_new(&*input_handle, parallel).await?); + Box::new(BytesPartialDecoderCache::async_new(&*input_handle, options).await?); }; let mut input_handle = { @@ -318,7 +321,7 @@ impl ArrayToBytesCodecTraits for CodecChain { let codec = &self.array_to_bytes; codec_index += 1; codec - .async_partial_decoder_opt(input_handle, array_representation, parallel) + .async_partial_decoder_opt(input_handle, array_representation, options) .await? }; @@ -331,14 +334,14 @@ impl ArrayToBytesCodecTraits for CodecChain { ArrayPartialDecoderCache::async_new( &*input_handle, array_representation.clone(), - parallel, + options, ) .await?, ); } codec_index += 1; input_handle = codec - .async_partial_decoder_opt(input_handle, array_representation, parallel) + .async_partial_decoder_opt(input_handle, array_representation, options) .await?; } @@ -347,7 +350,7 @@ impl ArrayToBytesCodecTraits for CodecChain { ArrayPartialDecoderCache::async_new( &*input_handle, array_representations.first().unwrap().clone(), - parallel, + options, ) .await?, ); @@ -379,11 +382,52 @@ impl ArrayToBytesCodecTraits for CodecChain { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for CodecChain { + fn recommended_concurrency( + &self, + decoded_representation: &ChunkRepresentation, + ) -> Result { + let mut recommended_concurrency = + RecommendedConcurrency::new(unsafe { NonZeroU64::new_unchecked(u64::MAX) }, unsafe { + NonZeroU64::new_unchecked(1) + }); + + let array_representations = + self.get_array_representations(decoded_representation.clone())?; + let bytes_representations = + self.get_bytes_representations(array_representations.last().unwrap())?; + + // bytes->bytes + for (codec, bytes_representation) in std::iter::zip( + self.bytes_to_bytes.iter().rev(), + bytes_representations.iter().rev().skip(1), + ) { + recommended_concurrency.merge(&codec.recommended_concurrency(bytes_representation)?); + } + + recommended_concurrency.merge( + &self + .array_to_bytes + .recommended_concurrency(array_representations.last().unwrap())?, + ); + + // array->array + for (codec, array_representation) in std::iter::zip( + self.array_to_array.iter().rev(), + array_representations.iter().rev().skip(1), + ) { + recommended_concurrency.merge(&codec.recommended_concurrency(array_representation)?); + } + + // FIXME: Check if maximum < efficient, if so set maximum to efficient + + Ok(recommended_concurrency) + } + fn encode_opt( &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &EncodeOptions, ) -> Result, CodecError> { if decoded_value.len() as u64 != decoded_representation.size() { return Err(CodecError::UnexpectedChunkDecodedSize( @@ -397,21 +441,21 @@ impl ArrayCodecTraits for CodecChain { let mut value = decoded_value; // array->array for codec in &self.array_to_array { - value = codec.encode_opt(value, &decoded_representation, parallel)?; + value = codec.encode_opt(value, &decoded_representation, options)?; decoded_representation = codec.compute_encoded_size(&decoded_representation)?; } // array->bytes value = self .array_to_bytes - .encode_opt(value, &decoded_representation, parallel)?; + .encode_opt(value, &decoded_representation, options)?; let mut decoded_representation = self .array_to_bytes .compute_encoded_size(&decoded_representation)?; // bytes->bytes for codec in &self.bytes_to_bytes { - value = codec.encode_opt(value, parallel)?; + value = codec.encode_opt(value, options)?; decoded_representation = codec.compute_encoded_size(&decoded_representation); } @@ -422,7 +466,7 @@ impl ArrayCodecTraits for CodecChain { &self, mut encoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { let array_representations = self.get_array_representations(decoded_representation.clone())?; @@ -434,14 +478,14 @@ impl ArrayCodecTraits for CodecChain { self.bytes_to_bytes.iter().rev(), bytes_representations.iter().rev().skip(1), ) { - encoded_value = codec.decode_opt(encoded_value, bytes_representation, parallel)?; + encoded_value = codec.decode_opt(encoded_value, bytes_representation, options)?; } // bytes->array encoded_value = self.array_to_bytes.decode_opt( encoded_value, array_representations.last().unwrap(), - parallel, + options, )?; // array->array @@ -449,7 +493,7 @@ impl ArrayCodecTraits for CodecChain { self.array_to_array.iter().rev(), array_representations.iter().rev().skip(1), ) { - encoded_value = codec.decode_opt(encoded_value, array_representation, parallel)?; + encoded_value = codec.decode_opt(encoded_value, array_representation, options)?; } if encoded_value.len() as u64 != decoded_representation.size() { @@ -467,7 +511,7 @@ impl ArrayCodecTraits for CodecChain { &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &EncodeOptions, ) -> Result, CodecError> { if decoded_value.len() as u64 != decoded_representation.size() { return Err(CodecError::UnexpectedChunkDecodedSize( @@ -482,7 +526,7 @@ impl ArrayCodecTraits for CodecChain { // array->array for codec in &self.array_to_array { value = codec - .async_encode_opt(value, &decoded_representation, parallel) + .async_encode_opt(value, &decoded_representation, options) .await?; decoded_representation = codec.compute_encoded_size(&decoded_representation)?; } @@ -490,7 +534,7 @@ impl ArrayCodecTraits for CodecChain { // array->bytes value = self .array_to_bytes - .async_encode_opt(value, &decoded_representation, parallel) + .async_encode_opt(value, &decoded_representation, options) .await?; let mut decoded_representation = self .array_to_bytes @@ -498,7 +542,7 @@ impl ArrayCodecTraits for CodecChain { // bytes->bytes for codec in &self.bytes_to_bytes { - value = codec.async_encode_opt(value, parallel).await?; + value = codec.async_encode_opt(value, options).await?; decoded_representation = codec.compute_encoded_size(&decoded_representation); } @@ -510,7 +554,7 @@ impl ArrayCodecTraits for CodecChain { &self, mut encoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { let array_representations = self.get_array_representations(decoded_representation.clone())?; @@ -523,7 +567,7 @@ impl ArrayCodecTraits for CodecChain { bytes_representations.iter().rev().skip(1), ) { encoded_value = codec - .async_decode_opt(encoded_value, bytes_representation, parallel) + .async_decode_opt(encoded_value, bytes_representation, options) .await?; } @@ -533,7 +577,7 @@ impl ArrayCodecTraits for CodecChain { .async_decode_opt( encoded_value, array_representations.last().unwrap(), - parallel, + options, ) .await?; @@ -543,7 +587,7 @@ impl ArrayCodecTraits for CodecChain { array_representations.iter().rev().skip(1), ) { encoded_value = codec - .async_decode_opt(encoded_value, array_representation, parallel) + .async_decode_opt(encoded_value, array_representation, options) .await?; } @@ -678,16 +722,16 @@ mod tests { } assert_eq!(bytes, decoded); - let encoded = codec - .par_encode(bytes.clone(), &chunk_representation) - .unwrap(); - let decoded = codec - .par_decode(encoded.clone(), &chunk_representation) - .unwrap(); - if not_just_bytes { - assert_ne!(encoded, decoded); - } - assert_eq!(bytes, decoded); + // let encoded = codec + // .par_encode(bytes.clone(), &chunk_representation) + // .unwrap(); + // let decoded = codec + // .par_decode(encoded.clone(), &chunk_representation) + // .unwrap(); + // if not_just_bytes { + // assert_ne!(encoded, decoded); + // } + // assert_eq!(bytes, decoded); let input_handle = Box::new(std::io::Cursor::new(encoded)); let partial_decoder = codec diff --git a/src/array/codec/array_to_bytes/pcodec/pcodec_codec.rs b/src/array/codec/array_to_bytes/pcodec/pcodec_codec.rs index 46ac8286..5c0edf93 100644 --- a/src/array/codec/array_to_bytes/pcodec/pcodec_codec.rs +++ b/src/array/codec/array_to_bytes/pcodec/pcodec_codec.rs @@ -1,12 +1,11 @@ -// Note: No validation that this codec is created *without* a specified endianness for multi-byte data types. - use pco::{ChunkConfig, FloatMultSpec, IntMultSpec, PagingSpec}; use crate::{ array::{ codec::{ ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, - BytesPartialDecoderTraits, CodecError, CodecTraits, + BytesPartialDecoderTraits, CodecError, CodecTraits, DecodeOptions, EncodeOptions, + PartialDecoderOptions, RecommendedConcurrency, }, transmute_from_bytes_vec, transmute_to_bytes_vec, BytesRepresentation, ChunkRepresentation, DataType, @@ -89,11 +88,19 @@ impl CodecTraits for PcodecCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for PcodecCodec { + fn recommended_concurrency( + &self, + _decoded_representation: &ChunkRepresentation, + ) -> Result { + // pcodec does not support parallel decode + Ok(RecommendedConcurrency::one()) + } + fn encode_opt( &self, decoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &EncodeOptions, ) -> Result, CodecError> { let data_type = decoded_representation.data_type(); macro_rules! pcodec_encode { @@ -136,7 +143,7 @@ impl ArrayCodecTraits for PcodecCodec { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { let data_type = decoded_representation.data_type(); macro_rules! pcodec_decode { @@ -180,7 +187,7 @@ impl ArrayToBytesCodecTraits for PcodecCodec { &self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(pcodec_partial_decoder::PcodecPartialDecoder::new( input_handle, @@ -193,7 +200,7 @@ impl ArrayToBytesCodecTraits for PcodecCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( pcodec_partial_decoder::AsyncPCodecPartialDecoder::new( diff --git a/src/array/codec/array_to_bytes/pcodec/pcodec_partial_decoder.rs b/src/array/codec/array_to_bytes/pcodec/pcodec_partial_decoder.rs index 222f679e..e81a66c6 100644 --- a/src/array/codec/array_to_bytes/pcodec/pcodec_partial_decoder.rs +++ b/src/array/codec/array_to_bytes/pcodec/pcodec_partial_decoder.rs @@ -1,6 +1,9 @@ use crate::{ array::{ - codec::{ArrayPartialDecoderTraits, ArraySubset, BytesPartialDecoderTraits, CodecError}, + codec::{ + ArrayPartialDecoderTraits, ArraySubset, BytesPartialDecoderTraits, CodecError, + PartialDecodeOptions, + }, ChunkRepresentation, DataType, }, array_subset::IncompatibleArraySubsetAndShapeError, @@ -105,9 +108,9 @@ impl ArrayPartialDecoderTraits for PcodecPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { - let decoded = self.input_handle.decode_opt(parallel)?; + let decoded = self.input_handle.decode_opt(options)?; do_partial_decode(decoded, decoded_regions, &self.decoded_representation) } } @@ -139,9 +142,9 @@ impl AsyncArrayPartialDecoderTraits for AsyncPCodecPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { - let decoded = self.input_handle.decode_opt(parallel).await?; + let decoded = self.input_handle.decode_opt(options).await?; do_partial_decode(decoded, decoded_regions, &self.decoded_representation) } } diff --git a/src/array/codec/array_to_bytes/sharding.rs b/src/array/codec/array_to_bytes/sharding.rs index 24935971..d7a7eb30 100644 --- a/src/array/codec/array_to_bytes/sharding.rs +++ b/src/array/codec/array_to_bytes/sharding.rs @@ -27,7 +27,7 @@ use thiserror::Error; use crate::{ array::{ - codec::{ArrayToBytesCodecTraits, Codec, CodecError, CodecPlugin}, + codec::{ArrayToBytesCodecTraits, Codec, CodecError, CodecPlugin, DecodeOptions}, BytesRepresentation, ChunkRepresentation, ChunkShape, DataType, FillValue, }, metadata::Metadata, @@ -108,11 +108,11 @@ fn decode_shard_index( encoded_shard_index: Vec, index_array_representation: &ChunkRepresentation, index_codecs: &dyn ArrayToBytesCodecTraits, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { // Decode the shard index let decoded_shard_index = - index_codecs.decode_opt(encoded_shard_index, index_array_representation, parallel)?; + index_codecs.decode_opt(encoded_shard_index, index_array_representation, options)?; Ok(decoded_shard_index .chunks_exact(core::mem::size_of::()) .map(|v| u64::from_ne_bytes(v.try_into().unwrap() /* safe */)) @@ -124,11 +124,11 @@ async fn async_decode_shard_index( encoded_shard_index: Vec, index_array_representation: &ChunkRepresentation, index_codecs: &dyn ArrayToBytesCodecTraits, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { // Decode the shard index let decoded_shard_index = index_codecs - .async_decode_opt(encoded_shard_index, index_array_representation, parallel) + .async_decode_opt(encoded_shard_index, index_array_representation, options) .await?; Ok(decoded_shard_index .chunks_exact(core::mem::size_of::()) @@ -141,10 +141,11 @@ mod tests { use crate::{ array::codec::{ bytes_to_bytes::test_unbounded::TestUnboundedCodec, ArrayCodecTraits, - BytesToBytesCodecTraits, + BytesToBytesCodecTraits, EncodeOptions, PartialDecodeOptions, }, array_subset::ArraySubset, }; + use std::num::NonZeroUsize; use super::*; @@ -197,7 +198,8 @@ mod tests { }"#; fn codec_sharding_round_trip_impl( - parallel: bool, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, unbounded: bool, index_at_end: bool, all_fill_value: bool, @@ -229,10 +231,10 @@ mod tests { .build(); let encoded = codec - .encode_opt(bytes.clone(), &chunk_representation, parallel) + .encode_opt(bytes.clone(), &chunk_representation, encode_options) .unwrap(); let decoded = codec - .decode_opt(encoded.clone(), &chunk_representation, parallel) + .decode_opt(encoded.clone(), &chunk_representation, decode_options) .unwrap(); assert_ne!(encoded, decoded); assert_eq!(bytes, decoded); @@ -246,8 +248,15 @@ mod tests { for all_fill_value in [true, false] { for unbounded in [true, false] { for parallel in [true, false] { + let mut encode_options = EncodeOptions::default(); + let mut decode_options = DecodeOptions::default(); + if !parallel { + encode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + decode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + } codec_sharding_round_trip_impl( - parallel, + &encode_options, + &decode_options, unbounded, all_fill_value, index_at_end, @@ -265,14 +274,23 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn codec_sharding_round_trip2() { + use std::num::NonZeroUsize; + use crate::array::codec::{Crc32cCodec, GzipCodec}; for index_at_end in [true, false] { for all_fill_value in [true, false] { for unbounded in [true, false] { for parallel in [true, false] { + let mut encode_options = EncodeOptions::default(); + let mut decode_options = DecodeOptions::default(); + if !parallel { + encode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + decode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + } codec_sharding_round_trip_impl( - parallel, + &encode_options, + &decode_options, unbounded, all_fill_value, index_at_end, @@ -289,7 +307,8 @@ mod tests { #[cfg(feature = "async")] async fn codec_sharding_async_round_trip_impl( - parallel: bool, + encode_options: &EncodeOptions, + decode_options: &DecodeOptions, unbounded: bool, index_at_end: bool, all_fill_value: bool, @@ -321,11 +340,11 @@ mod tests { .build(); let encoded = codec - .async_encode_opt(bytes.clone(), &chunk_representation, parallel) + .async_encode_opt(bytes.clone(), &chunk_representation, encode_options) .await .unwrap(); let decoded = codec - .async_decode_opt(encoded.clone(), &chunk_representation, parallel) + .async_decode_opt(encoded.clone(), &chunk_representation, decode_options) .await .unwrap(); assert_ne!(encoded, decoded); @@ -341,8 +360,15 @@ mod tests { for all_fill_value in [true, false] { for unbounded in [true, false] { for parallel in [true, false] { + let mut encode_options = EncodeOptions::default(); + let mut decode_options = DecodeOptions::default(); + if !parallel { + encode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + decode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + } codec_sharding_async_round_trip_impl( - parallel, + &encode_options, + &decode_options, unbounded, all_fill_value, index_at_end, @@ -356,7 +382,7 @@ mod tests { } fn codec_sharding_partial_decode( - parallel: bool, + options: &PartialDecodeOptions, unbounded: bool, index_at_end: bool, all_fill_value: bool, @@ -392,16 +418,14 @@ mod tests { .bytes_to_bytes_codecs(bytes_to_bytes_codecs) .build(); - let encoded = codec - .encode_opt(bytes.clone(), &chunk_representation, parallel) - .unwrap(); + let encoded = codec.encode(bytes.clone(), &chunk_representation).unwrap(); let decoded_regions = [ArraySubset::new_with_ranges(&[1..3, 0..1])]; let input_handle = Box::new(std::io::Cursor::new(encoded)); let partial_decoder = codec - .partial_decoder_opt(input_handle, &chunk_representation, parallel) + .partial_decoder_opt(input_handle, &chunk_representation, options) .unwrap(); let decoded_partial_chunk = partial_decoder - .partial_decode_opt(&decoded_regions, parallel) + .partial_decode_opt(&decoded_regions, options) .unwrap(); let decoded_partial_chunk: Vec = decoded_partial_chunk @@ -420,8 +444,14 @@ mod tests { for all_fill_value in [true, false] { for unbounded in [true, false] { for parallel in [true, false] { + let mut encode_options = EncodeOptions::default(); + let mut decode_options = DecodeOptions::default(); + if !parallel { + encode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + decode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + } codec_sharding_partial_decode( - parallel, + &decode_options, unbounded, all_fill_value, index_at_end, @@ -434,7 +464,7 @@ mod tests { #[cfg(feature = "async")] async fn codec_sharding_async_partial_decode( - parallel: bool, + options: &PartialDecodeOptions, unbounded: bool, index_at_end: bool, all_fill_value: bool, @@ -470,17 +500,15 @@ mod tests { .bytes_to_bytes_codecs(bytes_to_bytes_codecs) .build(); - let encoded = codec - .encode_opt(bytes.clone(), &chunk_representation, parallel) - .unwrap(); + let encoded = codec.encode(bytes.clone(), &chunk_representation).unwrap(); let decoded_regions = [ArraySubset::new_with_ranges(&[1..3, 0..1])]; let input_handle = Box::new(std::io::Cursor::new(encoded)); let partial_decoder = codec - .async_partial_decoder_opt(input_handle, &chunk_representation, parallel) + .async_partial_decoder_opt(input_handle, &chunk_representation, options) .await .unwrap(); let decoded_partial_chunk = partial_decoder - .partial_decode_opt(&decoded_regions, parallel) + .partial_decode_opt(&decoded_regions, options) .await .unwrap(); @@ -501,8 +529,14 @@ mod tests { for all_fill_value in [true, false] { for unbounded in [true, false] { for parallel in [true, false] { + let mut encode_options = EncodeOptions::default(); + let mut decode_options = DecodeOptions::default(); + if !parallel { + encode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + decode_options.set_concurrent_limit(NonZeroUsize::new(1).unwrap()); + } codec_sharding_async_partial_decode( - parallel, + &decode_options, unbounded, all_fill_value, index_at_end, diff --git a/src/array/codec/array_to_bytes/sharding/sharding_codec.rs b/src/array/codec/array_to_bytes/sharding/sharding_codec.rs index 896cd26e..4ee98ad2 100644 --- a/src/array/codec/array_to_bytes/sharding/sharding_codec.rs +++ b/src/array/codec/array_to_bytes/sharding/sharding_codec.rs @@ -5,7 +5,8 @@ use crate::{ chunk_shape_to_array_shape, codec::{ ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, - BytesPartialDecoderTraits, CodecChain, CodecError, CodecTraits, + BytesPartialDecoderTraits, CodecChain, CodecError, CodecTraits, DecodeOptions, + EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, }, transmute_to_bytes_vec, unravel_index, unsafe_cell_slice::UnsafeCellSlice, @@ -102,11 +103,22 @@ impl CodecTraits for ShardingCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for ShardingCodec { + fn recommended_concurrency( + &self, + decoded_representation: &ChunkRepresentation, + ) -> Result { + let chunks_per_shard = + calculate_chunks_per_shard(decoded_representation.shape(), self.chunk_shape.as_slice()) + .map_err(|e| CodecError::Other(e.to_string()))?; + let num_elements = chunks_per_shard.num_elements(); + Ok(RecommendedConcurrency::new(num_elements, num_elements)) + } + fn encode_opt( &self, decoded_value: Vec, shard_rep: &ChunkRepresentation, - parallel: bool, + options: &EncodeOptions, ) -> Result, CodecError> { if decoded_value.len() as u64 != shard_rep.size() { return Err(CodecError::UnexpectedChunkDecodedSize( @@ -126,18 +138,10 @@ impl ArrayCodecTraits for ShardingCodec { let chunk_bytes_representation = self.inner_codecs.compute_encoded_size(&chunk_rep)?; match chunk_bytes_representation { BytesRepresentation::BoundedSize(size) | BytesRepresentation::FixedSize(size) => { - if parallel { - self.par_encode_bounded(&decoded_value, shard_rep, &chunk_rep, size) - } else { - self.encode_bounded(&decoded_value, shard_rep, &chunk_rep, size) - } + self.encode_bounded(&decoded_value, shard_rep, &chunk_rep, size, options) } BytesRepresentation::UnboundedSize => { - if parallel { - self.par_encode_unbounded(&decoded_value, shard_rep, &chunk_rep) - } else { - self.encode_unbounded(&decoded_value, shard_rep, &chunk_rep) - } + self.encode_unbounded(&decoded_value, shard_rep, &chunk_rep, options) } } } @@ -146,17 +150,18 @@ impl ArrayCodecTraits for ShardingCodec { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { let chunks_per_shard = calculate_chunks_per_shard(decoded_representation.shape(), self.chunk_shape.as_slice()) .map_err(|e| CodecError::Other(e.to_string()))?; - let shard_index = self.decode_index(&encoded_value, chunks_per_shard.as_slice(), false)?; // FIXME: par decode index? + let shard_index = + self.decode_index(&encoded_value, chunks_per_shard.as_slice(), options)?; let chunks = self.decode_chunks( &encoded_value, &shard_index, decoded_representation, - parallel, + options, )?; Ok(chunks) } @@ -166,7 +171,7 @@ impl ArrayCodecTraits for ShardingCodec { &self, decoded_value: Vec, shard_rep: &ChunkRepresentation, - parallel: bool, + options: &EncodeOptions, ) -> Result, CodecError> { if decoded_value.len() as u64 != shard_rep.size() { return Err(CodecError::UnexpectedChunkDecodedSize( @@ -186,20 +191,26 @@ impl ArrayCodecTraits for ShardingCodec { let chunk_bytes_representation = self.inner_codecs.compute_encoded_size(&chunk_rep)?; match chunk_bytes_representation { BytesRepresentation::BoundedSize(size) | BytesRepresentation::FixedSize(size) => { - if parallel { - self.async_par_encode_bounded(&decoded_value, shard_rep, &chunk_rep, size) - .await + if options.is_parallel() { + self.async_par_encode_bounded( + &decoded_value, + shard_rep, + &chunk_rep, + size, + options, + ) + .await } else { - self.async_encode_bounded(&decoded_value, shard_rep, &chunk_rep, size) + self.async_encode_bounded(&decoded_value, shard_rep, &chunk_rep, size, options) .await } } BytesRepresentation::UnboundedSize => { - if parallel { - self.async_par_encode_unbounded(&decoded_value, shard_rep, &chunk_rep) + if options.is_parallel() { + self.async_par_encode_unbounded(&decoded_value, shard_rep, &chunk_rep, options) .await } else { - self.async_encode_unbounded(&decoded_value, shard_rep, &chunk_rep) + self.async_encode_unbounded(&decoded_value, shard_rep, &chunk_rep, options) .await } } @@ -211,20 +222,20 @@ impl ArrayCodecTraits for ShardingCodec { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { let chunks_per_shard = calculate_chunks_per_shard(decoded_representation.shape(), self.chunk_shape.as_slice()) .map_err(|e| CodecError::Other(e.to_string()))?; let shard_index = self - .async_decode_index(&encoded_value, chunks_per_shard.as_slice(), parallel) + .async_decode_index(&encoded_value, chunks_per_shard.as_slice(), options) .await?; let chunks = self .async_decode_chunks( &encoded_value, &shard_index, decoded_representation, - parallel, + options, ) .await?; Ok(chunks) @@ -237,7 +248,7 @@ impl ArrayToBytesCodecTraits for ShardingCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( sharding_partial_decoder::ShardingPartialDecoder::new( @@ -247,7 +258,7 @@ impl ArrayToBytesCodecTraits for ShardingCodec { &self.inner_codecs, &self.index_codecs, self.index_location, - parallel, + options, )?, )) } @@ -257,7 +268,7 @@ impl ArrayToBytesCodecTraits for ShardingCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( sharding_partial_decoder::AsyncShardingPartialDecoder::new( @@ -267,7 +278,7 @@ impl ArrayToBytesCodecTraits for ShardingCodec { &self.inner_codecs, &self.index_codecs, self.index_location, - parallel, + options, ) .await?, )) @@ -340,191 +351,195 @@ impl ShardingCodec { num_chunks * chunk_encoded_size + index_encoded_size } - /// Preallocate shard, encode and write chunks, then truncate shard - fn encode_bounded( - &self, - decoded_value: &[u8], - shard_representation: &ChunkRepresentation, - chunk_representation: &ChunkRepresentation, - chunk_size_bounded: u64, - ) -> Result, CodecError> { - debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); - - // Allocate an array for the shard - let chunks_per_shard = - calculate_chunks_per_shard(shard_representation.shape(), chunk_representation.shape()) - .map_err(|e| CodecError::Other(e.to_string()))?; - let index_decoded_representation = - sharding_index_decoded_representation(chunks_per_shard.as_slice()); - let index_encoded_size = - compute_index_encoded_size(&self.index_codecs, &index_decoded_representation)?; - let shard_size_bounded = Self::encoded_shard_bounded_size( - index_encoded_size, - chunk_size_bounded, - chunks_per_shard.as_slice(), - ); - let shard_size_bounded = usize::try_from(shard_size_bounded).unwrap(); - let mut shard = vec![core::mem::MaybeUninit::::uninit(); shard_size_bounded]; - - // Create array index - let chunks_per_shard = - calculate_chunks_per_shard(shard_representation.shape(), chunk_representation.shape()) - .map_err(|e| CodecError::Other(e.to_string()))?; - let index_decoded_representation = - sharding_index_decoded_representation(chunks_per_shard.as_slice()); - let mut shard_index = vec![u64::MAX; index_decoded_representation.num_elements_usize()]; - - // Iterate over chunk indices - let index_encoded_size = usize::try_from(index_encoded_size).unwrap(); - let mut encoded_shard_offset: usize = match self.index_location { - ShardingIndexLocation::Start => index_encoded_size, - ShardingIndexLocation::End => 0, - }; - { - let shard_slice = unsafe { - std::slice::from_raw_parts_mut(shard.as_mut_ptr().cast::(), shard.len()) - }; - let shard_shape = shard_representation.shape_u64(); - for (chunk_index, (_chunk_indices, chunk_subset)) in unsafe { - ArraySubset::new_with_shape(shard_shape.clone()) - .iter_chunks_unchecked(self.chunk_shape.as_slice()) - } - .enumerate() - { - let bytes = unsafe { - chunk_subset.extract_bytes_unchecked( - decoded_value, - &shard_shape, - shard_representation.element_size(), - ) - }; - if !chunk_representation.fill_value().equals_all(&bytes) { - let chunk_encoded = self.inner_codecs.encode(bytes, chunk_representation)?; - shard_index[chunk_index * 2] = u64::try_from(encoded_shard_offset).unwrap(); - shard_index[chunk_index * 2 + 1] = u64::try_from(chunk_encoded.len()).unwrap(); - shard_slice[encoded_shard_offset..encoded_shard_offset + chunk_encoded.len()] - .copy_from_slice(&chunk_encoded); - encoded_shard_offset += chunk_encoded.len(); - } - } - } - - // Truncate the shard - let shard_length = encoded_shard_offset - + match self.index_location { - ShardingIndexLocation::Start => 0, - ShardingIndexLocation::End => index_encoded_size, - }; - shard.truncate(shard_length); - - // Encode array index - let encoded_array_index = self.index_codecs.encode( - transmute_to_bytes_vec(shard_index), - &index_decoded_representation, - )?; - - // Add the index - { - let shard_slice = unsafe { - std::slice::from_raw_parts_mut(shard.as_mut_ptr().cast::(), shard.len()) - }; - match self.index_location { - ShardingIndexLocation::Start => { - shard_slice[..index_encoded_size].copy_from_slice(&encoded_array_index); - } - ShardingIndexLocation::End => { - shard_slice[shard_length - index_encoded_size..] - .copy_from_slice(&encoded_array_index); - } - } - } - #[allow(clippy::transmute_undefined_repr)] - let shard = unsafe { core::mem::transmute(shard) }; - Ok(shard) - } - - /// Encode inner chunks, then allocate shard, then write to shard - fn encode_unbounded( - &self, - decoded_value: &[u8], - shard_representation: &ChunkRepresentation, - chunk_representation: &ChunkRepresentation, - ) -> Result, CodecError> { - debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); - - // Create array index - let chunks_per_shard = - calculate_chunks_per_shard(shard_representation.shape(), chunk_representation.shape()) - .map_err(|e| CodecError::Other(e.to_string()))?; - let index_decoded_representation = - sharding_index_decoded_representation(chunks_per_shard.as_slice()); - let mut shard_index = vec![u64::MAX; index_decoded_representation.num_elements_usize()]; - - // Iterate over chunk indices - let mut shard_inner_chunks = Vec::new(); - let index_encoded_size = - compute_index_encoded_size(&self.index_codecs, &index_decoded_representation)?; - let mut encoded_shard_offset = match self.index_location { - ShardingIndexLocation::Start => index_encoded_size, - ShardingIndexLocation::End => 0, - }; - let shard_shape = shard_representation.shape_u64(); - for (chunk_index, (_chunk_indices, chunk_subset)) in unsafe { - ArraySubset::new_with_shape(shard_shape.clone()) - .iter_chunks_unchecked(self.chunk_shape.as_slice()) - } - .enumerate() - { - let bytes = unsafe { - chunk_subset.extract_bytes_unchecked( - decoded_value, - &shard_shape, - shard_representation.element_size(), - ) - }; - if !chunk_representation.fill_value().equals_all(&bytes) { - let chunk_encoded = self.inner_codecs.encode(bytes, chunk_representation)?; - shard_index[chunk_index * 2] = encoded_shard_offset; - shard_index[chunk_index * 2 + 1] = chunk_encoded.len().try_into().unwrap(); - encoded_shard_offset += chunk_encoded.len() as u64; - shard_inner_chunks.push(chunk_encoded); - } - } - - // Encode array index - let encoded_array_index = self.index_codecs.encode( - transmute_to_bytes_vec(shard_index), - &index_decoded_representation, - )?; - - // Encode the shard - let shard_inner_chunks_length = shard_inner_chunks.iter().map(Vec::len).sum::(); - let shard_size = shard_inner_chunks_length + encoded_array_index.len(); - let mut shard = Vec::with_capacity(shard_size); - match self.index_location { - ShardingIndexLocation::Start => { - shard.extend(encoded_array_index); - for chunk in shard_inner_chunks { - shard.extend(chunk); - } - } - ShardingIndexLocation::End => { - for chunk in shard_inner_chunks { - shard.extend(chunk); - } - shard.extend(encoded_array_index); - } - } - Ok(shard) - } + // /// Preallocate shard, encode and write chunks, then truncate shard + // fn encode_bounded( + // &self, + // decoded_value: &[u8], + // shard_representation: &ChunkRepresentation, + // chunk_representation: &ChunkRepresentation, + // chunk_size_bounded: u64, + // options: &EncodeOptions, + // ) -> Result, CodecError> { + // debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); + + // // Allocate an array for the shard + // let chunks_per_shard = + // calculate_chunks_per_shard(shard_representation.shape(), chunk_representation.shape()) + // .map_err(|e| CodecError::Other(e.to_string()))?; + // let index_decoded_representation = + // sharding_index_decoded_representation(chunks_per_shard.as_slice()); + // let index_encoded_size = + // compute_index_encoded_size(&self.index_codecs, &index_decoded_representation)?; + // let shard_size_bounded = Self::encoded_shard_bounded_size( + // index_encoded_size, + // chunk_size_bounded, + // chunks_per_shard.as_slice(), + // ); + // let shard_size_bounded = usize::try_from(shard_size_bounded).unwrap(); + // let mut shard = vec![core::mem::MaybeUninit::::uninit(); shard_size_bounded]; + + // // Create array index + // let chunks_per_shard = + // calculate_chunks_per_shard(shard_representation.shape(), chunk_representation.shape()) + // .map_err(|e| CodecError::Other(e.to_string()))?; + // let index_decoded_representation = + // sharding_index_decoded_representation(chunks_per_shard.as_slice()); + // let mut shard_index = vec![u64::MAX; index_decoded_representation.num_elements_usize()]; + + // // Iterate over chunk indices + // let index_encoded_size = usize::try_from(index_encoded_size).unwrap(); + // let mut encoded_shard_offset: usize = match self.index_location { + // ShardingIndexLocation::Start => index_encoded_size, + // ShardingIndexLocation::End => 0, + // }; + // { + // let shard_slice = unsafe { + // std::slice::from_raw_parts_mut(shard.as_mut_ptr().cast::(), shard.len()) + // }; + // let shard_shape = shard_representation.shape_u64(); + // for (chunk_index, (_chunk_indices, chunk_subset)) in unsafe { + // ArraySubset::new_with_shape(shard_shape.clone()) + // .iter_chunks_unchecked(self.chunk_shape.as_slice()) + // } + // .enumerate() + // { + // let bytes = unsafe { + // chunk_subset.extract_bytes_unchecked( + // decoded_value, + // &shard_shape, + // shard_representation.element_size(), + // ) + // }; + // if !chunk_representation.fill_value().equals_all(&bytes) { + // let chunk_encoded = self.inner_codecs.encode(bytes, chunk_representation)?; + // shard_index[chunk_index * 2] = u64::try_from(encoded_shard_offset).unwrap(); + // shard_index[chunk_index * 2 + 1] = u64::try_from(chunk_encoded.len()).unwrap(); + // shard_slice[encoded_shard_offset..encoded_shard_offset + chunk_encoded.len()] + // .copy_from_slice(&chunk_encoded); + // encoded_shard_offset += chunk_encoded.len(); + // } + // } + // } + + // // Truncate the shard + // let shard_length = encoded_shard_offset + // + match self.index_location { + // ShardingIndexLocation::Start => 0, + // ShardingIndexLocation::End => index_encoded_size, + // }; + // shard.truncate(shard_length); + + // // Encode array index + // let encoded_array_index = self.index_codecs.encode( + // transmute_to_bytes_vec(shard_index), + // &index_decoded_representation, + // )?; + + // // Add the index + // { + // let shard_slice = unsafe { + // std::slice::from_raw_parts_mut(shard.as_mut_ptr().cast::(), shard.len()) + // }; + // match self.index_location { + // ShardingIndexLocation::Start => { + // shard_slice[..index_encoded_size].copy_from_slice(&encoded_array_index); + // } + // ShardingIndexLocation::End => { + // shard_slice[shard_length - index_encoded_size..] + // .copy_from_slice(&encoded_array_index); + // } + // } + // } + // #[allow(clippy::transmute_undefined_repr)] + // let shard = unsafe { core::mem::transmute(shard) }; + // Ok(shard) + // } + + // /// Encode inner chunks, then allocate shard, then write to shard + // fn encode_unbounded( + // &self, + // decoded_value: &[u8], + // shard_representation: &ChunkRepresentation, + // chunk_representation: &ChunkRepresentation, + // options: &EncodeOptions, + // ) -> Result, CodecError> { + // debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); + + // // Create array index + // let chunks_per_shard = + // calculate_chunks_per_shard(shard_representation.shape(), chunk_representation.shape()) + // .map_err(|e| CodecError::Other(e.to_string()))?; + // let index_decoded_representation = + // sharding_index_decoded_representation(chunks_per_shard.as_slice()); + // let mut shard_index = vec![u64::MAX; index_decoded_representation.num_elements_usize()]; + + // // Iterate over chunk indices + // let mut shard_inner_chunks = Vec::new(); + // let index_encoded_size = + // compute_index_encoded_size(&self.index_codecs, &index_decoded_representation)?; + // let mut encoded_shard_offset = match self.index_location { + // ShardingIndexLocation::Start => index_encoded_size, + // ShardingIndexLocation::End => 0, + // }; + // let shard_shape = shard_representation.shape_u64(); + // for (chunk_index, (_chunk_indices, chunk_subset)) in unsafe { + // ArraySubset::new_with_shape(shard_shape.clone()) + // .iter_chunks_unchecked(self.chunk_shape.as_slice()) + // } + // .enumerate() + // { + // let bytes = unsafe { + // chunk_subset.extract_bytes_unchecked( + // decoded_value, + // &shard_shape, + // shard_representation.element_size(), + // ) + // }; + // if !chunk_representation.fill_value().equals_all(&bytes) { + // let chunk_encoded = self.inner_codecs.encode(bytes, chunk_representation)?; + // shard_index[chunk_index * 2] = encoded_shard_offset; + // shard_index[chunk_index * 2 + 1] = chunk_encoded.len().try_into().unwrap(); + // encoded_shard_offset += chunk_encoded.len() as u64; + // shard_inner_chunks.push(chunk_encoded); + // } + // } + + // // Encode array index + // let encoded_array_index = self.index_codecs.encode( + // transmute_to_bytes_vec(shard_index), + // &index_decoded_representation, + // )?; + + // // Encode the shard + // let shard_inner_chunks_length = shard_inner_chunks.iter().map(Vec::len).sum::(); + // let shard_size = shard_inner_chunks_length + encoded_array_index.len(); + // let mut shard = Vec::with_capacity(shard_size); + // match self.index_location { + // ShardingIndexLocation::Start => { + // shard.extend(encoded_array_index); + // for chunk in shard_inner_chunks { + // shard.extend(chunk); + // } + // } + // ShardingIndexLocation::End => { + // for chunk in shard_inner_chunks { + // shard.extend(chunk); + // } + // shard.extend(encoded_array_index); + // } + // } + // Ok(shard) + // } /// Preallocate shard, encode and write chunks (in parallel), then truncate shard - fn par_encode_bounded( + #[allow(clippy::too_many_lines)] + fn encode_bounded( &self, decoded_value: &[u8], shard_representation: &ChunkRepresentation, chunk_representation: &ChunkRepresentation, chunk_size_bounded: u64, + options: &EncodeOptions, ) -> Result, CodecError> { debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); // already validated in par_encode @@ -563,16 +578,18 @@ impl ShardingCodec { let shard_slice = UnsafeCellSlice::new(shard_slice); let shard_index_slice = UnsafeCellSlice::new(&mut shard_index); let shard_shape = shard_representation.shape_u64(); - (0..chunks_per_shard + let n_chunks = chunks_per_shard .as_slice() .iter() - .map(|i| i.get()) - .product::()) - .into_par_iter() - .try_for_each(|chunk_index| { + .map(|i| usize::try_from(i.get()).unwrap()) + .product::(); + rayon_iter_concurrent_limit::iter_concurrent_limit!( + options.concurrent_limit().get(), + (0..n_chunks).into_par_iter(), + try_for_each, + |chunk_index| { let chunk_subset = - self.chunk_index_to_subset(chunk_index, chunks_per_shard.as_slice()); - let chunk_index = usize::try_from(chunk_index).unwrap(); + self.chunk_index_to_subset(chunk_index as u64, chunks_per_shard.as_slice()); let bytes = unsafe { chunk_subset.extract_bytes_unchecked( decoded_value, @@ -582,7 +599,8 @@ impl ShardingCodec { }; if !chunk_representation.fill_value().equals_all(&bytes) { let chunk_encoded = - self.inner_codecs.encode(bytes, chunk_representation)?; + self.inner_codecs + .encode_opt(bytes, chunk_representation, options)?; // FIXME: Adjust options for inner codec decoding let chunk_offset = encoded_shard_offset .fetch_add(chunk_encoded.len(), std::sync::atomic::Ordering::Relaxed); @@ -606,7 +624,8 @@ impl ShardingCodec { } } Ok(()) - })?; + } + )?; } // Truncate shard @@ -618,9 +637,10 @@ impl ShardingCodec { shard.truncate(shard_length); // Encode and write array index - let encoded_array_index = self.index_codecs.par_encode( + let encoded_array_index = self.index_codecs.encode_opt( transmute_to_bytes_vec(shard_index), &index_decoded_representation, + options, )?; { let shard_slice = unsafe { @@ -642,11 +662,15 @@ impl ShardingCodec { } /// Encode inner chunks (in parallel), then allocate shard, then write to shard (in parallel) - fn par_encode_unbounded( + // TODO: Collecting chunks then allocating shard can use a lot of memory, have a low memory variant + // TODO: Also benchmark performance with just performing an alloc like 1x decoded size and writing directly into it, growing if needed + #[allow(clippy::too_many_lines)] + fn encode_unbounded( &self, decoded_value: &[u8], shard_representation: &ChunkRepresentation, chunk_representation: &ChunkRepresentation, + options: &EncodeOptions, ) -> Result, CodecError> { debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); // already validated in par_encode @@ -661,32 +685,40 @@ impl ShardingCodec { // Find chunks that are not entirely the fill value and collect their decoded bytes let shard_shape = shard_representation.shape_u64(); - let encoded_chunks: Vec<(u64, Vec)> = (0..chunks_per_shard + let n_chunks = chunks_per_shard .as_slice() .iter() - .map(|i| i.get()) - .product::()) - .into_par_iter() - .filter_map(|chunk_index| { - let chunk_subset = - self.chunk_index_to_subset(chunk_index, chunks_per_shard.as_slice()); - let bytes = unsafe { - chunk_subset.extract_bytes_unchecked( - decoded_value, - &shard_shape, - shard_representation.element_size(), - ) - }; - if chunk_representation.fill_value().equals_all(&bytes) { - None - } else { - let encoded_chunk = self.inner_codecs.encode(bytes, chunk_representation); - match encoded_chunk { - Ok(encoded_chunk) => Some(Ok((chunk_index, encoded_chunk))), - Err(err) => Some(Err(err)), + .map(|i| usize::try_from(i.get()).unwrap()) + .product::(); + + let encoded_chunks: Vec<(usize, Vec)> = + rayon_iter_concurrent_limit::iter_concurrent_limit!( + options.concurrent_limit().get(), + (0..n_chunks).into_par_iter(), + filter_map, + |chunk_index| { + let chunk_subset = + self.chunk_index_to_subset(chunk_index as u64, chunks_per_shard.as_slice()); + let bytes = unsafe { + chunk_subset.extract_bytes_unchecked( + decoded_value, + &shard_shape, + shard_representation.element_size(), + ) + }; + if chunk_representation.fill_value().equals_all(&bytes) { + None + } else { + let encoded_chunk = + self.inner_codecs + .encode_opt(bytes, chunk_representation, options); // FIXME: Adjust options for inner chunk encoding + match encoded_chunk { + Ok(encoded_chunk) => Some(Ok((chunk_index, encoded_chunk))), + Err(err) => Some(Err(err)), + } } } - }) + ) .collect::, _>>()?; // Allocate the shard @@ -705,16 +737,17 @@ impl ShardingCodec { }; // Write shard and update shard index - { + if !encoded_chunks.is_empty() { let shard_slice = unsafe { std::slice::from_raw_parts_mut(shard.as_mut_ptr().cast::(), shard.len()) }; let shard_slice = UnsafeCellSlice::new(shard_slice); let shard_index_slice = UnsafeCellSlice::new(&mut shard_index); - encoded_chunks - .into_par_iter() - .for_each(|(chunk_index, chunk_encoded)| { - let chunk_index = usize::try_from(chunk_index).unwrap(); + rayon_iter_concurrent_limit::iter_concurrent_limit!( + options.concurrent_limit().get(), + encoded_chunks.into_par_iter(), + for_each, + |(chunk_index, chunk_encoded)| { let chunk_offset = encoded_shard_offset .fetch_add(chunk_encoded.len(), std::sync::atomic::Ordering::Relaxed); unsafe { @@ -727,13 +760,15 @@ impl ShardingCodec { shard_unsafe[chunk_offset..chunk_offset + chunk_encoded.len()] .copy_from_slice(&chunk_encoded); } - }); + } + ); } // Write shard index - let encoded_array_index = self.index_codecs.par_encode( + let encoded_array_index = self.index_codecs.encode_opt( transmute_to_bytes_vec(shard_index), &index_decoded_representation, + options, )?; { let shard_slice = unsafe { @@ -761,6 +796,7 @@ impl ShardingCodec { decoded_value: &[u8], shard_representation: &ChunkRepresentation, chunk_representation: &ChunkRepresentation, + options: &EncodeOptions, ) -> Result, CodecError> { debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); @@ -797,7 +833,7 @@ impl ShardingCodec { if !chunk_representation.fill_value().equals_all(&bytes) { let chunk_encoded = self .inner_codecs - .async_encode_opt(bytes, chunk_representation, false) + .async_encode_opt(bytes, chunk_representation, options) .await?; shard_index[chunk_index * 2] = encoded_shard_offset; shard_index[chunk_index * 2 + 1] = chunk_encoded.len().try_into().unwrap(); @@ -812,7 +848,7 @@ impl ShardingCodec { .async_encode_opt( transmute_to_bytes_vec(shard_index), &index_decoded_representation, - false, + options, ) .await?; @@ -845,6 +881,7 @@ impl ShardingCodec { shard_representation: &ChunkRepresentation, chunk_representation: &ChunkRepresentation, chunk_size_bounded: u64, + options: &EncodeOptions, ) -> Result, CodecError> { debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); @@ -901,7 +938,7 @@ impl ShardingCodec { if !chunk_representation.fill_value().equals_all(&bytes) { let chunk_encoded = self .inner_codecs - .async_encode_opt(bytes, chunk_representation, false) + .async_encode_opt(bytes, chunk_representation, options) .await?; shard_index[chunk_index * 2] = u64::try_from(encoded_shard_offset).unwrap(); shard_index[chunk_index * 2 + 1] = u64::try_from(chunk_encoded.len()).unwrap(); @@ -926,7 +963,7 @@ impl ShardingCodec { .async_encode_opt( transmute_to_bytes_vec(shard_index), &index_decoded_representation, - false, + options, ) .await?; @@ -959,6 +996,7 @@ impl ShardingCodec { shard_representation: &ChunkRepresentation, chunk_representation: &ChunkRepresentation, chunk_size_bounded: u64, + options: &EncodeOptions, ) -> Result, CodecError> { debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); // already validated in par_encode @@ -1021,7 +1059,7 @@ impl ShardingCodec { Some(async move { let chunk_encoded = self .inner_codecs - .async_encode_opt(bytes, chunk_representation, false) + .async_encode_opt(bytes, chunk_representation, options) .await?; Ok::<_, CodecError>((chunk_index, chunk_encoded)) }) @@ -1074,7 +1112,7 @@ impl ShardingCodec { .async_encode_opt( transmute_to_bytes_vec(shard_index), &index_decoded_representation, - true, + options, ) .await?; { @@ -1105,6 +1143,7 @@ impl ShardingCodec { decoded_value: &[u8], shard_representation: &ChunkRepresentation, chunk_representation: &ChunkRepresentation, + options: &EncodeOptions, ) -> Result, CodecError> { debug_assert_eq!(decoded_value.len() as u64, shard_representation.size()); // already validated in par_encode @@ -1145,7 +1184,7 @@ impl ShardingCodec { .map(|(chunk_index, bytes)| async move { let encoded_chunk = self .inner_codecs - .async_encode_opt(bytes, chunk_representation, false) + .async_encode_opt(bytes, chunk_representation, options) .await?; Ok::<_, CodecError>((chunk_index, encoded_chunk)) }), @@ -1211,7 +1250,7 @@ impl ShardingCodec { .async_encode_opt( transmute_to_bytes_vec(shard_index), &index_decoded_representation, - true, + options, ) .await?; { @@ -1237,7 +1276,7 @@ impl ShardingCodec { &self, encoded_shard: &[u8], chunks_per_shard: &[NonZeroU64], - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { // Get index array representation and encoded size let index_array_representation = sharding_index_decoded_representation(chunks_per_shard); @@ -1267,7 +1306,7 @@ impl ShardingCodec { encoded_shard_index.to_vec(), &index_array_representation, &self.index_codecs, - parallel, + options, ) } @@ -1276,7 +1315,7 @@ impl ShardingCodec { &self, encoded_shard: &[u8], chunks_per_shard: &[NonZeroU64], - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { // Get index array representation and encoded size let index_array_representation = sharding_index_decoded_representation(chunks_per_shard); @@ -1306,17 +1345,18 @@ impl ShardingCodec { encoded_shard_index.to_vec(), &index_array_representation, &self.index_codecs, - parallel, + options, ) .await } + #[allow(clippy::too_many_lines)] fn decode_chunks( &self, encoded_shard: &[u8], shard_index: &[u64], shard_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { // Allocate an array for the output let mut shard = vec![MaybeUninit::::uninit(); shard_representation.size_usize()]; @@ -1332,19 +1372,35 @@ impl ShardingCodec { ) }; + let any_empty = shard_index + .into_par_iter() + .any(|offset_or_size| *offset_or_size == u64::MAX); + let fill_value_chunk = if any_empty { + Some( + chunk_representation + .fill_value() + .as_ne_bytes() + .repeat(chunk_representation.num_elements_usize()), + ) + } else { + None + }; + let shard_shape = shard_representation.shape_u64(); - if parallel { + if options.is_parallel() { let chunks_per_shard = calculate_chunks_per_shard( shard_representation.shape(), chunk_representation.shape(), ) .map_err(|e| CodecError::Other(e.to_string()))?; - let shard_slice = UnsafeCellSlice::new(shard_slice); - (0..chunks_per_shard + let num_chunks = chunks_per_shard .as_slice() .iter() .map(|i| i.get()) - .product::()) + .product::(); + let shard_slice = UnsafeCellSlice::new(shard_slice); + let element_size = chunk_representation.element_size() as u64; + (0..num_chunks) .into_par_iter() .try_for_each(|chunk_index| { let chunk_subset = @@ -1352,37 +1408,47 @@ impl ShardingCodec { let chunk_index = usize::try_from(chunk_index).unwrap(); let shard_slice = unsafe { shard_slice.get() }; + let mut copy_to_subset = |decoded_chunk: &[u8]| { + // Copy to subset of shard + let mut data_idx = 0; + let contiguous_iterator = unsafe { + chunk_subset.iter_contiguous_linearised_indices_unchecked(&shard_shape) + }; + let length = usize::try_from( + contiguous_iterator.contiguous_elements() * element_size, + ) + .unwrap(); + for (index, _) in unsafe { + chunk_subset.iter_contiguous_linearised_indices_unchecked(&shard_shape) + } { + let shard_offset = usize::try_from(index * element_size).unwrap(); + shard_slice[shard_offset..shard_offset + length] + .copy_from_slice(&decoded_chunk[data_idx..data_idx + length]); + data_idx += length; + } + }; + // Read the offset/size let offset = shard_index[chunk_index * 2]; let size = shard_index[chunk_index * 2 + 1]; - let decoded_chunk = if offset == u64::MAX && size == u64::MAX { - // Can fill values be populated faster than repeat? - chunk_representation - .fill_value() - .as_ne_bytes() - .repeat(chunk_representation.num_elements_usize()) + if offset == u64::MAX && size == u64::MAX { + if let Some(fill_value_chunk) = &fill_value_chunk { + copy_to_subset(fill_value_chunk.as_slice()); + } + // else is unreachable } else { let offset: usize = offset.try_into().unwrap(); // safe let size: usize = size.try_into().unwrap(); // safe let encoded_chunk_slice = encoded_shard[offset..offset + size].to_vec(); // NOTE: Intentionally using single threaded decode, since parallelisation is in the loop - self.inner_codecs - .decode(encoded_chunk_slice, &chunk_representation)? + let decoded = self.inner_codecs.decode_opt( + encoded_chunk_slice, + &chunk_representation, + options, + )?; // FIXME: Adjust options for inner chunk decoding + copy_to_subset(&decoded); }; - // Copy to subset of shard - let mut data_idx = 0; - let element_size = chunk_representation.element_size() as u64; - for (index, num_elements) in unsafe { - chunk_subset.iter_contiguous_linearised_indices_unchecked(&shard_shape) - } { - let shard_offset = usize::try_from(index * element_size).unwrap(); - let length = usize::try_from(num_elements * element_size).unwrap(); - shard_slice[shard_offset..shard_offset + length] - .copy_from_slice(&decoded_chunk[data_idx..data_idx + length]); - data_idx += length; - } - Ok::<_, CodecError>(()) })?; } else { @@ -1406,8 +1472,11 @@ impl ShardingCodec { let offset: usize = offset.try_into().unwrap(); // safe let size: usize = size.try_into().unwrap(); // safe let encoded_chunk_slice = encoded_shard[offset..offset + size].to_vec(); - self.inner_codecs - .decode(encoded_chunk_slice, &chunk_representation)? + self.inner_codecs.decode_opt( + encoded_chunk_slice, + &chunk_representation, + options, + )? // FIXME: Adjust options for inner chunk encoding }; // Copy to subset of shard @@ -1436,7 +1505,7 @@ impl ShardingCodec { encoded_shard: &[u8], shard_index: &[u64], shard_representation: &ChunkRepresentation, - parallel: bool, + options: &DecodeOptions, ) -> Result, CodecError> { // Allocate an array for the output let mut shard = vec![MaybeUninit::::uninit(); shard_representation.size_usize()]; @@ -1454,7 +1523,7 @@ impl ShardingCodec { // Decode chunks let shard_shape = shard_representation.shape_u64(); - if parallel { + if options.is_parallel() { let chunks_per_shard = calculate_chunks_per_shard( shard_representation.shape(), chunk_representation.shape(), @@ -1500,7 +1569,11 @@ impl ShardingCodec { let encoded_chunk_slice = encoded_shard[*offset..*offset + *size].to_vec(); self.inner_codecs - .async_decode_opt(encoded_chunk_slice, &chunk_representation, false) + .async_decode_opt( + encoded_chunk_slice, + &chunk_representation, + options, + ) // FIXME .await .map(move |decoded_chunk| (chunk_subset, decoded_chunk)) } @@ -1579,7 +1652,7 @@ impl ShardingCodec { let size: usize = size.try_into().unwrap(); // safe let encoded_chunk_slice = encoded_shard[offset..offset + size].to_vec(); self.inner_codecs - .async_decode_opt(encoded_chunk_slice, &chunk_representation, false) + .async_decode_opt(encoded_chunk_slice, &chunk_representation, options) .await? }; diff --git a/src/array/codec/array_to_bytes/sharding/sharding_partial_decoder.rs b/src/array/codec/array_to_bytes/sharding/sharding_partial_decoder.rs index a4878eb8..23c74dba 100644 --- a/src/array/codec/array_to_bytes/sharding/sharding_partial_decoder.rs +++ b/src/array/codec/array_to_bytes/sharding/sharding_partial_decoder.rs @@ -9,6 +9,7 @@ use crate::{ codec::{ ArrayPartialDecoderTraits, ArraySubset, ArrayToBytesCodecTraits, ByteIntervalPartialDecoder, BytesPartialDecoderTraits, CodecChain, CodecError, + PartialDecodeOptions, PartialDecoderOptions, }, ravel_indices, unsafe_cell_slice::UnsafeCellSlice, @@ -49,7 +50,7 @@ impl<'a> ShardingPartialDecoder<'a> { inner_codecs: &'a CodecChain, index_codecs: &'a CodecChain, index_location: ShardingIndexLocation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result { let shard_index = Self::decode_shard_index( &*input_handle, @@ -57,7 +58,7 @@ impl<'a> ShardingPartialDecoder<'a> { index_location, chunk_shape.as_slice(), &decoded_representation, - parallel, + options, )?; Ok(Self { input_handle, @@ -75,7 +76,7 @@ impl<'a> ShardingPartialDecoder<'a> { index_location: ShardingIndexLocation, chunk_shape: &[NonZeroU64], decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &PartialDecoderOptions, ) -> Result>, CodecError> { let shard_shape = decoded_representation.shape(); let chunk_representation = unsafe { @@ -105,7 +106,7 @@ impl<'a> ShardingPartialDecoder<'a> { }; let encoded_shard_index = input_handle - .partial_decode_opt(&[index_byte_range], parallel)? + .partial_decode_opt(&[index_byte_range], options)? .map(|mut v| v.remove(0)); Ok(match encoded_shard_index { @@ -113,7 +114,7 @@ impl<'a> ShardingPartialDecoder<'a> { encoded_shard_index, &index_array_representation, index_codecs, - parallel, + options, )?), None => None, }) @@ -124,118 +125,7 @@ impl ArrayPartialDecoderTraits for ShardingPartialDecoder<'_> { fn partial_decode_opt( &self, array_subsets: &[ArraySubset], - parallel: bool, - ) -> Result>, CodecError> { - if parallel { - self.par_partial_decode(array_subsets) - } else { - self.partial_decode(array_subsets) - } - } - - fn partial_decode(&self, array_subsets: &[ArraySubset]) -> Result>, CodecError> { - let Some(shard_index) = &self.shard_index else { - return Ok(array_subsets - .iter() - .map(|array_subset| { - self.decoded_representation - .fill_value() - .as_ne_bytes() - .repeat(array_subset.num_elements_usize()) - }) - .collect()); - }; - - let chunk_representation = unsafe { - ChunkRepresentation::new_unchecked( - self.chunk_grid.chunk_shape().to_vec(), - self.decoded_representation.data_type().clone(), - self.decoded_representation.fill_value().clone(), - ) - }; - - let chunks_per_shard = calculate_chunks_per_shard( - self.decoded_representation.shape(), - chunk_representation.shape(), - ) - .map_err(|e| CodecError::Other(e.to_string()))?; - let chunks_per_shard = chunk_shape_to_array_shape(chunks_per_shard.as_slice()); - - let element_size = self.decoded_representation.element_size(); - let element_size_u64 = element_size as u64; - let fill_value = chunk_representation.fill_value().as_ne_bytes(); - - let mut out = Vec::with_capacity(array_subsets.len()); - for array_subset in array_subsets { - let array_subset_size = - usize::try_from(array_subset.num_elements() * element_size_u64).unwrap(); - let mut out_array_subset = vec![0; array_subset_size]; - - // Decode those chunks if required and put in chunk cache - for (chunk_indices, chunk_subset) in - unsafe { array_subset.iter_chunks_unchecked(chunk_representation.shape()) } - { - let shard_index_index: usize = - usize::try_from(ravel_indices(&chunk_indices, &chunks_per_shard) * 2).unwrap(); - let offset = shard_index[shard_index_index]; - let size = shard_index[shard_index_index + 1]; - - let overlap = unsafe { array_subset.overlap_unchecked(&chunk_subset) }; - let decoded_bytes = if offset == u64::MAX && size == u64::MAX { - // The chunk is just the fill value - fill_value.repeat(chunk_subset.num_elements_usize()) - } else { - // The chunk must be decoded - let partial_decoder = self.inner_codecs.partial_decoder( - Box::new(ByteIntervalPartialDecoder::new( - &*self.input_handle, - offset, - size, - )), - &chunk_representation, - )?; - let array_subset_in_chunk_subset = - unsafe { overlap.relative_to_unchecked(chunk_subset.start()) }; - - // Partial decoding is actually really slow with the blosc codec! Assume sharded chunks are small, and just decode the whole thing and extract bytes - // TODO: Make this behaviour optional? - // partial_decoder - // .partial_decode(&[array_subset_in_chunk_subset.clone()])? - // .remove(0) - let decoded_chunk = partial_decoder - .partial_decode(&[ArraySubset::new_with_shape( - chunk_subset.shape().to_vec(), - )])? - .remove(0); - array_subset_in_chunk_subset - .extract_bytes(&decoded_chunk, chunk_subset.shape(), element_size) - .unwrap() - }; - - // Copy decoded bytes to the output - let chunk_subset_in_array_subset = - unsafe { overlap.relative_to_unchecked(array_subset.start()) }; - let mut decoded_offset = 0; - for (array_subset_element_index, num_elements) in unsafe { - chunk_subset_in_array_subset - .iter_contiguous_linearised_indices_unchecked(array_subset.shape()) - } { - let output_offset = - usize::try_from(array_subset_element_index * element_size_u64).unwrap(); - let length = usize::try_from(num_elements * element_size_u64).unwrap(); - out_array_subset[output_offset..output_offset + length] - .copy_from_slice(&decoded_bytes[decoded_offset..decoded_offset + length]); - decoded_offset += length; - } - } - out.push(out_array_subset); - } - Ok(out) - } - - fn par_partial_decode( - &self, - array_subsets: &[ArraySubset], + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let Some(shard_index) = &self.shard_index else { return Ok(array_subsets @@ -275,6 +165,7 @@ impl ArrayPartialDecoderTraits for ShardingPartialDecoder<'_> { let out_array_subset_slice = UnsafeCellSlice::new(out_array_subset.as_mut_slice()); // Decode those chunks if required + // FIXME: Concurrent limit here unsafe { array_subset.iter_chunks_unchecked(chunk_representation.shape()) } .par_bridge() .try_for_each(|(chunk_indices, chunk_subset)| { @@ -296,17 +187,18 @@ impl ArrayPartialDecoderTraits for ShardingPartialDecoder<'_> { fill_value.repeat(array_subset_in_chunk_subset.num_elements_usize()) } else { // The chunk must be decoded - let partial_decoder = self.inner_codecs.partial_decoder( + let partial_decoder = self.inner_codecs.partial_decoder_opt( Box::new(ByteIntervalPartialDecoder::new( &*self.input_handle, offset, size, )), &chunk_representation, + options, // FIXME: Adjust options for partial decoding )?; // NOTE: Intentionally using single threaded decode, since parallelisation is in the loop partial_decoder - .partial_decode(&[array_subset_in_chunk_subset])? + .partial_decode_opt(&[array_subset_in_chunk_subset], options)? // FIXME: Adjust options for partial decoding .remove(0) }; @@ -333,6 +225,106 @@ impl ArrayPartialDecoderTraits for ShardingPartialDecoder<'_> { } Ok(out) } + + // fn partial_decode(&self, array_subsets: &[ArraySubset]) -> Result>, CodecError> { + // let Some(shard_index) = &self.shard_index else { + // return Ok(array_subsets + // .iter() + // .map(|array_subset| { + // self.decoded_representation + // .fill_value() + // .as_ne_bytes() + // .repeat(array_subset.num_elements_usize()) + // }) + // .collect()); + // }; + + // let chunk_representation = unsafe { + // ChunkRepresentation::new_unchecked( + // self.chunk_grid.chunk_shape().to_vec(), + // self.decoded_representation.data_type().clone(), + // self.decoded_representation.fill_value().clone(), + // ) + // }; + + // let chunks_per_shard = calculate_chunks_per_shard( + // self.decoded_representation.shape(), + // chunk_representation.shape(), + // ) + // .map_err(|e| CodecError::Other(e.to_string()))?; + // let chunks_per_shard = chunk_shape_to_array_shape(chunks_per_shard.as_slice()); + + // let element_size = self.decoded_representation.element_size(); + // let element_size_u64 = element_size as u64; + // let fill_value = chunk_representation.fill_value().as_ne_bytes(); + + // let mut out = Vec::with_capacity(array_subsets.len()); + // for array_subset in array_subsets { + // let array_subset_size = + // usize::try_from(array_subset.num_elements() * element_size_u64).unwrap(); + // let mut out_array_subset = vec![0; array_subset_size]; + + // // Decode those chunks if required and put in chunk cache + // for (chunk_indices, chunk_subset) in + // unsafe { array_subset.iter_chunks_unchecked(chunk_representation.shape()) } + // { + // let shard_index_index: usize = + // usize::try_from(ravel_indices(&chunk_indices, &chunks_per_shard) * 2).unwrap(); + // let offset = shard_index[shard_index_index]; + // let size = shard_index[shard_index_index + 1]; + + // let overlap = unsafe { array_subset.overlap_unchecked(&chunk_subset) }; + // let decoded_bytes = if offset == u64::MAX && size == u64::MAX { + // // The chunk is just the fill value + // fill_value.repeat(chunk_subset.num_elements_usize()) + // } else { + // // The chunk must be decoded + // let partial_decoder = self.inner_codecs.partial_decoder( + // Box::new(ByteIntervalPartialDecoder::new( + // &*self.input_handle, + // offset, + // size, + // )), + // &chunk_representation, + // )?; + // let array_subset_in_chunk_subset = + // unsafe { overlap.relative_to_unchecked(chunk_subset.start()) }; + + // // Partial decoding is actually really slow with the blosc codec! Assume sharded chunks are small, and just decode the whole thing and extract bytes + // // TODO: Make this behaviour optional? + // // partial_decoder + // // .partial_decode(&[array_subset_in_chunk_subset.clone()])? + // // .remove(0) + // let decoded_chunk = partial_decoder + // .partial_decode(&[ArraySubset::new_with_shape( + // chunk_subset.shape().to_vec(), + // )])? + // .remove(0); + // array_subset_in_chunk_subset + // .extract_bytes(&decoded_chunk, chunk_subset.shape(), element_size) + // .unwrap() + // }; + + // // Copy decoded bytes to the output + // let chunk_subset_in_array_subset = + // unsafe { overlap.relative_to_unchecked(array_subset.start()) }; + // let mut decoded_offset = 0; + // for (array_subset_element_index, num_elements) in unsafe { + // chunk_subset_in_array_subset + // .iter_contiguous_linearised_indices_unchecked(array_subset.shape()) + // } { + // let output_offset = + // usize::try_from(array_subset_element_index * element_size_u64).unwrap(); + // let length = usize::try_from(num_elements * element_size_u64).unwrap(); + // out_array_subset[output_offset..output_offset + length] + // .copy_from_slice(&decoded_bytes[decoded_offset..decoded_offset + length]); + // decoded_offset += length; + // } + // } + // out.push(out_array_subset); + // } + // Ok(out) + // } } #[cfg(feature = "async")] @@ -355,7 +347,7 @@ impl<'a> AsyncShardingPartialDecoder<'a> { inner_codecs: &'a CodecChain, index_codecs: &'a CodecChain, index_location: ShardingIndexLocation, - parallel: bool, + options: &PartialDecodeOptions, ) -> Result, CodecError> { let shard_index = Self::decode_shard_index( &*input_handle, @@ -363,7 +355,7 @@ impl<'a> AsyncShardingPartialDecoder<'a> { index_location, chunk_shape.as_slice(), &decoded_representation, - parallel, + options, ) .await?; Ok(Self { @@ -382,7 +374,7 @@ impl<'a> AsyncShardingPartialDecoder<'a> { index_location: ShardingIndexLocation, chunk_shape: &[NonZeroU64], decoded_representation: &ChunkRepresentation, - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let shard_shape = decoded_representation.shape(); let chunk_representation = unsafe { @@ -412,7 +404,7 @@ impl<'a> AsyncShardingPartialDecoder<'a> { }; let encoded_shard_index = input_handle - .partial_decode_opt(&[index_byte_range], parallel) + .partial_decode_opt(&[index_byte_range], options) .await? .map(|mut v| v.remove(0)); @@ -422,7 +414,7 @@ impl<'a> AsyncShardingPartialDecoder<'a> { encoded_shard_index, &index_array_representation, index_codecs, - parallel, + options, ) .await?, ), @@ -434,117 +426,11 @@ impl<'a> AsyncShardingPartialDecoder<'a> { #[cfg(feature = "async")] #[async_trait::async_trait] impl AsyncArrayPartialDecoderTraits for AsyncShardingPartialDecoder<'_> { - async fn partial_decode_opt( - &self, - array_subsets: &[ArraySubset], - parallel: bool, - ) -> Result>, CodecError> { - if parallel { - self.par_partial_decode(array_subsets).await - } else { - self.partial_decode(array_subsets).await - } - } - - async fn partial_decode( - &self, - array_subsets: &[ArraySubset], - ) -> Result>, CodecError> { - let Some(shard_index) = &self.shard_index else { - return Ok(array_subsets - .iter() - .map(|array_subset| { - self.decoded_representation - .fill_value() - .as_ne_bytes() - .repeat(array_subset.num_elements_usize()) - }) - .collect()); - }; - - let chunk_representation = unsafe { - ChunkRepresentation::new_unchecked( - self.chunk_grid.chunk_shape().to_vec(), - self.decoded_representation.data_type().clone(), - self.decoded_representation.fill_value().clone(), - ) - }; - - let chunks_per_shard = calculate_chunks_per_shard( - self.decoded_representation.shape(), - chunk_representation.shape(), - ) - .map_err(|e| CodecError::Other(e.to_string()))?; - let chunks_per_shard = chunk_shape_to_array_shape(chunks_per_shard.as_slice()); - - let element_size = self.decoded_representation.element_size() as u64; - let fill_value = chunk_representation.fill_value().as_ne_bytes(); - - let mut out = Vec::with_capacity(array_subsets.len()); - for array_subset in array_subsets { - let array_subset_size = - usize::try_from(array_subset.num_elements() * element_size).unwrap(); - let mut out_array_subset = vec![0; array_subset_size]; - - // Decode those chunks if required and put in chunk cache - for (chunk_indices, chunk_subset) in - unsafe { array_subset.iter_chunks_unchecked(chunk_representation.shape()) } - { - let shard_index_index: usize = - usize::try_from(ravel_indices(&chunk_indices, &chunks_per_shard) * 2).unwrap(); - let offset = shard_index[shard_index_index]; - let size = shard_index[shard_index_index + 1]; - - let overlap = unsafe { array_subset.overlap_unchecked(&chunk_subset) }; - let decoded_bytes = if offset == u64::MAX && size == u64::MAX { - // The chunk is just the fill value - fill_value.repeat(chunk_subset.num_elements_usize()) - } else { - // The chunk must be decoded - let partial_decoder = self - .inner_codecs - .async_partial_decoder( - Box::new(AsyncByteIntervalPartialDecoder::new( - &*self.input_handle, - offset, - size, - )), - &chunk_representation, - ) - .await?; - let array_subset_in_chunk_subset = - unsafe { overlap.relative_to_unchecked(chunk_subset.start()) }; - partial_decoder - .partial_decode(&[array_subset_in_chunk_subset.clone()]) - .await? - .remove(0) - }; - - // Copy decoded bytes to the output - let chunk_subset_in_array_subset = - unsafe { overlap.relative_to_unchecked(array_subset.start()) }; - let mut decoded_offset = 0; - for (array_subset_element_index, num_elements) in unsafe { - chunk_subset_in_array_subset - .iter_contiguous_linearised_indices_unchecked(array_subset.shape()) - } { - let output_offset = - usize::try_from(array_subset_element_index * element_size).unwrap(); - let length = usize::try_from(num_elements * element_size).unwrap(); - out_array_subset[output_offset..output_offset + length] - .copy_from_slice(&decoded_bytes[decoded_offset..decoded_offset + length]); - decoded_offset += length; - } - } - out.push(out_array_subset); - } - Ok(out) - } - #[allow(clippy::too_many_lines)] - async fn par_partial_decode( + async fn partial_decode_opt( &self, array_subsets: &[ArraySubset], + options: &PartialDecodeOptions, ) -> Result>, CodecError> { let Some(shard_index) = &self.shard_index else { return Ok(array_subsets @@ -617,13 +503,14 @@ impl AsyncArrayPartialDecoderTraits for AsyncShardingPartialDecoder<'_> { }; let partial_decoder = self .inner_codecs - .async_partial_decoder( + .async_partial_decoder_opt( Box::new(AsyncByteIntervalPartialDecoder::new( &*self.input_handle, u64::try_from(*offset).unwrap(), u64::try_from(*size).unwrap(), )), &chunk_representation, + options, // FIXME: Adjust options for partial decoding ) .await?; let overlap = unsafe { array_subset.overlap_unchecked(chunk_subset) }; @@ -636,9 +523,10 @@ impl AsyncArrayPartialDecoderTraits for AsyncShardingPartialDecoder<'_> { // .await? // .remove(0); let decoded_chunk = partial_decoder - .partial_decode(&[ArraySubset::new_with_shape( - chunk_subset.shape().to_vec(), - )]) + .partial_decode_opt( + &[ArraySubset::new_with_shape(chunk_subset.shape().to_vec())], + options, + ) // FIXME: Adjust options for partial decoding .await? .remove(0); let decoded_chunk = array_subset_in_chunk_subset diff --git a/src/array/codec/array_to_bytes/zfp/zfp_codec.rs b/src/array/codec/array_to_bytes/zfp/zfp_codec.rs index 188b8670..1ec37cda 100644 --- a/src/array/codec/array_to_bytes/zfp/zfp_codec.rs +++ b/src/array/codec/array_to_bytes/zfp/zfp_codec.rs @@ -1,13 +1,17 @@ use zfp_sys::{ - zfp_compress, zfp_exec_policy_zfp_exec_omp, zfp_stream_maximum_size, zfp_stream_rewind, - zfp_stream_set_bit_stream, zfp_stream_set_execution, + zfp_compress, + zfp_stream_maximum_size, + zfp_stream_rewind, + zfp_stream_set_bit_stream, + // zfp_exec_policy_zfp_exec_omp, zfp_stream_set_execution }; use crate::{ array::{ codec::{ ArrayCodecTraits, ArrayPartialDecoderTraits, ArrayToBytesCodecTraits, - BytesPartialDecoderTraits, CodecError, CodecTraits, + BytesPartialDecoderTraits, CodecError, CodecTraits, DecodeOptions, EncodeOptions, + PartialDecoderOptions, RecommendedConcurrency, }, BytesRepresentation, ChunkRepresentation, DataType, }, @@ -123,11 +127,19 @@ impl CodecTraits for ZfpCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl ArrayCodecTraits for ZfpCodec { + fn recommended_concurrency( + &self, + _decoded_representation: &ChunkRepresentation, + ) -> Result { + // TODO: zfp supports multi thread, when is it optimal to kick in? + Ok(RecommendedConcurrency::one()) + } + fn encode_opt( &self, mut decoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + _options: &EncodeOptions, ) -> Result, CodecError> { let Some(zfp_type) = zarr_data_type_to_zfp_data_type(decoded_representation.data_type()) else { @@ -161,12 +173,13 @@ impl ArrayCodecTraits for ZfpCodec { zfp_stream_rewind(zfp.as_zfp_stream()); // needed? } - if parallel { - // Number of threads is set automatically - unsafe { - zfp_stream_set_execution(zfp.as_zfp_stream(), zfp_exec_policy_zfp_exec_omp); - } - } + // FIXME + // if parallel { + // // Number of threads is set automatically + // unsafe { + // zfp_stream_set_execution(zfp.as_zfp_stream(), zfp_exec_policy_zfp_exec_omp); + // } + // } // Compress array let size = unsafe { zfp_compress(zfp.as_zfp_stream(), field.as_zfp_field()) }; @@ -182,7 +195,7 @@ impl ArrayCodecTraits for ZfpCodec { &self, encoded_value: Vec, decoded_representation: &ChunkRepresentation, - parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { let Some(zfp_type) = zarr_data_type_to_zfp_data_type(decoded_representation.data_type()) else { @@ -195,7 +208,7 @@ impl ArrayCodecTraits for ZfpCodec { zfp_type, encoded_value, decoded_representation, - parallel, + false, // FIXME ) } } @@ -206,7 +219,7 @@ impl ArrayToBytesCodecTraits for ZfpCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(zfp_partial_decoder::ZfpPartialDecoder::new( input_handle, @@ -220,7 +233,7 @@ impl ArrayToBytesCodecTraits for ZfpCodec { &'a self, input_handle: Box, decoded_representation: &ChunkRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(zfp_partial_decoder::AsyncZfpPartialDecoder::new( input_handle, diff --git a/src/array/codec/array_to_bytes/zfp/zfp_partial_decoder.rs b/src/array/codec/array_to_bytes/zfp/zfp_partial_decoder.rs index d71ddcb8..0015575d 100644 --- a/src/array/codec/array_to_bytes/zfp/zfp_partial_decoder.rs +++ b/src/array/codec/array_to_bytes/zfp/zfp_partial_decoder.rs @@ -2,7 +2,9 @@ use zfp_sys::zfp_type; use crate::{ array::{ - codec::{ArrayPartialDecoderTraits, BytesPartialDecoderTraits, CodecError}, + codec::{ + ArrayPartialDecoderTraits, BytesPartialDecoderTraits, CodecError, PartialDecodeOptions, + }, ChunkRepresentation, }, array_subset::ArraySubset, @@ -51,9 +53,9 @@ impl ArrayPartialDecoderTraits for ZfpPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel)?; + let encoded_value = self.input_handle.decode_opt(options)?; let mut out = Vec::with_capacity(decoded_regions.len()); let chunk_shape = self.decoded_representation.shape_u64(); match encoded_value { @@ -63,7 +65,7 @@ impl ArrayPartialDecoderTraits for ZfpPartialDecoder<'_> { self.zfp_type, encoded_value, &self.decoded_representation, - parallel, + false, // FIXME )?; for array_subset in decoded_regions { let byte_ranges = unsafe { @@ -132,9 +134,9 @@ impl AsyncArrayPartialDecoderTraits for AsyncZfpPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel).await?; + let encoded_value = self.input_handle.decode_opt(options).await?; let chunk_shape = self.decoded_representation.shape_u64(); let mut out = Vec::with_capacity(decoded_regions.len()); match encoded_value { @@ -144,7 +146,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncZfpPartialDecoder<'_> { self.zfp_type, encoded_value, &self.decoded_representation, - parallel, + false, // FIXME )?; for array_subset in decoded_regions { let byte_ranges = unsafe { diff --git a/src/array/codec/byte_interval_partial_decoder.rs b/src/array/codec/byte_interval_partial_decoder.rs index 4d9d07e1..1dc9e653 100644 --- a/src/array/codec/byte_interval_partial_decoder.rs +++ b/src/array/codec/byte_interval_partial_decoder.rs @@ -1,6 +1,6 @@ use crate::byte_range::{ByteLength, ByteOffset, ByteRange}; -use super::{BytesPartialDecoderTraits, CodecError}; +use super::{BytesPartialDecoderTraits, CodecError, PartialDecodeOptions}; #[cfg(feature = "async")] use super::AsyncBytesPartialDecoderTraits; @@ -33,7 +33,7 @@ impl<'a> BytesPartialDecoderTraits for ByteIntervalPartialDecoder<'a> { fn partial_decode_opt( &self, byte_ranges: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { let byte_ranges: Vec = byte_ranges .iter() @@ -53,7 +53,7 @@ impl<'a> BytesPartialDecoderTraits for ByteIntervalPartialDecoder<'a> { ), }) .collect(); - self.inner.partial_decode_opt(&byte_ranges, parallel) + self.inner.partial_decode_opt(&byte_ranges, options) } } @@ -89,7 +89,7 @@ impl<'a> AsyncBytesPartialDecoderTraits for AsyncByteIntervalPartialDecoder<'a> async fn partial_decode_opt( &self, byte_ranges: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { let byte_ranges: Vec = byte_ranges .iter() @@ -109,6 +109,6 @@ impl<'a> AsyncBytesPartialDecoderTraits for AsyncByteIntervalPartialDecoder<'a> ), }) .collect(); - self.inner.partial_decode_opt(&byte_ranges, parallel).await + self.inner.partial_decode_opt(&byte_ranges, options).await } } diff --git a/src/array/codec/bytes_to_bytes/blosc/blosc_codec.rs b/src/array/codec/bytes_to_bytes/blosc/blosc_codec.rs index f284e1e8..03bd59ec 100644 --- a/src/array/codec/bytes_to_bytes/blosc/blosc_codec.rs +++ b/src/array/codec/bytes_to_bytes/blosc/blosc_codec.rs @@ -4,7 +4,10 @@ use blosc_sys::{blosc_get_complib_info, BLOSC_MAX_OVERHEAD}; use crate::{ array::{ - codec::{BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits}, + codec::{ + BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits, + DecodeOptions, EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, + }, BytesRepresentation, }, metadata::Metadata, @@ -135,34 +138,46 @@ impl CodecTraits for BloscCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl BytesToBytesCodecTraits for BloscCodec { - fn encode_opt(&self, decoded_value: Vec, parallel: bool) -> Result, CodecError> { - if parallel { - let n_threads = std::thread::available_parallelism().unwrap().get(); - self.do_encode(&decoded_value, n_threads) - } else { - self.do_encode(&decoded_value, 1) - } + fn recommended_concurrency( + &self, + _decoded_representation: &BytesRepresentation, + ) -> Result { + // TODO: Dependent on the block size, recommended concurrency could be > 1 + Ok(RecommendedConcurrency::one()) + } + + fn encode_opt( + &self, + decoded_value: Vec, + _options: &EncodeOptions, + ) -> Result, CodecError> { + // if options.is_parallel() { + // let n_threads = std::thread::available_parallelism().unwrap().get(); + // self.do_encode(&decoded_value, n_threads) + // } else { + self.do_encode(&decoded_value, 1) + // } } fn decode_opt( &self, encoded_value: Vec, _decoded_representation: &BytesRepresentation, - parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { - if parallel { - let n_threads = std::thread::available_parallelism().unwrap().get(); - Self::do_decode(&encoded_value, n_threads) - } else { - Self::do_decode(&encoded_value, 1) - } + // if options.is_parallel() { + // let n_threads = std::thread::available_parallelism().unwrap().get(); + // Self::do_decode(&encoded_value, n_threads) + // } else { + Self::do_decode(&encoded_value, 1) + // } } fn partial_decoder_opt<'a>( &'a self, input_handle: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _parallel: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(blosc_partial_decoder::BloscPartialDecoder::new( input_handle, @@ -174,7 +189,7 @@ impl BytesToBytesCodecTraits for BloscCodec { &'a self, input_handle: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _parallel: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( blosc_partial_decoder::AsyncBloscPartialDecoder::new(input_handle), diff --git a/src/array/codec/bytes_to_bytes/blosc/blosc_partial_decoder.rs b/src/array/codec/bytes_to_bytes/blosc/blosc_partial_decoder.rs index e05c2bd9..e521aad1 100644 --- a/src/array/codec/bytes_to_bytes/blosc/blosc_partial_decoder.rs +++ b/src/array/codec/bytes_to_bytes/blosc/blosc_partial_decoder.rs @@ -1,5 +1,8 @@ use crate::{ - array::codec::{bytes_to_bytes::blosc::blosc_nbytes, BytesPartialDecoderTraits, CodecError}, + array::codec::{ + bytes_to_bytes::blosc::blosc_nbytes, BytesPartialDecoderTraits, CodecError, + PartialDecodeOptions, + }, byte_range::ByteRange, }; @@ -23,9 +26,9 @@ impl BytesPartialDecoderTraits for BloscPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel)?; + let encoded_value = self.input_handle.decode_opt(options)?; let Some(encoded_value) = encoded_value else { return Ok(None); }; @@ -74,9 +77,9 @@ impl AsyncBytesPartialDecoderTraits for AsyncBloscPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel).await?; + let encoded_value = self.input_handle.decode_opt(options).await?; let Some(encoded_value) = encoded_value else { return Ok(None); }; diff --git a/src/array/codec/bytes_to_bytes/bz2/bz2_codec.rs b/src/array/codec/bytes_to_bytes/bz2/bz2_codec.rs index b73fa833..fe81d802 100644 --- a/src/array/codec/bytes_to_bytes/bz2/bz2_codec.rs +++ b/src/array/codec/bytes_to_bytes/bz2/bz2_codec.rs @@ -2,7 +2,10 @@ use std::io::Read; use crate::{ array::{ - codec::{BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits}, + codec::{ + BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits, + DecodeOptions, EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, + }, BytesRepresentation, }, metadata::Metadata, @@ -58,7 +61,19 @@ impl CodecTraits for Bz2Codec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl BytesToBytesCodecTraits for Bz2Codec { - fn encode_opt(&self, decoded_value: Vec, _parallel: bool) -> Result, CodecError> { + fn recommended_concurrency( + &self, + _decoded_representation: &BytesRepresentation, + ) -> Result { + // bz2 does not support parallel decode + Ok(RecommendedConcurrency::one()) + } + + fn encode_opt( + &self, + decoded_value: Vec, + _options: &EncodeOptions, + ) -> Result, CodecError> { let mut encoder = bzip2::read::BzEncoder::new(decoded_value.as_slice(), self.compression); let mut out: Vec = Vec::new(); encoder.read_to_end(&mut out)?; @@ -69,7 +84,7 @@ impl BytesToBytesCodecTraits for Bz2Codec { &self, encoded_value: Vec, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { let mut decoder = bzip2::read::BzDecoder::new(encoded_value.as_slice()); let mut out: Vec = Vec::new(); @@ -81,7 +96,7 @@ impl BytesToBytesCodecTraits for Bz2Codec { &'a self, input_handle: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(bz2_partial_decoder::Bz2PartialDecoder::new( input_handle, @@ -93,7 +108,7 @@ impl BytesToBytesCodecTraits for Bz2Codec { &'a self, input_handle: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(bz2_partial_decoder::AsyncBz2PartialDecoder::new( input_handle, diff --git a/src/array/codec/bytes_to_bytes/bz2/bz2_partial_decoder.rs b/src/array/codec/bytes_to_bytes/bz2/bz2_partial_decoder.rs index 662d0c9b..197e7a25 100644 --- a/src/array/codec/bytes_to_bytes/bz2/bz2_partial_decoder.rs +++ b/src/array/codec/bytes_to_bytes/bz2/bz2_partial_decoder.rs @@ -1,7 +1,7 @@ use std::io::Read; use crate::{ - array::codec::{BytesPartialDecoderTraits, CodecError}, + array::codec::{BytesPartialDecoderTraits, CodecError, PartialDecodeOptions}, byte_range::{extract_byte_ranges, ByteRange}, }; @@ -23,9 +23,9 @@ impl BytesPartialDecoderTraits for Bz2PartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel)?; + let encoded_value = self.input_handle.decode_opt(options)?; let Some(encoded_value) = encoded_value else { return Ok(None); }; @@ -60,9 +60,9 @@ impl AsyncBytesPartialDecoderTraits for AsyncBz2PartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel).await?; + let encoded_value = self.input_handle.decode_opt(options).await?; let Some(encoded_value) = encoded_value else { return Ok(None); }; diff --git a/src/array/codec/bytes_to_bytes/crc32c/crc32c_codec.rs b/src/array/codec/bytes_to_bytes/crc32c/crc32c_codec.rs index fc5f1643..88d9f75a 100644 --- a/src/array/codec/bytes_to_bytes/crc32c/crc32c_codec.rs +++ b/src/array/codec/bytes_to_bytes/crc32c/crc32c_codec.rs @@ -1,6 +1,9 @@ use crate::{ array::{ - codec::{BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits}, + codec::{ + BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits, + DecodeOptions, EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, + }, BytesRepresentation, }, metadata::Metadata, @@ -49,10 +52,17 @@ impl CodecTraits for Crc32cCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl BytesToBytesCodecTraits for Crc32cCodec { + fn recommended_concurrency( + &self, + _decoded_representation: &BytesRepresentation, + ) -> Result { + Ok(RecommendedConcurrency::one()) + } + fn encode_opt( &self, mut decoded_value: Vec, - _parallel: bool, + _options: &EncodeOptions, ) -> Result, CodecError> { let checksum = crc32c::crc32c(&decoded_value).to_le_bytes(); decoded_value.reserve_exact(checksum.len()); @@ -64,7 +74,7 @@ impl BytesToBytesCodecTraits for Crc32cCodec { &self, mut encoded_value: Vec, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { if encoded_value.len() >= CHECKSUM_SIZE { if crate::config::global_config().validate_checksums() { @@ -87,7 +97,7 @@ impl BytesToBytesCodecTraits for Crc32cCodec { &'a self, input_handle: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(crc32c_partial_decoder::Crc32cPartialDecoder::new( input_handle, @@ -99,7 +109,7 @@ impl BytesToBytesCodecTraits for Crc32cCodec { &'a self, input_handle: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( crc32c_partial_decoder::AsyncCrc32cPartialDecoder::new(input_handle), diff --git a/src/array/codec/bytes_to_bytes/crc32c/crc32c_partial_decoder.rs b/src/array/codec/bytes_to_bytes/crc32c/crc32c_partial_decoder.rs index eccbfc1c..b0a9f230 100644 --- a/src/array/codec/bytes_to_bytes/crc32c/crc32c_partial_decoder.rs +++ b/src/array/codec/bytes_to_bytes/crc32c/crc32c_partial_decoder.rs @@ -1,5 +1,5 @@ use crate::{ - array::codec::{BytesPartialDecoderTraits, CodecError}, + array::codec::{BytesPartialDecoderTraits, CodecError, PartialDecodeOptions}, byte_range::ByteRange, }; @@ -24,11 +24,11 @@ impl BytesPartialDecoderTraits for Crc32cPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { let bytes = self .input_handle - .partial_decode_opt(decoded_regions, parallel)?; + .partial_decode_opt(decoded_regions, options)?; let Some(mut bytes) = bytes else { return Ok(None); }; @@ -73,11 +73,11 @@ impl AsyncBytesPartialDecoderTraits for AsyncCrc32cPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { let bytes = self .input_handle - .partial_decode_opt(decoded_regions, parallel) + .partial_decode_opt(decoded_regions, options) .await?; let Some(mut bytes) = bytes else { return Ok(None); diff --git a/src/array/codec/bytes_to_bytes/gzip/gzip_codec.rs b/src/array/codec/bytes_to_bytes/gzip/gzip_codec.rs index 3e38af13..12c8d9c2 100644 --- a/src/array/codec/bytes_to_bytes/gzip/gzip_codec.rs +++ b/src/array/codec/bytes_to_bytes/gzip/gzip_codec.rs @@ -4,7 +4,10 @@ use flate2::bufread::{GzDecoder, GzEncoder}; use crate::{ array::{ - codec::{BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits}, + codec::{ + BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits, + DecodeOptions, EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, + }, BytesRepresentation, }, metadata::Metadata, @@ -64,7 +67,18 @@ impl CodecTraits for GzipCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl BytesToBytesCodecTraits for GzipCodec { - fn encode_opt(&self, decoded_value: Vec, _parallel: bool) -> Result, CodecError> { + fn recommended_concurrency( + &self, + _decoded_representation: &BytesRepresentation, + ) -> Result { + Ok(RecommendedConcurrency::one()) + } + + fn encode_opt( + &self, + decoded_value: Vec, + _options: &EncodeOptions, + ) -> Result, CodecError> { let mut encoder = GzEncoder::new( Cursor::new(decoded_value), flate2::Compression::new(self.compression_level.as_u32()), @@ -78,7 +92,7 @@ impl BytesToBytesCodecTraits for GzipCodec { &self, encoded_value: Vec, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { let mut decoder = GzDecoder::new(Cursor::new(encoded_value)); let mut out: Vec = Vec::new(); @@ -90,7 +104,7 @@ impl BytesToBytesCodecTraits for GzipCodec { &self, r: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(gzip_partial_decoder::GzipPartialDecoder::new(r))) } @@ -100,7 +114,7 @@ impl BytesToBytesCodecTraits for GzipCodec { &'a self, r: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( gzip_partial_decoder::AsyncGzipPartialDecoder::new(r), diff --git a/src/array/codec/bytes_to_bytes/gzip/gzip_partial_decoder.rs b/src/array/codec/bytes_to_bytes/gzip/gzip_partial_decoder.rs index bdad5e26..d9b159a9 100644 --- a/src/array/codec/bytes_to_bytes/gzip/gzip_partial_decoder.rs +++ b/src/array/codec/bytes_to_bytes/gzip/gzip_partial_decoder.rs @@ -3,7 +3,7 @@ use std::io::{Cursor, Read}; use flate2::bufread::GzDecoder; use crate::{ - array::codec::{BytesPartialDecoderTraits, CodecError}, + array::codec::{BytesPartialDecoderTraits, CodecError, PartialDecodeOptions}, byte_range::{extract_byte_ranges, ByteRange}, }; @@ -26,9 +26,9 @@ impl BytesPartialDecoderTraits for GzipPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel)?; + let encoded_value = self.input_handle.decode_opt(options)?; let Some(encoded_value) = encoded_value else { return Ok(None); }; @@ -64,9 +64,9 @@ impl AsyncBytesPartialDecoderTraits for AsyncGzipPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel).await?; + let encoded_value = self.input_handle.decode_opt(options).await?; let Some(encoded_value) = encoded_value else { return Ok(None); }; diff --git a/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_codec.rs b/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_codec.rs index c57088be..ebd68075 100644 --- a/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_codec.rs +++ b/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_codec.rs @@ -1,6 +1,9 @@ use crate::{ array::{ - codec::{BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits}, + codec::{ + BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits, + DecodeOptions, EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, + }, BytesRepresentation, }, metadata::Metadata, @@ -41,7 +44,19 @@ impl CodecTraits for TestUnboundedCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl BytesToBytesCodecTraits for TestUnboundedCodec { - fn encode_opt(&self, decoded_value: Vec, _parallel: bool) -> Result, CodecError> { + /// Return the maximum internal concurrency supported for the requested decoded representation. + fn recommended_concurrency( + &self, + _decoded_representation: &BytesRepresentation, + ) -> Result { + Ok(RecommendedConcurrency::one()) + } + + fn encode_opt( + &self, + decoded_value: Vec, + _options: &EncodeOptions, + ) -> Result, CodecError> { Ok(decoded_value) } @@ -49,7 +64,7 @@ impl BytesToBytesCodecTraits for TestUnboundedCodec { &self, encoded_value: Vec, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { Ok(encoded_value) } @@ -58,7 +73,7 @@ impl BytesToBytesCodecTraits for TestUnboundedCodec { &self, r: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( test_unbounded_partial_decoder::TestUnboundedPartialDecoder::new(r), @@ -70,7 +85,7 @@ impl BytesToBytesCodecTraits for TestUnboundedCodec { &'a self, r: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( test_unbounded_partial_decoder::AsyncTestUnboundedPartialDecoder::new(r), diff --git a/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_partial_decoder.rs b/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_partial_decoder.rs index 47b24755..59f3b5e2 100644 --- a/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_partial_decoder.rs +++ b/src/array/codec/bytes_to_bytes/test_unbounded/test_unbounded_partial_decoder.rs @@ -1,5 +1,5 @@ use crate::{ - array::codec::{BytesPartialDecoderTraits, CodecError}, + array::codec::{BytesPartialDecoderTraits, CodecError, PartialDecodeOptions}, byte_range::{extract_byte_ranges, ByteRange}, }; @@ -22,9 +22,9 @@ impl BytesPartialDecoderTraits for TestUnboundedPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel)?; + let encoded_value = self.input_handle.decode_opt(options)?; let Some(encoded_value) = encoded_value else { return Ok(None); }; @@ -56,9 +56,9 @@ impl AsyncBytesPartialDecoderTraits for AsyncTestUnboundedPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel).await?; + let encoded_value = self.input_handle.decode_opt(options).await?; let Some(encoded_value) = encoded_value else { return Ok(None); }; diff --git a/src/array/codec/bytes_to_bytes/zstd.rs b/src/array/codec/bytes_to_bytes/zstd.rs index 8ebac0ac..d292799a 100644 --- a/src/array/codec/bytes_to_bytes/zstd.rs +++ b/src/array/codec/bytes_to_bytes/zstd.rs @@ -67,9 +67,9 @@ mod tests { let decoded = codec.decode(encoded, &bytes_representation).unwrap(); assert_eq!(bytes, decoded); - let encoded = codec.par_encode(bytes.clone()).unwrap(); - let decoded = codec.par_decode(encoded, &bytes_representation).unwrap(); - assert_eq!(bytes, decoded); + // let encoded = codec.par_encode(bytes.clone()).unwrap(); + // let decoded = codec.par_decode(encoded, &bytes_representation).unwrap(); + // assert_eq!(bytes, decoded); } #[test] diff --git a/src/array/codec/bytes_to_bytes/zstd/zstd_codec.rs b/src/array/codec/bytes_to_bytes/zstd/zstd_codec.rs index 61357346..8f7112f9 100644 --- a/src/array/codec/bytes_to_bytes/zstd/zstd_codec.rs +++ b/src/array/codec/bytes_to_bytes/zstd/zstd_codec.rs @@ -2,7 +2,10 @@ use zstd::zstd_safe; use crate::{ array::{ - codec::{BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits}, + codec::{ + BytesPartialDecoderTraits, BytesToBytesCodecTraits, CodecError, CodecTraits, + DecodeOptions, EncodeOptions, PartialDecoderOptions, RecommendedConcurrency, + }, BytesRepresentation, }, metadata::Metadata, @@ -61,14 +64,26 @@ impl CodecTraits for ZstdCodec { #[cfg_attr(feature = "async", async_trait::async_trait)] impl BytesToBytesCodecTraits for ZstdCodec { - fn encode_opt(&self, decoded_value: Vec, parallel: bool) -> Result, CodecError> { + fn recommended_concurrency( + &self, + _decoded_representation: &BytesRepresentation, + ) -> Result { + // TODO: zstd supports multithread, but at what point is it good to kick in? + Ok(RecommendedConcurrency::one()) + } + + fn encode_opt( + &self, + decoded_value: Vec, + _options: &EncodeOptions, + ) -> Result, CodecError> { let mut result = Vec::::new(); let mut encoder = zstd::Encoder::new(&mut result, self.compression)?; encoder.include_checksum(self.checksum)?; - if parallel { - let n_threads = std::thread::available_parallelism().unwrap().get(); - encoder.multithread(u32::try_from(n_threads).unwrap())?; // TODO: Check overhead of zstd par_encode - } + // if parallel { + // let n_threads = std::thread::available_parallelism().unwrap().get(); + // encoder.multithread(u32::try_from(n_threads).unwrap())?; // TODO: Check overhead of zstd par_encode + // } std::io::copy(&mut decoded_value.as_slice(), &mut encoder)?; encoder.finish()?; Ok(result) @@ -78,7 +93,7 @@ impl BytesToBytesCodecTraits for ZstdCodec { &self, encoded_value: Vec, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &DecodeOptions, ) -> Result, CodecError> { zstd::decode_all(encoded_value.as_slice()).map_err(CodecError::IOError) } @@ -87,7 +102,7 @@ impl BytesToBytesCodecTraits for ZstdCodec { &self, r: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new(zstd_partial_decoder::ZstdPartialDecoder::new(r))) } @@ -97,7 +112,7 @@ impl BytesToBytesCodecTraits for ZstdCodec { &'a self, r: Box, _decoded_representation: &BytesRepresentation, - _parallel: bool, + _options: &PartialDecoderOptions, ) -> Result, CodecError> { Ok(Box::new( zstd_partial_decoder::AsyncZstdPartialDecoder::new(r), diff --git a/src/array/codec/bytes_to_bytes/zstd/zstd_partial_decoder.rs b/src/array/codec/bytes_to_bytes/zstd/zstd_partial_decoder.rs index 8f60bc60..187fad9d 100644 --- a/src/array/codec/bytes_to_bytes/zstd/zstd_partial_decoder.rs +++ b/src/array/codec/bytes_to_bytes/zstd/zstd_partial_decoder.rs @@ -1,5 +1,5 @@ use crate::{ - array::codec::{BytesPartialDecoderTraits, CodecError}, + array::codec::{BytesPartialDecoderTraits, CodecError, PartialDecodeOptions}, byte_range::{extract_byte_ranges, ByteRange}, }; @@ -22,9 +22,9 @@ impl BytesPartialDecoderTraits for ZstdPartialDecoder<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel)?; + let encoded_value = self.input_handle.decode_opt(options)?; let Some(encoded_value) = encoded_value else { return Ok(None); }; @@ -59,9 +59,9 @@ impl AsyncBytesPartialDecoderTraits for AsyncZstdPartialDecoder<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - let encoded_value = self.input_handle.decode_opt(parallel).await?; + let encoded_value = self.input_handle.decode_opt(options).await?; let Some(encoded_value) = encoded_value else { return Ok(None); }; diff --git a/src/array/codec/options.rs b/src/array/codec/options.rs new file mode 100644 index 00000000..b72be45d --- /dev/null +++ b/src/array/codec/options.rs @@ -0,0 +1,73 @@ +//! Options for codec encoding and decoding. + +use std::num::NonZeroUsize; + +/// Encode options. +pub struct EncodeOptions { + concurrent_limit: NonZeroUsize, +} + +impl Default for EncodeOptions { + fn default() -> Self { + Self { + concurrent_limit: std::thread::available_parallelism().unwrap(), + } + } +} + +impl EncodeOptions { + /// Return the concurrent limit. + #[must_use] + pub fn concurrent_limit(&self) -> NonZeroUsize { + self.concurrent_limit + } + + /// Set the concurrent limit. + pub fn set_concurrent_limit(&mut self, concurrent_limit: NonZeroUsize) { + self.concurrent_limit = concurrent_limit; + } + + /// FIXME: Temporary, remove + #[must_use] + pub fn is_parallel(&self) -> bool { + self.concurrent_limit.get() > 1 + } +} + +/// Decode options. +pub struct DecodeOptions { + concurrent_limit: NonZeroUsize, +} + +impl Default for DecodeOptions { + fn default() -> Self { + Self { + concurrent_limit: std::thread::available_parallelism().unwrap(), + } + } +} + +impl DecodeOptions { + /// Return the concurrent limit. + #[must_use] + pub fn concurrent_limit(&self) -> NonZeroUsize { + self.concurrent_limit + } + + /// Set the concurrent limit. + pub fn set_concurrent_limit(&mut self, concurrent_limit: NonZeroUsize) { + self.concurrent_limit = concurrent_limit; + } + + /// FIXME: Temporary, remove + #[must_use] + pub fn is_parallel(&self) -> bool { + self.concurrent_limit.get() > 1 + } +} + +/// Partial decoder options. +pub type PartialDecoderOptions = DecodeOptions; + +/// Partial decode options. +pub type PartialDecodeOptions = DecodeOptions; diff --git a/src/array/codec/partial_decoder_cache.rs b/src/array/codec/partial_decoder_cache.rs index 2e8d9834..7a5fed6f 100644 --- a/src/array/codec/partial_decoder_cache.rs +++ b/src/array/codec/partial_decoder_cache.rs @@ -1,5 +1,7 @@ //! A cache for partial decoders. +// TODO: Move BytesPartialDecoderCache and ArrayPartialDecoderCache into separate files + use std::marker::PhantomData; use crate::{ @@ -8,7 +10,10 @@ use crate::{ byte_range::{extract_byte_ranges, ByteRange}, }; -use super::{ArrayPartialDecoderTraits, ArraySubset, BytesPartialDecoderTraits, CodecError}; +use super::{ + ArrayPartialDecoderTraits, ArraySubset, BytesPartialDecoderTraits, CodecError, + PartialDecodeOptions, +}; #[cfg(feature = "async")] use super::{AsyncArrayPartialDecoderTraits, AsyncBytesPartialDecoderTraits}; @@ -26,10 +31,10 @@ impl<'a> BytesPartialDecoderCache<'a> { /// Returns a [`CodecError`] if caching fails. pub fn new( input_handle: &dyn BytesPartialDecoderTraits, - parallel: bool, + options: &PartialDecodeOptions, ) -> Result { let cache = input_handle - .partial_decode_opt(&[ByteRange::FromStart(0, None)], parallel)? + .partial_decode_opt(&[ByteRange::FromStart(0, None)], options)? .map(|mut bytes| bytes.remove(0)); Ok(Self { cache, @@ -44,10 +49,10 @@ impl<'a> BytesPartialDecoderCache<'a> { /// Returns a [`CodecError`] if caching fails. pub async fn async_new( input_handle: &dyn AsyncBytesPartialDecoderTraits, - parallel: bool, + options: &PartialDecodeOptions, ) -> Result, CodecError> { let cache = input_handle - .partial_decode_opt(&[ByteRange::FromStart(0, None)], parallel) + .partial_decode_opt(&[ByteRange::FromStart(0, None)], options) .await? .map(|mut bytes| bytes.remove(0)); Ok(Self { @@ -61,7 +66,7 @@ impl BytesPartialDecoderTraits for BytesPartialDecoderCache<'_> { fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - _parallel: bool, + _options: &PartialDecodeOptions, ) -> Result>>, CodecError> { Ok(match &self.cache { Some(bytes) => Some( @@ -79,9 +84,9 @@ impl AsyncBytesPartialDecoderTraits for BytesPartialDecoderCache<'_> { async fn partial_decode_opt( &self, decoded_regions: &[ByteRange], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>>, CodecError> { - BytesPartialDecoderTraits::partial_decode_opt(self, decoded_regions, parallel) + BytesPartialDecoderTraits::partial_decode_opt(self, decoded_regions, options) } } @@ -100,14 +105,14 @@ impl<'a> ArrayPartialDecoderCache<'a> { pub fn new( input_handle: &dyn ArrayPartialDecoderTraits, decoded_representation: ChunkRepresentation, - parallel: bool, + options: &PartialDecodeOptions, ) -> Result { let cache = input_handle .partial_decode_opt( &[ArraySubset::new_with_shape( decoded_representation.shape_u64(), )], - parallel, + options, )? .remove(0); Ok(Self { @@ -125,14 +130,14 @@ impl<'a> ArrayPartialDecoderCache<'a> { pub async fn async_new( input_handle: &dyn AsyncArrayPartialDecoderTraits, decoded_representation: ChunkRepresentation, - parallel: bool, + options: &PartialDecodeOptions, ) -> Result, CodecError> { let cache = input_handle .partial_decode_opt( &[ArraySubset::new_with_shape( decoded_representation.shape_u64(), )], - parallel, + options, ) .await? .remove(0); @@ -148,7 +153,7 @@ impl<'a> ArrayPartialDecoderTraits for ArrayPartialDecoderCache<'a> { fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - _parallel: bool, + _options: &PartialDecodeOptions, ) -> Result>, CodecError> { let mut out: Vec> = Vec::with_capacity(decoded_regions.len()); let array_shape = self.decoded_representation.shape_u64(); @@ -175,8 +180,8 @@ impl<'a> AsyncArrayPartialDecoderTraits for ArrayPartialDecoderCache<'a> { async fn partial_decode_opt( &self, decoded_regions: &[ArraySubset], - parallel: bool, + options: &PartialDecodeOptions, ) -> Result>, CodecError> { - ArrayPartialDecoderTraits::partial_decode_opt(self, decoded_regions, parallel) + ArrayPartialDecoderTraits::partial_decode_opt(self, decoded_regions, options) } }