Skip to main content

vortex_array/display/extractors/
buffer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5
6use humansize::DECIMAL;
7use humansize::format_size;
8
9use crate::ArrayRef;
10use crate::display::extractor::IndentedFormatter;
11use crate::display::extractor::TreeContext;
12use crate::display::extractor::TreeExtractor;
13
14/// Extractor that adds buffer detail lines.
15pub struct BufferExtractor {
16    /// Whether to show buffer-level percentage of parent nbytes.
17    pub show_percent: bool,
18}
19
20impl TreeExtractor for BufferExtractor {
21    fn write_details(
22        &self,
23        array: &ArrayRef,
24        _ctx: &TreeContext,
25        f: &mut IndentedFormatter<'_, '_>,
26    ) -> fmt::Result {
27        let (indent, f) = f.parts();
28        let nbytes = array.nbytes();
29        for (name, buffer) in array.named_buffers() {
30            let loc = if buffer.is_on_device() {
31                "device"
32            } else if buffer.is_on_host() {
33                "host"
34            } else {
35                "location-unknown"
36            };
37            let align = if buffer.is_on_host() {
38                buffer.as_host().alignment().to_string()
39            } else {
40                String::new()
41            };
42
43            if self.show_percent {
44                let buffer_percent = if nbytes == 0 {
45                    0.0
46                } else {
47                    100_f64 * buffer.len() as f64 / nbytes as f64
48                };
49                writeln!(
50                    f,
51                    "{indent}buffer: {} {loc} {} (align={}) ({:.2}%)",
52                    name,
53                    format_size(buffer.len(), DECIMAL),
54                    align,
55                    buffer_percent,
56                )?;
57            } else {
58                writeln!(
59                    f,
60                    "{indent}buffer: {} {loc} {} (align={})",
61                    name,
62                    format_size(buffer.len(), DECIMAL),
63                    align,
64                )?;
65            }
66        }
67        Ok(())
68    }
69}