mito2/
batch_size.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Utilities for choosing an execution batch size from source statistics.
16
17use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
18
19/// Approximate decoded bytes to produce in one batch.
20pub(crate) const TARGET_BATCH_BYTES: usize = 64 * 1024 * 1024;
21
22/// Estimates a batch size from `(num_rows, estimated_bytes)` pairs.
23///
24/// Uses the widest source with valid statistics and targets approximately 64 MiB of decoded data.
25/// The result is clamped to `1..=DEFAULT_READ_BATCH_SIZE`. Returns the default read batch size if
26/// no source contains usable statistics.
27pub fn estimate_batch_size(sources: impl IntoIterator<Item = (u64, u64)>) -> usize {
28    let max_row_width = sources
29        .into_iter()
30        .filter_map(|(rows, bytes)| estimate_row_width(rows, bytes))
31        .max();
32
33    let Some(row_width) = max_row_width else {
34        return DEFAULT_READ_BATCH_SIZE;
35    };
36
37    (TARGET_BATCH_BYTES as u64 / row_width).clamp(1, DEFAULT_READ_BATCH_SIZE as u64) as usize
38}
39
40/// Returns the ceiling of the estimated bytes per row.
41fn estimate_row_width(rows: u64, bytes: u64) -> Option<u64> {
42    if rows == 0 || bytes == 0 {
43        return None;
44    }
45
46    Some(bytes / rows + u64::from(!bytes.is_multiple_of(rows)))
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn test_estimate_batch_size_without_stats() {
55        assert_eq!(DEFAULT_READ_BATCH_SIZE, estimate_batch_size([]));
56        assert_eq!(
57            DEFAULT_READ_BATCH_SIZE,
58            estimate_batch_size([(0, 100), (100, 0)])
59        );
60    }
61
62    #[test]
63    fn test_estimate_batch_size_for_narrow_and_wide_rows() {
64        assert_eq!(DEFAULT_READ_BATCH_SIZE, estimate_batch_size([(100, 100)]));
65        assert_eq!(256, estimate_batch_size([(1, 256 * 1024)]));
66        assert_eq!(1, estimate_batch_size([(1, TARGET_BATCH_BYTES as u64 + 1)]));
67    }
68
69    #[test]
70    fn test_estimate_batch_size_uses_widest_known_source() {
71        assert_eq!(
72            256,
73            estimate_batch_size([(100, 0), (100, 100), (4, 1024 * 1024)])
74        );
75    }
76
77    #[test]
78    fn test_estimate_batch_size_saturates() {
79        assert_eq!(
80            DEFAULT_READ_BATCH_SIZE,
81            estimate_batch_size([(u64::MAX, u64::MAX)])
82        );
83        assert_eq!(1, estimate_batch_size([(1, u64::MAX)]));
84        assert_eq!(1, estimate_batch_size([(2, u64::MAX)]));
85    }
86}
mito2/batch_size.rs

mito2/
batch_size.rs
