Skip to main content

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}