Skip to main content

rich_data/
window.rs

1//! Windows of rows for `rich_ext`'s virtualised tables (#260).
2//!
3//! [`VirtualTable`] fetches rows by index
4//! through [`VirtualRows`]. [`Rows`] implements it directly. A [`RowSource`]
5//! reads forward only, so [`RowWindow`] reads one pass of it and keeps just
6//! the window, the first rows (to sample column widths from) and the count:
7//! a file of any length costs one window of memory.
8//!
9//! ```
10//! use rich::Console;
11//! use rich_data::window::RowWindow;
12//! use rich_data::{Rows, Value};
13//!
14//! let mut rows = Rows::new(["n", "square"]);
15//! for n in 0..100_000i64 {
16//!     rows.push([Value::Int(n), Value::Int(n * n)]);
17//! }
18//! let window = RowWindow::reader()
19//!     .offset(50_000)
20//!     .len(2)
21//!     .read(rows.into_source())
22//!     .unwrap();
23//! assert_eq!((window.start(), window.rows().len(), window.total()), (50_000, 2, Some(100_000)));
24//!
25//! let out = Console::builder().width(40).build().render_export(&window.to_virtual_table());
26//! assert!(out.contains("│ 50001 │ 2500100001 │"), "{out}");
27//! assert!(out.ends_with("rows 50,001–50,002 of 100,000\n"), "{out}");
28//! ```
29
30use std::collections::VecDeque;
31
32use rich_ext::table::virtualized::{DEFAULT_SAMPLE, MAX_HEIGHT, MAX_SAMPLE};
33use rich_ext::table::{VirtualRows, VirtualTable};
34
35use crate::sql::typed_columns;
36use crate::{DataError, Row, RowSource, Rows, Schema, Value};
37
38impl VirtualRows for Rows {
39    fn row_count(&self) -> Option<usize> {
40        Some(self.len())
41    }
42
43    fn row(&self, index: usize) -> Option<Vec<Value>> {
44        self.rows().get(index).cloned()
45    }
46
47    fn rows(&self, start: usize, len: usize) -> Vec<Vec<Value>> {
48        let rows = Rows::rows(self);
49        let start = start.min(rows.len());
50        let end = start.saturating_add(len).min(rows.len());
51        rows[start..end].to_vec()
52    }
53}
54
55/// How [`RowWindow`] reads a source; from [`RowWindow::reader`].
56#[derive(Clone, Debug)]
57pub struct WindowReader {
58    offset: usize,
59    len: usize,
60    sample: usize,
61    count: bool,
62}
63
64impl Default for WindowReader {
65    fn default() -> Self {
66        WindowReader {
67            offset: 0,
68            len: 20,
69            sample: DEFAULT_SAMPLE,
70            count: true,
71        }
72    }
73}
74
75impl WindowReader {
76    /// Start the window at row `offset` (0-based, default 0).
77    #[must_use]
78    pub fn offset(mut self, offset: usize) -> Self {
79        self.offset = offset;
80        self
81    }
82
83    /// Keep `rows` rows (default 20, at most `MAX_HEIGHT`).
84    #[must_use]
85    pub fn len(mut self, rows: usize) -> Self {
86        self.len = rows.min(MAX_HEIGHT);
87        self
88    }
89
90    /// Also keep the first `rows` rows, for column widths and alignment
91    /// (default `DEFAULT_SAMPLE`, at most `MAX_SAMPLE`).
92    #[must_use]
93    pub fn sample(mut self, rows: usize) -> Self {
94        self.sample = rows.min(MAX_SAMPLE);
95        self
96    }
97
98    /// Read to the end to count the rows (default on). Off, reading stops one
99    /// row past the window, so the position says `of 1,041+`.
100    #[must_use]
101    pub fn count(mut self, count: bool) -> Self {
102        self.count = count;
103        self
104    }
105
106    /// Read `source` once, keeping the window. An offset past the end gives
107    /// the last full window, as a [`VirtualTable`] does. A row that cannot be
108    /// read is an error, wherever it is.
109    pub fn read(&self, mut source: impl RowSource) -> Result<RowWindow, DataError> {
110        let columns = source.columns().to_vec();
111        let schema = source.schema().cloned();
112        let width = columns.len();
113        let end = self.offset.saturating_add(self.len);
114        let mut head = Vec::new();
115        let mut ring: VecDeque<Row> = VecDeque::with_capacity(self.len);
116        let mut next = None;
117        let mut index = 0usize;
118        while let Some(row) = source.next_row() {
119            let mut row = row?;
120            row.resize(width, Value::Null);
121            if index < self.sample {
122                head.push(row.clone());
123            }
124            if index < end {
125                if self.len > 0 {
126                    if ring.len() == self.len {
127                        ring.pop_front();
128                    }
129                    ring.push_back(row);
130                }
131            } else if !self.count {
132                next = Some(row);
133                break;
134            }
135            index += 1;
136        }
137        let start = index.min(end) - ring.len();
138        Ok(RowWindow {
139            columns,
140            schema,
141            head,
142            start,
143            rows: ring.into(),
144            total: next.is_none().then_some(index),
145            next,
146        })
147    }
148}
149
150/// One window of a [`RowSource`], with its first rows and its count; see
151/// the [module docs](self).
152#[derive(Clone, Debug)]
153pub struct RowWindow {
154    columns: Vec<String>,
155    schema: Option<Schema>,
156    head: Vec<Row>,
157    start: usize,
158    rows: Vec<Row>,
159    total: Option<usize>,
160    next: Option<Row>,
161}
162
163impl RowWindow {
164    /// How to read: offset, length, sample and counting.
165    pub fn reader() -> WindowReader {
166        WindowReader::default()
167    }
168
169    /// The rows `offset..offset + len` of `source`, counting the rest.
170    pub fn read(source: impl RowSource, offset: usize, len: usize) -> Result<Self, DataError> {
171        Self::reader().offset(offset).len(len).read(source)
172    }
173
174    /// The column names.
175    pub fn columns(&self) -> &[String] {
176        &self.columns
177    }
178
179    /// The source's schema, if it had one.
180    pub fn schema(&self) -> Option<&Schema> {
181        self.schema.as_ref()
182    }
183
184    /// The 0-based index of the window's first row.
185    pub fn start(&self) -> usize {
186        self.start
187    }
188
189    /// The window's rows.
190    pub fn rows(&self) -> &[Row] {
191        &self.rows
192    }
193
194    /// How many rows the source had, when it was read to the end.
195    pub fn total(&self) -> Option<usize> {
196        self.total
197    }
198
199    /// The window as a virtualised table: typed alignment from the schema
200    /// (or the sampled cells), positioned on the window.
201    pub fn to_virtual_table(&self) -> VirtualTable<&RowWindow> {
202        let sample = if self.head.is_empty() {
203            &self.rows
204        } else {
205            &self.head
206        };
207        VirtualTable::new(
208            typed_columns(&self.columns, self.schema.as_ref(), sample),
209            self,
210        )
211        .offset(self.start)
212        .height(self.rows.len())
213        .fit_window(true)
214    }
215}
216
217impl VirtualRows for RowWindow {
218    fn row_count(&self) -> Option<usize> {
219        self.total
220    }
221
222    fn row(&self, index: usize) -> Option<Vec<Value>> {
223        if let Some(offset) = index.checked_sub(self.start) {
224            if let Some(row) = self.rows.get(offset) {
225                return Some(row.clone());
226            }
227            if offset == self.rows.len() {
228                if let Some(next) = &self.next {
229                    return Some(next.clone());
230                }
231            }
232        }
233        self.head.get(index).cloned()
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    fn numbers(count: i64) -> Rows {
242        let mut rows = Rows::new(["n"]);
243        for n in 0..count {
244            rows.push([Value::Int(n)]);
245        }
246        rows
247    }
248
249    #[test]
250    fn a_window_past_the_end_is_the_last_full_window() {
251        let window = RowWindow::read(numbers(10).into_source(), 50, 3).unwrap();
252        assert_eq!(window.start(), 7);
253        assert_eq!(window.rows()[0], [Value::Int(7)]);
254        assert_eq!(window.total(), Some(10));
255        let short = RowWindow::read(numbers(2).into_source(), 0, 5).unwrap();
256        assert_eq!((short.start(), short.rows().len()), (0, 2));
257    }
258
259    #[test]
260    fn without_counting_one_row_past_the_window_is_kept() {
261        let window = RowWindow::reader()
262            .offset(4)
263            .len(2)
264            .count(false)
265            .read(numbers(1_000).into_source())
266            .unwrap();
267        assert_eq!(window.total(), None);
268        assert_eq!(window.row(6), Some(vec![Value::Int(6)]));
269        assert_eq!(window.row(7), None);
270        let table = window.to_virtual_table();
271        assert_eq!(table.page().position(false), "rows 5–6 of 7+");
272    }
273
274    #[test]
275    fn the_head_and_the_window_answer_by_index() {
276        let window = RowWindow::reader()
277            .offset(500)
278            .len(2)
279            .sample(3)
280            .read(numbers(1_000).into_source())
281            .unwrap();
282        assert_eq!(window.row(2), Some(vec![Value::Int(2)]));
283        assert_eq!(window.row(3), None);
284        assert_eq!(window.row(501), Some(vec![Value::Int(501)]));
285        assert_eq!(VirtualRows::rows(&window, 0, 10).len(), 3);
286    }
287
288    #[test]
289    fn read_errors_surface() {
290        struct Broken(usize);
291        impl RowSource for Broken {
292            fn columns(&self) -> &[String] {
293                &[]
294            }
295            fn next_row(&mut self) -> Option<Result<Row, DataError>> {
296                self.0 += 1;
297                Some(if self.0 < 3 {
298                    Ok(Vec::new())
299                } else {
300                    Err(DataError::at(self.0, "bad row"))
301                })
302            }
303        }
304        let error = RowWindow::read(Broken(0), 0, 1).unwrap_err();
305        assert_eq!(error.to_string(), "line 3: bad row");
306    }
307}