icechunk 2.0.2

Transactional storage engine for Zarr designed for use on cloud object storage
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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

#![allow(clippy::expect_used, clippy::unwrap_used)]

use std::path::PathBuf;
use std::sync::Arc;
use std::vec;

use icechunk::{
    Repository, RepositoryConfig, Storage,
    config::{
        ManifestConfig, ManifestSplitCondition, ManifestSplitDim,
        ManifestSplitDimCondition, ManifestSplittingConfig,
    },
    format::{ChunkIndices, Path, manifest::ChunkPayload, snapshot::ArrayShape},
    storage::ObjectStorage,
};

use bytes::Bytes;
use itertools::Itertools as _;
use tokio::sync::RwLock;

use clap::Parser;
use test_log::tracing_subscriber;
use test_log::tracing_subscriber::{
    layer::SubscriberExt as _, util::SubscriberInitExt as _,
};
use tracing::info;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Location for the new repo
    #[arg(default_value = "test_repo")]
    repo_dir: PathBuf,

    /// Number of `chunk_refs` to set
    #[arg(default_value_t = 1_000_000)]
    count: u32,

    /// Number of `chunk_refs` per manifest
    #[arg(default_value_t = 50_000_000)]
    num_chunks: u32,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::new(
            std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()),
        ))
        .with(tracing_subscriber::fmt::layer().compact())
        .init();

    const CHUNK_SIZE: usize = 20_000;

    let args = Args::parse();
    let repo_dir = args.repo_dir;
    let repo_size = args.count;
    let num_chunks = args.num_chunks;

    let storage: Arc<dyn Storage + Send + Sync> = Arc::new(
        ObjectStorage::new_local_filesystem(repo_dir.as_path())
            .await
            .expect("Creating local storage failed"),
    );
    let creds = Default::default();

    let split_sizes = Some(vec![(
        ManifestSplitCondition::PathMatches { regex: r".*".to_string() },
        vec![ManifestSplitDim { condition: ManifestSplitDimCondition::Any, num_chunks }],
    )]);
    let manifest_config = ManifestConfig {
        splitting: Some(ManifestSplittingConfig { split_sizes }),
        ..ManifestConfig::default()
    };

    let repo = Repository::create(
        Some(RepositoryConfig {
            inline_chunk_threshold_bytes: Some(0),
            manifest: Some(manifest_config),
            ..Default::default()
        }),
        storage,
        creds,
        None,
        true,
    )
    .await
    .expect("Failed to initialize repository");

    let ds = Arc::new(RwLock::new(repo.writable_session("main").await?));

    let shape = ArrayShape::new([(repo_size as u64, 1)]).unwrap();
    let user_data = Bytes::new();

    let new_array_path: Path = "/array".try_into().unwrap();

    ds.write()
        .await
        .add_array(new_array_path.clone(), shape, None, user_data)
        .await
        .unwrap();

    let payload = ChunkPayload::Inline(Bytes::from_static(&[42u8]));

    let mut count = 0;
    for chnk in &(0..repo_size).chunks(CHUNK_SIZE) {
        let mut set = tokio::task::JoinSet::new();

        for idx in chnk {
            let new_array_path = new_array_path.clone();
            let payload = payload.clone();
            let ds = Arc::clone(&ds);

            set.spawn(async move {
                ds.write()
                    .await
                    .set_chunk_ref(new_array_path, ChunkIndices(vec![idx]), Some(payload))
                    .await
            });
        }

        while let Some(res) = set.join_next().await {
            res.unwrap().unwrap();
        }

        count += CHUNK_SIZE;
        if count % 1_000_000 == 0 {
            info!("{count} chunk_ref written");
        }
    }

    info!("starting commit");
    ds.write()
        .await
        .commit("really big changeset")
        .max_concurrent_nodes(8)
        .execute()
        .await?;

    Ok(())
}