krishiv_plan/optimizer/
small_file.rs

1//! Small-file scan planner.
2
3/// Per-file metadata used by [`SmallFilePlanner`].
4#[non_exhaustive]
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct FileStats {
7    pub path: String,
8    pub size_bytes: u64,
9}
10
11/// Advice produced by [`SmallFilePlanner`]: a list of scan groups where each
12/// group of file paths should be handled by a single executor task.
13#[non_exhaustive]
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct SplitPlanAdvice {
16    /// Each inner `Vec` is one task's worth of files.
17    pub task_groups: Vec<Vec<String>>,
18}
19
20/// Plans scan parallelism for a set of files.
21///
22/// When individual files are smaller than `target_bytes`, multiple files are
23/// grouped into a single task so each task processes roughly `target_bytes` of
24/// data. Files larger than `target_bytes` each get their own task (splitting
25/// within a file is not yet supported).
26pub struct SmallFilePlanner {
27    target_bytes: u64,
28}
29
30impl SmallFilePlanner {
31    /// Create a planner with the given target bytes per task.
32    pub fn new(target_bytes: u64) -> Self {
33        Self { target_bytes }
34    }
35
36    /// Produce a scan plan for the given file list.
37    ///
38    /// Files are grouped greedily: accumulate until the next file would push the
39    /// group over `target_bytes`, then start a new group. This ensures each
40    /// group is at most `target_bytes + max_single_file_bytes`.
41    pub fn plan(&self, files: &[FileStats]) -> SplitPlanAdvice {
42        if files.is_empty() {
43            return SplitPlanAdvice {
44                task_groups: Vec::new(),
45            };
46        }
47
48        let mut groups: Vec<Vec<String>> = Vec::new();
49        let mut current: Vec<String> = Vec::new();
50        let mut current_bytes = 0u128;
51        let target_bytes = u128::from(self.target_bytes);
52
53        for file in files {
54            let file_bytes = u128::from(file.size_bytes);
55            if !current.is_empty() && current_bytes + file_bytes > target_bytes {
56                groups.push(std::mem::take(&mut current));
57                current_bytes = 0;
58            }
59            current.push(file.path.clone());
60            current_bytes += file_bytes;
61        }
62        if !current.is_empty() {
63            groups.push(current);
64        }
65
66        SplitPlanAdvice {
67            task_groups: groups,
68        }
69    }
70}
krishiv_plan/optimizer/small_file.rs

krishiv_plan/optimizer/
small_file.rs
