From bac776f47906ee6b1e65df32ed24d77b23a8777b Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Wed, 1 Jan 2025 11:34:04 +1100 Subject: [PATCH] fix: workaround for trait upcasting --- .../array/array_sync_sharded_readable_ext.rs | 93 +++++++++++++++---- .../array_to_bytes/sharding/sharding_codec.rs | 8 +- zarrs/src/lib.rs | 2 - 3 files changed, 79 insertions(+), 24 deletions(-) diff --git a/zarrs/src/array/array_sync_sharded_readable_ext.rs b/zarrs/src/array/array_sync_sharded_readable_ext.rs index 0062aad2..461df909 100644 --- a/zarrs/src/array/array_sync_sharded_readable_ext.rs +++ b/zarrs/src/array/array_sync_sharded_readable_ext.rs @@ -1,14 +1,15 @@ -use std::any::Any; use std::{collections::HashMap, sync::Arc}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; use rayon_iter_concurrent_limit::iter_concurrent_limit; use unsafe_cell_slice::UnsafeCellSlice; +use zarrs_metadata::v3::array::codec::sharding::ShardingCodecConfiguration; use zarrs_storage::byte_range::ByteRange; use zarrs_storage::StorageHandle; use super::array_bytes::{merge_chunks_vlen, update_bytes_flen}; use super::codec::array_to_bytes::sharding::ShardingPartialDecoder; +use super::codec::{CodecError, ShardingCodec}; use super::element::ElementOwned; use super::{ codec::CodecOptions, concurrency::concurrency_chunks_and_codec, Array, ArrayError, @@ -19,7 +20,29 @@ use crate::array::codec::StoragePartialDecoder; use crate::storage::ReadableStorageTraits; use crate::{array::codec::ArrayPartialDecoderTraits, array_subset::ArraySubset}; -type PartialDecoderHashMap = HashMap, Arc>; +// TODO: Remove with trait upcasting +#[derive(Clone)] +enum MaybeShardingPartialDecoder { + Sharding(Arc), + Other(Arc), +} + +impl MaybeShardingPartialDecoder { + fn partial_decode( + &self, + array_subsets: &[ArraySubset], + options: &CodecOptions, + ) -> Result>, CodecError> { + match self { + Self::Sharding(partial_decoder) => { + partial_decoder.partial_decode(array_subsets, options) + } + Self::Other(partial_decoder) => partial_decoder.partial_decode(array_subsets, options), + } + } +} + +type PartialDecoderHashMap = HashMap, MaybeShardingPartialDecoder>; /// A cache used for methods in the [`ArrayShardedReadableExt`] trait. pub struct ArrayShardedReadableExtCache { @@ -86,7 +109,7 @@ impl ArrayShardedReadableExtCache { &self, array: &Array, shard_indices: &[u64], - ) -> Result, ArrayError> { + ) -> Result { let mut cache = self.cache.lock().unwrap(); if let Some(partial_decoder) = cache.get(shard_indices) { Ok(partial_decoder.clone()) @@ -100,21 +123,47 @@ impl ArrayShardedReadableExtCache { storage_transformer, array.chunk_key(shard_indices), )); + + // --- Workaround for lack off trait upcasting --- let chunk_representation = array.chunk_array_representation(shard_indices)?; - let partial_decoder = array + let sharding_codec_metadata = array .codecs() .array_to_bytes_codec() - .clone() - .partial_decoder( + .create_metadata() + .expect("valid sharding metadata"); + let sharding_codec_configuration = sharding_codec_metadata + .to_configuration::() + .expect("valid sharding configuration"); + let sharding_codec = Arc::new( + ShardingCodec::new_with_configuration(&sharding_codec_configuration).expect( + "supported sharding codec configuration, already instantiated in array", + ), + ); + let partial_decoder = + MaybeShardingPartialDecoder::Sharding(Arc::new(ShardingPartialDecoder::new( input_handle, - &chunk_representation, + chunk_representation, + sharding_codec.chunk_shape.clone(), + sharding_codec.inner_codecs.clone(), + &sharding_codec.index_codecs, + sharding_codec.index_location, &CodecOptions::default(), - )?; + )?)); + // // TODO: Trait upcasting + // let partial_decoder = array + // .codecs() + // .array_to_bytes_codec() + // .clone() + // .partial_decoder( + // input_handle, + // &chunk_representation, + // &CodecOptions::default(), + // )?; cache.insert(shard_indices.to_vec(), partial_decoder.clone()); Ok(partial_decoder) } else { - let partial_decoder: Arc = - array.partial_decoder(shard_indices)?; + let partial_decoder = + MaybeShardingPartialDecoder::Other(array.partial_decoder(shard_indices)?); cache.insert(shard_indices.to_vec(), partial_decoder.clone()); Ok(partial_decoder) } @@ -310,10 +359,14 @@ impl ArrayShardedReadableExt let (shard_indices, chunk_indices) = inner_chunk_shard_index_and_chunk_index(self, cache, inner_chunk_indices)?; let partial_decoder = cache.retrieve(self, &shard_indices)?; - let partial_decoder: Arc = partial_decoder.clone(); - let partial_decoder = partial_decoder - .downcast::() - .expect("array is exclusively sharded"); + let MaybeShardingPartialDecoder::Sharding(partial_decoder) = partial_decoder else { + unreachable!("exlusively sharded") + }; + // TODO: trait upcasting + // let partial_decoder: Arc = partial_decoder.clone(); + // let partial_decoder = partial_decoder + // .downcast::() + // .expect("array is exclusively sharded"); Ok(partial_decoder.inner_chunk_byte_range(&chunk_indices)?) } else { @@ -332,10 +385,14 @@ impl ArrayShardedReadableExt let (shard_indices, chunk_indices) = inner_chunk_shard_index_and_chunk_index(self, cache, inner_chunk_indices)?; let partial_decoder = cache.retrieve(self, &shard_indices)?; - let partial_decoder: Arc = partial_decoder.clone(); - let partial_decoder = partial_decoder - .downcast::() - .expect("array is exclusively sharded"); + let MaybeShardingPartialDecoder::Sharding(partial_decoder) = partial_decoder else { + unreachable!("exlusively sharded") + }; + // TODO: trait upcasting + // let partial_decoder: Arc = partial_decoder.clone(); + // let partial_decoder = partial_decoder + // .downcast::() + // .expect("array is exclusively sharded"); Ok(partial_decoder .retrieve_inner_chunk_encoded(&chunk_indices)? 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 3411ca40..f6ec782f 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 @@ -38,13 +38,13 @@ use unsafe_cell_slice::UnsafeCellSlice; #[derive(Clone, Debug)] pub struct ShardingCodec { /// An array of integers specifying the shape of the inner chunks in a shard along each dimension of the outer array. - chunk_shape: ChunkShape, + pub(crate) chunk_shape: ChunkShape, /// The codecs used to encode and decode inner chunks. - inner_codecs: Arc, + pub(crate) inner_codecs: Arc, /// The codecs used to encode and decode the shard index. - index_codecs: Arc, + pub(crate) index_codecs: Arc, /// Specifies whether the shard index is located at the beginning or end of the file. - index_location: ShardingIndexLocation, + pub(crate) index_location: ShardingIndexLocation, } impl ShardingCodec { diff --git a/zarrs/src/lib.rs b/zarrs/src/lib.rs index 359cfe5f..5f9b9482 100644 --- a/zarrs/src/lib.rs +++ b/zarrs/src/lib.rs @@ -184,8 +184,6 @@ //! Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in the work by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions. #![cfg_attr(docsrs, feature(doc_auto_cfg))] -#![feature(trait_upcasting)] // FIXME: Remove - pub mod array; pub mod array_subset; pub mod config;