dataprof_core/sampling/chunk_size.rs
1use sysinfo::System;
2
3/// How much data a streaming engine reads per chunk.
4///
5/// **The unit is bytes, everywhere.** Chunk size bounds how much of the source
6/// is resident at once, so it is expressed in the unit that actually bounds
7/// memory rather than in rows, whose width varies per dataset. Every engine and
8/// binding agrees: `ChunkSize::Fixed(65_536)` and Python's `chunk_size=65536`
9/// both mean 64 KiB per chunk, whether the source is a file, a byte stream, or
10/// a URL.
11///
12/// Chunk size never changes *what* a profile contains — only the granularity at
13/// which the source is read, progress is emitted, and chunk-level stop
14/// conditions are evaluated.
15#[derive(Debug, Clone, Default)]
16pub enum ChunkSize {
17 /// Fixed chunk size in **bytes**.
18 Fixed(usize),
19
20 /// Let the engine choose the chunk size (default).
21 ///
22 /// Each engine resolves this against what it knows: the incremental engine
23 /// derives a size from its memory limit and the file size, while the async
24 /// reader — whose source has no length to adapt to — uses a fixed working-set
25 /// target. Unlike [`Fixed`](Self::Fixed), the resulting size is not a
26 /// guarantee, so it is not something to assert against.
27 #[default]
28 Adaptive,
29
30 /// Custom sizing function, given the source size in bytes and returning a
31 /// chunk size in bytes. Cannot derive Debug/Clone with a function pointer.
32 Custom(fn(u64) -> usize),
33}
34
35impl ChunkSize {
36 /// Resolve to a concrete chunk size in bytes for a source of the given size.
37 ///
38 /// This is a standalone helper, not the path any engine takes: `Fixed` and
39 /// `Custom` resolve exactly as an engine would, but `Adaptive` here is
40 /// derived from *system* available memory, whereas an engine derives it
41 /// from its own configured memory limit. Do not use this to predict what an
42 /// engine will do with `Adaptive`.
43 pub fn calculate(&self, file_size_bytes: u64) -> usize {
44 match self {
45 ChunkSize::Fixed(size) => *size,
46 ChunkSize::Adaptive => self.adaptive_size(file_size_bytes),
47 ChunkSize::Custom(func) => func(file_size_bytes),
48 }
49 }
50
51 /// Derive a chunk size from system available memory and the source size.
52 fn adaptive_size(&self, file_size_bytes: u64) -> usize {
53 let mut system = System::new_all();
54 system.refresh_memory();
55
56 let available_memory = system.available_memory();
57
58 // Use max 10% of available memory for each chunk
59 let bytes_per_chunk = (available_memory / 10).max(64 * 1024 * 1024) as usize; // Min 64MB
60
61 // Adjust based on file size
62 let file_size_mb = file_size_bytes / (1024 * 1024);
63
64 if file_size_mb > 10_000 {
65 // Very large files: smaller chunks to avoid memory pressure
66 bytes_per_chunk / 2
67 } else {
68 bytes_per_chunk
69 }
70 }
71}