Skip to main content

polars_io/
configs.rs

1use std::num::NonZeroUsize;
2use std::sync::LazyLock;
3
4pub fn env_upload_chunk_size() -> Option<NonZeroUsize> {
5    std::env::var("POLARS_UPLOAD_CHUNK_SIZE").ok().map(|s| {
6        s.parse::<NonZeroUsize>()
7            .unwrap_or_else(|_| panic!("invalid value for POLARS_UPLOAD_CHUNK_SIZE: {s}"))
8    })
9}
10
11pub fn env_partitioned_upload_chunk_size() -> Option<NonZeroUsize> {
12    std::env::var("POLARS_PARTITIONED_UPLOAD_CHUNK_SIZE")
13        .ok()
14        .map(|s| {
15            s.parse::<NonZeroUsize>().unwrap_or_else(|_| {
16                panic!("invalid value for POLARS_PARTITIONED_UPLOAD_CHUNK_SIZE: {s}")
17            })
18        })
19}
20
21/// Max concurrent tasks within a single cloud writer.
22pub fn env_upload_concurrency() -> Option<NonZeroUsize> {
23    std::env::var("POLARS_UPLOAD_CONCURRENCY").ok().map(|s| {
24        s.parse::<NonZeroUsize>()
25            .unwrap_or_else(|_| panic!("invalid value for POLARS_UPLOAD_CONCURRENCY: {s}"))
26    })
27}
28
29pub fn env_partitioned_upload_concurrency() -> Option<NonZeroUsize> {
30    std::env::var("POLARS_PARTITIONED_UPLOAD_CONCURRENCY")
31        .ok()
32        .map(|s| {
33            s.parse::<NonZeroUsize>().unwrap_or_else(|_| {
34                panic!("invalid value for POLARS_PARTITIONED_UPLOAD_CONCURRENCY: {s}")
35            })
36        })
37}
38
39/// Runs of this many values whose total bytes are <= `copy_buffer_reserve_size` will be copied into
40/// a single contiguous chunk.
41pub(crate) fn cloud_writer_coalesce_run_length() -> usize {
42    return *COALESCE_RUN_LENGTH;
43
44    static COALESCE_RUN_LENGTH: LazyLock<usize> = LazyLock::new(|| {
45        let mut v: usize = 64;
46
47        if let Ok(s) = std::env::var("POLARS_CLOUD_WRITER_COALESCE_RUN_LENGTH") {
48            v = s
49                .parse::<usize>()
50                .ok()
51                .filter(|x| *x >= 2)
52                .unwrap_or_else(|| {
53                    panic!("invalid value for POLARS_CLOUD_WRITER_COALESCE_RUN_LENGTH: {s}")
54                })
55        }
56
57        if polars_core::config::verbose() {
58            eprintln!("cloud_writer coalesce_run_length: {v}")
59        }
60
61        v
62    });
63}
64
65pub(crate) fn cloud_writer_copy_buffer_size() -> NonZeroUsize {
66    return *COPY_BUFFER_SIZE;
67
68    static COPY_BUFFER_SIZE: LazyLock<NonZeroUsize> = LazyLock::new(|| {
69        let mut v: NonZeroUsize = const { NonZeroUsize::new(16 * 1024 * 1024).unwrap() };
70
71        if let Ok(s) = std::env::var("POLARS_CLOUD_WRITER_COPY_BUFFER_SIZE") {
72            v = s.parse::<NonZeroUsize>().unwrap_or_else(|_| {
73                panic!("invalid value for POLARS_CLOUD_WRITER_COPY_BUFFER_SIZE: {s}")
74            })
75        }
76
77        if polars_core::config::verbose() {
78            eprintln!("cloud_writer copy_buffer_size: {v}")
79        }
80
81        v
82    });
83}