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    pub fn offset(mut self, offset: usize) -> Self {
78        self.offset = offset;
79        self
80    }
81
82    /// Keep `rows` rows (default 20, at most `MAX_HEIGHT`).
83    pub fn len(mut self, rows: usize) -> Self {
84        self.len = rows.min(MAX_HEIGHT);
85        self
86    }
87
88    /// Also keep the first `rows` rows, for column widths and alignment
89    /// (default `DEFAULT_SAMPLE`, at most `MAX_SAMPLE`).
90    pub fn sample(mut self, rows: usize) -> Self {
91        self.sample = rows.min(MAX_SAMPLE);
92        self
93    }
94
95    /// Read to the end to count the rows (default on). Off, reading stops one
96    /// row past the window, so the position says `of 1,041+`.
97    pub fn count(mut self, count: bool) -> Self {
98        self.count = count;
99        self
100    }
101
102    /// Read `source` once, keeping the window. An offset past the end gives
103    /// the last full window, as a [`VirtualTable`] does. A row that cannot be
104    /// read is an error, wherever it is.
105    pub fn read(&self, mut source: impl RowSource) -> Result<RowWindow, DataError> {
106        let columns = source.columns().to_vec();
107        let schema = source.schema().cloned();
108        let width = columns.len();
109        let end = self.offset.saturating_add(self.len);
110        let mut head = Vec::new();
111        let mut ring: VecDeque<Row> = VecDeque::with_capacity(self.len);
112        let mut next = None;
113        let mut index = 0usize;
114        while let Some(row) = source.next_row() {
115            let mut row = row?;
116            row.resize(width, Value::Null);
117            if index < self.sample {
118                head.push(row.clone());
119            }
120            if index < end {
121                if self.len > 0 {
122                    if ring.len() == self.len {
123                        ring.pop_front();
124                    }
125                    ring.push_back(row);
126                }
127            } else if !self.count {
128                next = Some(row);
129                break;
130            }
131            index += 1;
132        }
133        let start = index.min(end) - ring.len();
134        Ok(RowWindow {
135            columns,
136            schema,
137            head,
138            start,
139            rows: ring.into(),
140            total: next.is_none().then_some(index),
141            next,
142        })
143    }
144}
145
146/// One window of a [`RowSource`], with its first rows and its count; see
147/// the [module docs](self).
148#[derive(Clone, Debug)]
149pub struct RowWindow {
150    columns: Vec<String>,
151    schema: Option<Schema>,
152    head: Vec<Row>,
153    start: usize,
154    rows: Vec<Row>,
155    total: Option<usize>,
156    next: Option<Row>,
157}
158
159impl RowWindow {
160    /// How to read: offset, length, sample and counting.
161    pub fn reader() -> WindowReader {
162        WindowReader::default()
163    }
164
165    /// The rows `offset..offset + len` of `source`, counting the rest.
166    pub fn read(source: impl RowSource, offset: usize, len: usize) -> Result<Self, DataError> {
167        Self::reader().offset(offset).len(len).read(source)
168    }
169
170    /// The column names.
171    pub fn columns(&self) -> &[String] {
172        &self.columns
173    }
174
175    /// The source's schema, if it had one.
176    pub fn schema(&self) -> Option<&Schema> {
177        self.schema.as_ref()
178    }
179
180    /// The 0-based index of the window's first row.
181    pub fn start(&self) -> usize {
182        self.start
183    }
184
185    /// The window's rows.
186    pub fn rows(&self) -> &[Row] {
187        &self.rows
188    }
189
190    /// How many rows the source had, when it was read to the end.
191    pub fn total(&self) -> Option<usize> {
192        self.total
193    }
194
195    /// The window as a virtualised table: typed alignment from the schema
196    /// (or the sampled cells), positioned on the window.
197    pub fn to_virtual_table(&self) -> VirtualTable<&RowWindow> {
198        let sample = if self.head.is_empty() {
199            &self.rows
200        } else {
201            &self.head
202        };
203        VirtualTable::new(
204            typed_columns(&self.columns, self.schema.as_ref(), sample),
205            self,
206        )
207        .offset(self.start)
208        .height(self.rows.len())
209        .fit_window(true)
210    }
211}
212
213impl VirtualRows for RowWindow {
214    fn row_count(&self) -> Option<usize> {
215        self.total
216    }
217
218    fn row(&self, index: usize) -> Option<Vec<Value>> {
219        if let Some(offset) = index.checked_sub(self.start) {
220            if let Some(row) = self.rows.get(offset) {
221                return Some(row.clone());
222            }
223            if offset == self.rows.len() {
224                if let Some(next) = &self.next {
225                    return Some(next.clone());
226                }
227            }
228        }
229        self.head.get(index).cloned()
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236
237    fn numbers(count: i64) -> Rows {
238        let mut rows = Rows::new(["n"]);
239        for n in 0..count {
240            rows.push([Value::Int(n)]);
241        }
242        rows
243    }
244
245    #[test]
246    fn a_window_past_the_end_is_the_last_full_window() {
247        let window = RowWindow::read(numbers(10).into_source(), 50, 3).unwrap();
248        assert_eq!(window.start(), 7);
249        assert_eq!(window.rows()[0], [Value::Int(7)]);
250        assert_eq!(window.total(), Some(10));
251        let short = RowWindow::read(numbers(2).into_source(), 0, 5).unwrap();
252        assert_eq!((short.start(), short.rows().len()), (0, 2));
253    }
254
255    #[test]
256    fn without_counting_one_row_past_the_window_is_kept() {
257        let window = RowWindow::reader()
258            .offset(4)
259            .len(2)
260            .count(false)
261            .read(numbers(1_000).into_source())
262            .unwrap();
263        assert_eq!(window.total(), None);
264        assert_eq!(window.row(6), Some(vec![Value::Int(6)]));
265        assert_eq!(window.row(7), None);
266        let table = window.to_virtual_table();
267        assert_eq!(table.page().position(false), "rows 5–6 of 7+");
268    }
269
270    #[test]
271    fn the_head_and_the_window_answer_by_index() {
272        let window = RowWindow::reader()
273            .offset(500)
274            .len(2)
275            .sample(3)
276            .read(numbers(1_000).into_source())
277            .unwrap();
278        assert_eq!(window.row(2), Some(vec![Value::Int(2)]));
279        assert_eq!(window.row(3), None);
280        assert_eq!(window.row(501), Some(vec![Value::Int(501)]));
281        assert_eq!(VirtualRows::rows(&window, 0, 10).len(), 3);
282    }
283
284    #[test]
285    fn read_errors_surface() {
286        struct Broken(usize);
287        impl RowSource for Broken {
288            fn columns(&self) -> &[String] {
289                &[]
290            }
291            fn next_row(&mut self) -> Option<Result<Row, DataError>> {
292                self.0 += 1;
293                Some(if self.0 < 3 {
294                    Ok(Vec::new())
295                } else {
296                    Err(DataError::at(self.0, "bad row"))
297                })
298            }
299        }
300        let error = RowWindow::read(Broken(0), 0, 1).unwrap_err();
301        assert_eq!(error.to_string(), "line 3: bad row");
302    }
303}