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 parquet::file::metadata::ParquetMetaData;
18
19use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE;
20
21/// Approximate decoded bytes to produce in one batch.
22pub(crate) const TARGET_BATCH_BYTES: usize = 64 * 1024 * 1024;
23
24/// Returns `(num_rows, uncompressed_bytes)` of the largest row group in a Parquet file, which
25/// bounds the decoded size of a batch read from it.
26pub(crate) fn largest_row_group_statistic(metadata: &ParquetMetaData) -> Option<(u64, u64)> {
27    metadata
28        .row_groups()
29        .iter()
30        .map(|row_group| {
31            let uncompressed_bytes = row_group
32                .columns()
33                .iter()
34                .map(|column| column.uncompressed_size() as u64)
35                .sum();
36            (row_group.num_rows() as u64, uncompressed_bytes)
37        })
38        .max_by_key(|(_, uncompressed_bytes)| *uncompressed_bytes)
39}
40
41/// Estimates a batch size from `(num_rows, estimated_bytes)` pairs.
42///
43/// Uses the widest source with valid statistics and targets approximately 64 MiB of decoded data.
44/// The result is clamped to `1..=DEFAULT_READ_BATCH_SIZE`. Returns the default read batch size if
45/// no source contains usable statistics.
46pub fn estimate_batch_size(sources: impl IntoIterator<Item = (u64, u64)>) -> usize {
47    let max_row_width = sources
48        .into_iter()
49        .filter_map(|(rows, bytes)| estimate_row_width(rows, bytes))
50        .max();
51
52    let Some(row_width) = max_row_width else {
53        return DEFAULT_READ_BATCH_SIZE;
54    };
55
56    (TARGET_BATCH_BYTES as u64 / row_width).clamp(1, DEFAULT_READ_BATCH_SIZE as u64) as usize
57}
58
59/// Returns the ceiling of the estimated bytes per row.
60fn estimate_row_width(rows: u64, bytes: u64) -> Option<u64> {
61    if rows == 0 || bytes == 0 {
62        return None;
63    }
64
65    Some(bytes / rows + u64::from(!bytes.is_multiple_of(rows)))
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn test_estimate_batch_size_without_stats() {
74        assert_eq!(DEFAULT_READ_BATCH_SIZE, estimate_batch_size([]));
75        assert_eq!(
76            DEFAULT_READ_BATCH_SIZE,
77            estimate_batch_size([(0, 100), (100, 0)])
78        );
79    }
80
81    #[test]
82    fn test_estimate_batch_size_for_narrow_and_wide_rows() {
83        assert_eq!(DEFAULT_READ_BATCH_SIZE, estimate_batch_size([(100, 100)]));
84        assert_eq!(256, estimate_batch_size([(1, 256 * 1024)]));
85        assert_eq!(1, estimate_batch_size([(1, TARGET_BATCH_BYTES as u64 + 1)]));
86    }
87
88    #[test]
89    fn test_estimate_batch_size_uses_widest_known_source() {
90        assert_eq!(
91            256,
92            estimate_batch_size([(100, 0), (100, 100), (4, 1024 * 1024)])
93        );
94    }
95
96    #[test]
97    fn test_estimate_batch_size_saturates() {
98        assert_eq!(
99            DEFAULT_READ_BATCH_SIZE,
100            estimate_batch_size([(u64::MAX, u64::MAX)])
101        );
102        assert_eq!(1, estimate_batch_size([(1, u64::MAX)]));
103        assert_eq!(1, estimate_batch_size([(2, u64::MAX)]));
104    }
105}
mito2/batch_size.rs

mito2/
batch_size.rs
