zarrs_icechunk 0.5.1

icechunk store support for the zarrs crate
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use std::{num::NonZeroU64, sync::Arc};

use icechunk::{repository::VersionInfo, Repository, RepositoryConfig};
use zarrs::array::{data_type, ArrayBuilder, ArraySubset};
use zarrs_icechunk::AsyncIcechunkStore;

#[tokio::test]
async fn icechunk_array() -> Result<(), Box<dyn std::error::Error>> {
    let storage = icechunk::new_in_memory_storage().await?;
    let config = RepositoryConfig::default();
    let repo = Repository::create(
        Some(config),
        storage,
        Default::default(),
        Default::default(),
        true,
    )
    .await?;

    let array_path = "/array";
    let data_type = data_type::uint8();
    let mut builder = ArrayBuilder::new(
        vec![4, 4], // array shape
        vec![2, 2], // regular chunk shape
        data_type.clone(),
        0u8,
    );
    builder.bytes_to_bytes_codecs(vec![]);
    builder.array_to_bytes_codec(Arc::new(
        zarrs::array::codec::array_to_bytes::sharding::ShardingCodecBuilder::new(
            vec![NonZeroU64::new(1).unwrap(), NonZeroU64::new(1).unwrap()],
            &data_type,
        )
        .bytes_to_bytes_codecs(vec![Arc::new(zarrs::array::codec::GzipCodec::new(5)?)])
        .build(),
    ));

    let session = repo.writable_session("main").await?;
    let store = Arc::new(AsyncIcechunkStore::new(session));
    let array = builder.build(store.clone(), array_path).unwrap();

    array.async_store_metadata().await?;

    assert_eq!(array.data_type(), &data_type::uint8());
    assert_eq!(array.fill_value().as_ne_bytes(), &[0u8]);
    assert_eq!(array.shape(), &[4, 4]);
    assert_eq!(
        array.chunk_shape(&[0, 0]).unwrap(),
        vec![NonZeroU64::new(2).unwrap(), NonZeroU64::new(2).unwrap()]
    );
    assert_eq!(array.chunk_grid_shape(), &[2, 2]);

    // 1  2 | 0  0
    // 0  0 | 0  0
    // -----|-----
    // 0  0 | 0  0
    // 0  0 | 0  0
    array.async_store_chunk(&[0, 0], &[1u8, 2, 0, 0]).await?;
    let snapshot0 = store.session().write().await.commit("a").execute().await?;

    let session = repo.writable_session("main").await?;
    let store = Arc::new(AsyncIcechunkStore::new(session));
    let array = builder.build(store.clone(), array_path).unwrap();

    // 1  2 | 3  4
    // 5  6 | 7  8
    // -----|-----
    // 9 10 | 0  0
    // 0  0 | 0  0
    array.async_store_chunk(&[0, 1], &[3u8, 4, 7, 8]).await?;
    array
        .async_store_array_subset(
            &ArraySubset::new_with_ranges(&[1..3, 0..2]),
            &[5u8, 6, 9, 10],
        )
        .await?;
    assert!(array
        .async_retrieve_chunk::<Vec<u8>>(&[0, 0, 0])
        .await
        .is_err());
    assert_eq!(
        array.async_retrieve_chunk::<Vec<u8>>(&[0, 0]).await?,
        vec![1u8, 2, 5, 6]
    );
    let _snapshot1 = store.session().write().await.commit("b").execute().await?;

    let session = repo
        .readonly_session(&VersionInfo::SnapshotId(snapshot0))
        .await?;
    let store = Arc::new(AsyncIcechunkStore::new(session));
    let array = builder.build(store.clone(), array_path).unwrap();

    assert_eq!(
        array.async_retrieve_chunk::<Vec<u8>>(&[0, 0]).await?,
        vec![1u8, 2, 0, 0]
    );

    Ok(())
}