From 65542ccb0ac5727930a7687c5bd6c2e84e5b5f5e Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Thu, 19 Sep 2024 12:01:24 +1000 Subject: [PATCH] Change `Contiguous[Linearised]Indices` iterators to return indices only The fixed number of contiguous indices can be separately requested --- CHANGELOG.md | 5 +++- zarrs/src/array/array_bytes.rs | 2 +- .../array_to_bytes/sharding/sharding_codec.rs | 3 +- zarrs/src/array_subset.rs | 4 +-- zarrs/src/array_subset/iterators.rs | 16 ++++++---- .../iterators/contiguous_indices_iterator.rs | 29 ++++++++++++++----- .../contiguous_linearised_indices_iterator.rs | 25 +++++++++++++--- 7 files changed, 61 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e5b23335..3303445d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,13 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed +- **Breaking**: Change `Contiguous[Linearised]Indices` iterators to return indices only + ### Fixed - Fixed an unnecessary copy in `Array::[async_]retrieve_chunk_if_exists_opt` - Fixed `CodecOptions` not being forwarded in `Array::retrieve_chunk_subset_opt` on the fast path - Fixed missing fast path in `Array::[async_]retrieve_chunk_subset_opt` ### Removed -- Removed `ArraySubset::extract_elements[_unchecked]` which were unused and incorrect +- **Breaking**: Removed `ArraySubset::extract_elements[_unchecked]` which were unused and incorrect ## [0.17.0-beta.1] - 2024-09-16 diff --git a/zarrs/src/array/array_bytes.rs b/zarrs/src/array/array_bytes.rs index c59f0602..2be37255 100644 --- a/zarrs/src/array/array_bytes.rs +++ b/zarrs/src/array/array_bytes.rs @@ -281,7 +281,7 @@ pub fn update_bytes_flen( let length = contiguous_indices.contiguous_elements_usize() * data_type_size; let mut decoded_offset = 0; // TODO: Par iteration? - for (array_subset_element_index, _num_elements) in &contiguous_indices { + for array_subset_element_index in &contiguous_indices { let output_offset = usize::try_from(array_subset_element_index).unwrap() * data_type_size; debug_assert!((output_offset + length) <= output_bytes.len()); debug_assert!((decoded_offset + length) <= subset_bytes.len()); diff --git a/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs b/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs index df12e643..8df9b195 100644 --- a/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs +++ b/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs @@ -292,7 +292,8 @@ impl ArrayToBytesCodecTraits for ShardingCodec { chunk_subset .contiguous_linearised_indices_unchecked(&shard_shape) }; - for (index, elements) in &contiguous_iterator { + let elements = contiguous_iterator.contiguous_elements(); + for index in &contiguous_iterator { debug_assert_eq!( fv.len() as u64, elements * data_type_size as u64 diff --git a/zarrs/src/array_subset.rs b/zarrs/src/array_subset.rs index 087743cc..2e297cf0 100644 --- a/zarrs/src/array_subset.rs +++ b/zarrs/src/array_subset.rs @@ -285,7 +285,7 @@ impl ArraySubset { let mut byte_ranges: Vec = Vec::new(); let contiguous_indices = self.contiguous_linearised_indices(array_shape)?; let byte_length = contiguous_indices.contiguous_elements_usize() * element_size; - for (array_index, _contiguous_elements) in &contiguous_indices { + for array_index in &contiguous_indices { let byte_index = array_index * element_size as u64; byte_ranges.push(ByteRange::FromStart(byte_index, Some(byte_length as u64))); } @@ -305,7 +305,7 @@ impl ArraySubset { let mut byte_ranges: Vec = Vec::new(); let contiguous_indices = self.contiguous_linearised_indices_unchecked(array_shape); let byte_length = contiguous_indices.contiguous_elements_usize() * element_size; - for (array_index, _contiguous_elements) in &contiguous_indices { + for array_index in &contiguous_indices { let byte_index = array_index * element_size as u64; byte_ranges.push(ByteRange::FromStart(byte_index, Some(byte_length as u64))); } diff --git a/zarrs/src/array_subset/iterators.rs b/zarrs/src/array_subset/iterators.rs index c6e8313b..30126cc6 100644 --- a/zarrs/src/array_subset/iterators.rs +++ b/zarrs/src/array_subset/iterators.rs @@ -102,7 +102,8 @@ mod tests { let indices = subset.contiguous_indices(&[2, 2]).unwrap(); let mut iter = indices.into_iter(); assert_eq!(iter.size_hint(), (1, Some(1))); - assert_eq!(iter.next(), Some((vec![0, 0], 4))); + assert_eq!(iter.contiguous_elements(), 4); + assert_eq!(iter.next(), Some(vec![0, 0])); assert_eq!(iter.next(), None); } @@ -115,8 +116,9 @@ mod tests { assert_eq!(indices.contiguous_elements_usize(), 2); let mut iter = indices.iter(); assert_eq!(iter.size_hint(), (2, Some(2))); - assert_eq!(iter.next_back(), Some((vec![2, 1], 2))); - assert_eq!(iter.next(), Some((vec![1, 1], 2))); + assert_eq!(iter.contiguous_elements(), 2); + assert_eq!(iter.next_back(), Some(vec![2, 1])); + assert_eq!(iter.next(), Some(vec![1, 1])); assert_eq!(iter.next(), None); } @@ -126,7 +128,8 @@ mod tests { let indices = subset.contiguous_indices(&[3, 1, 2, 2]).unwrap(); let mut iter = indices.into_iter(); assert_eq!(iter.size_hint(), (1, Some(1))); - assert_eq!(iter.next(), Some((vec![1, 0, 0, 0], 8))); + assert_eq!(iter.contiguous_elements(), 8); + assert_eq!(iter.next(), Some(vec![1, 0, 0, 0])); assert_eq!(iter.next(), None); } @@ -143,8 +146,9 @@ mod tests { // 8 9 10 11 // 12 13 14 15 assert_eq!(iter.size_hint(), (2, Some(2))); - assert_eq!(iter.next_back(), Some((9, 2))); - assert_eq!(iter.next(), Some((5, 2))); + assert_eq!(iter.contiguous_elements(), 2); + assert_eq!(iter.next_back(), Some(9)); + assert_eq!(iter.next(), Some(5)); assert_eq!(iter.next(), None); } diff --git a/zarrs/src/array_subset/iterators/contiguous_indices_iterator.rs b/zarrs/src/array_subset/iterators/contiguous_indices_iterator.rs index 118e984d..2cec7b6f 100644 --- a/zarrs/src/array_subset/iterators/contiguous_indices_iterator.rs +++ b/zarrs/src/array_subset/iterators/contiguous_indices_iterator.rs @@ -129,7 +129,7 @@ impl ContiguousIndices { } impl<'a> IntoIterator for &'a ContiguousIndices { - type Item = (ArrayIndices, u64); + type Item = ArrayIndices; type IntoIter = ContiguousIndicesIterator<'a>; fn into_iter(self) -> Self::IntoIter { @@ -148,13 +148,28 @@ pub struct ContiguousIndicesIterator<'a> { contiguous_elements: u64, } +impl ContiguousIndicesIterator<'_> { + /// Return the number of contiguous elements (fixed on each iteration). + #[must_use] + pub fn contiguous_elements(&self) -> u64 { + self.contiguous_elements + } + + /// Return the number of contiguous elements (fixed on each iteration). + /// + /// # Panics + /// Panics if the number of contiguous elements exceeds [`usize::MAX`]. + #[must_use] + pub fn contiguous_elements_usize(&self) -> usize { + usize::try_from(self.contiguous_elements).unwrap() + } +} + impl Iterator for ContiguousIndicesIterator<'_> { - type Item = (ArrayIndices, u64); + type Item = ArrayIndices; fn next(&mut self) -> Option { - self.inner - .next() - .map(|indices| (indices, self.contiguous_elements)) + self.inner.next() } fn size_hint(&self) -> (usize, Option) { @@ -164,9 +179,7 @@ impl Iterator for ContiguousIndicesIterator<'_> { impl DoubleEndedIterator for ContiguousIndicesIterator<'_> { fn next_back(&mut self) -> Option { - self.inner - .next_back() - .map(|indices| (indices, self.contiguous_elements)) + self.inner.next_back() } } diff --git a/zarrs/src/array_subset/iterators/contiguous_linearised_indices_iterator.rs b/zarrs/src/array_subset/iterators/contiguous_linearised_indices_iterator.rs index 94a320d8..c3ac9263 100644 --- a/zarrs/src/array_subset/iterators/contiguous_linearised_indices_iterator.rs +++ b/zarrs/src/array_subset/iterators/contiguous_linearised_indices_iterator.rs @@ -92,7 +92,7 @@ impl ContiguousLinearisedIndices { } impl<'a> IntoIterator for &'a ContiguousLinearisedIndices { - type Item = (u64, u64); + type Item = u64; type IntoIter = ContiguousLinearisedIndicesIterator<'a>; fn into_iter(self) -> Self::IntoIter { @@ -111,13 +111,30 @@ pub struct ContiguousLinearisedIndicesIterator<'a> { array_shape: &'a [u64], } +impl ContiguousLinearisedIndicesIterator<'_> { + /// Return the number of contiguous elements (fixed on each iteration). + #[must_use] + pub fn contiguous_elements(&self) -> u64 { + self.inner.contiguous_elements() + } + + /// Return the number of contiguous elements (fixed on each iteration). + /// + /// # Panics + /// Panics if the number of contiguous elements exceeds [`usize::MAX`]. + #[must_use] + pub fn contiguous_elements_usize(&self) -> usize { + self.inner.contiguous_elements_usize() + } +} + impl Iterator for ContiguousLinearisedIndicesIterator<'_> { - type Item = (u64, u64); + type Item = u64; fn next(&mut self) -> Option { self.inner .next() - .map(|(indices, elements)| (ravel_indices(&indices, self.array_shape), elements)) + .map(|indices| ravel_indices(&indices, self.array_shape)) } fn size_hint(&self) -> (usize, Option) { @@ -129,7 +146,7 @@ impl DoubleEndedIterator for ContiguousLinearisedIndicesIterator<'_> { fn next_back(&mut self) -> Option { self.inner .next_back() - .map(|(indices, elements)| (ravel_indices(&indices, self.array_shape), elements)) + .map(|indices| ravel_indices(&indices, self.array_shape)) } }