polars_io/
configs.rs

1use std::num::NonZeroUsize;
2use std::sync::LazyLock;
3
4use polars_core::config::verbose;
5use polars_utils::sys::total_memory;
6
7pub fn upload_chunk_size() -> usize {
8    return *UPLOAD_CHUNK_SIZE;
9
10    static UPLOAD_CHUNK_SIZE: LazyLock<usize> = LazyLock::new(|| {
11        let mut v: usize = 32 * 1024 * 1024;
12
13        if let Ok(s) = std::env::var("POLARS_UPLOAD_CHUNK_SIZE") {
14            v = s
15                .parse::<usize>()
16                .unwrap_or_else(|_| panic!("invalid value for POLARS_UPLOAD_CHUNK_SIZE: {s}"))
17        }
18
19        if verbose() {
20            eprintln!("upload_chunk_size: {v}")
21        }
22
23        v
24    });
25}
26
27pub fn partitioned_upload_chunk_size() -> usize {
28    return *PARTITIONED_UPLOAD_CHUNK_SIZE;
29
30    static PARTITIONED_UPLOAD_CHUNK_SIZE: LazyLock<usize> = LazyLock::new(|| {
31        let mut v: usize = 6 * 1024 * 1024;
32
33        if let Ok(s) = std::env::var("POLARS_PARTITIONED_UPLOAD_CHUNK_SIZE") {
34            v = s.parse::<usize>().unwrap_or_else(|_| {
35                panic!("invalid value for POLARS_PARTITIONED_UPLOAD_CHUNK_SIZE: {s}")
36            })
37        }
38
39        if verbose() {
40            eprintln!("partitioned_upload_chunk_size: {v}")
41        }
42
43        v
44    });
45}
46
47/// Max concurrent tasks within a single cloud writer.
48pub fn upload_concurrency() -> NonZeroUsize {
49    return *UPLOAD_CONCURRENCY;
50
51    static UPLOAD_CONCURRENCY: LazyLock<NonZeroUsize> = LazyLock::new(|| {
52        let buffer_limit: usize = (total_memory() / 32) as _;
53
54        let mut v: NonZeroUsize =
55            NonZeroUsize::new(usize::clamp(buffer_limit / upload_chunk_size(), 8, 256)).unwrap();
56
57        if let Ok(s) = std::env::var("POLARS_UPLOAD_CONCURRENCY") {
58            v = s
59                .parse::<NonZeroUsize>()
60                .unwrap_or_else(|_| panic!("invalid value for POLARS_UPLOAD_CONCURRENCY: {s}"))
61        }
62
63        if verbose() {
64            eprintln!("upload_concurrency: {v}")
65        }
66
67        v
68    });
69}
70
71pub fn partitioned_upload_concurrency() -> NonZeroUsize {
72    return *PARTITIONED_UPLOAD_CONCURRENCY;
73
74    static PARTITIONED_UPLOAD_CONCURRENCY: LazyLock<NonZeroUsize> = LazyLock::new(|| {
75        let mut v: NonZeroUsize = NonZeroUsize::new(64).unwrap();
76
77        if let Ok(s) = std::env::var("POLARS_PARTITIONED_UPLOAD_CONCURRENCY") {
78            v = s.parse::<NonZeroUsize>().unwrap_or_else(|_| {
79                panic!("invalid value for POLARS_PARTITIONED_UPLOAD_CONCURRENCY: {s}")
80            })
81        }
82
83        if verbose() {
84            eprintln!("partitioned_upload_concurrency: {v}")
85        }
86
87        v
88    });
89}
90
91/// Runs of this many values whose total bytes are <= `copy_buffer_reserve_size` will be copied into
92/// a single contiguous chunk.
93pub(crate) fn cloud_writer_coalesce_run_length() -> usize {
94    return *COALESCE_RUN_LENGTH;
95
96    static COALESCE_RUN_LENGTH: LazyLock<usize> = LazyLock::new(|| {
97        let mut v: usize = 64;
98
99        if let Ok(s) = std::env::var("POLARS_CLOUD_WRITER_COALESCE_RUN_LENGTH") {
100            v = s
101                .parse::<usize>()
102                .ok()
103                .filter(|x| *x >= 2)
104                .unwrap_or_else(|| {
105                    panic!("invalid value for POLARS_CLOUD_WRITER_COALESCE_RUN_LENGTH: {s}")
106                })
107        }
108
109        if polars_core::config::verbose() {
110            eprintln!("cloud_writer coalesce_run_length: {v}")
111        }
112
113        v
114    });
115}
116
117pub(crate) fn cloud_writer_copy_buffer_size() -> NonZeroUsize {
118    return *COPY_BUFFER_SIZE;
119
120    static COPY_BUFFER_SIZE: LazyLock<NonZeroUsize> = LazyLock::new(|| {
121        let mut v: NonZeroUsize = const { NonZeroUsize::new(16 * 1024 * 1024).unwrap() };
122
123        if let Ok(s) = std::env::var("POLARS_CLOUD_WRITER_COPY_BUFFER_SIZE") {
124            v = s.parse::<NonZeroUsize>().unwrap_or_else(|_| {
125                panic!("invalid value for POLARS_CLOUD_WRITER_COPY_BUFFER_SIZE: {s}")
126            })
127        }
128
129        if polars_core::config::verbose() {
130            eprintln!("cloud_writer copy_buffer_size: {v}")
131        }
132
133        v
134    });
135}