datui_lib/formats/
row_index.rs

1//! Lazy frames decoded over a row index, for readers whose rows sit at places worked
2//! out from the row number: audio frames, fixed records, and the like.
3//!
4//! The plan has no scan in it. It is a row index over a frame with no columns, only a
5//! height, and one elementwise expression per column that decodes that column's values
6//! for the rows the index names, from an `Arc` of the source. So a slice anywhere
7//! decodes only its own rows, a query that names one column decodes only that column,
8//! and the streaming engine decodes a morsel at a time: filters, sorts and group-bys
9//! stream. The in-memory engine still builds the index whole up to a slice's end, 4
10//! bytes a row, which is why an untouched view reads its window straight from the
11//! source through [`crate::formats::pushdown::Windowed`].
12//!
13//! A Polars `AnonymousScan` gives none of this: Polars hands it no row offset, and the
14//! streaming engine of Polars 0.55 cannot run one.
15
16use polars::prelude::*;
17use std::sync::Arc;
18
19/// The row index the plan decodes from; the select over it leaves it out.
20pub const INDEX: &str = "__datui_row";
21
22/// The most rows a source can show: Polars counts rows in 32 bits.
23pub const MAX_ROWS: usize = IdxSize::MAX as usize;
24
25/// A table whose values are decoded from their row numbers.
26pub trait RowSource: Send + Sync + 'static {
27    /// Rows on hand, at most [`MAX_ROWS`].
28    fn height(&self) -> usize;
29
30    fn schema(&self) -> SchemaRef;
31
32    /// The schema's `column`th column for the rows `index` names, in that order. The
33    /// index comes from the plan's row index, so it has no nulls and every row is
34    /// under [`Self::height`]; a source still refuses one that is not, rather than
35    /// read past its bytes.
36    fn decode(&self, column: usize, index: &IdxCa) -> PolarsResult<Column>;
37}
38
39/// `source` as a lazy frame that Polars can stream, slice and prune. See the module.
40pub fn lazy<S: RowSource>(source: &Arc<S>) -> LazyFrame {
41    lazy_with_height(source, source.height())
42}
43
44/// [`lazy`] with `height` rows, for a source that grows: the frame of no columns at
45/// its root is replaced with a taller one as rows arrive (`crate::formats::lines::bound`).
46pub fn lazy_with_height<S: RowSource>(source: &Arc<S>, height: usize) -> LazyFrame {
47    frame(source, height, false)
48}
49
50/// [`lazy_with_height`], each row carrying its place in the source in [`INDEX`]: the
51/// number `#` shows, kept through a sort or a filter. Hidden from the view as the
52/// dataset's row index always is.
53pub fn lazy_numbered<S: RowSource>(source: &Arc<S>, height: usize) -> LazyFrame {
54    frame(source, height, true)
55}
56
57fn frame<S: RowSource>(source: &Arc<S>, height: usize, numbered: bool) -> LazyFrame {
58    let height = DataFrame::empty_with_height(height.min(MAX_ROWS)).lazy();
59    frame_over(source, height, numbered)
60}
61
62/// [`lazy_numbered`] over `height`, a frame of no columns whose height is the rows:
63/// for a source whose height is known only when the frame runs.
64pub fn lazy_numbered_over<S: RowSource>(source: &Arc<S>, height: LazyFrame) -> LazyFrame {
65    frame_over(source, height, true)
66}
67
68fn frame_over<S: RowSource>(source: &Arc<S>, height: LazyFrame, numbered: bool) -> LazyFrame {
69    let base = height.with_row_index(INDEX, None);
70    let mut exprs: Vec<Expr> = source
71        .schema()
72        .iter()
73        .enumerate()
74        .map(|(column, (name, dtype))| {
75            let source = Arc::clone(source);
76            let field = Field::new(name.clone(), dtype.clone());
77            col(INDEX)
78                .map(
79                    move |c: Column| source.decode(column, c.as_materialized_series().idx()?),
80                    move |_, _| Ok(field.clone()),
81                )
82                .alias(name.clone())
83        })
84        .collect();
85    if numbered {
86        exprs.push(col(INDEX));
87    }
88    base.select(exprs)
89}
90
91/// The rows `index` names as row numbers, or an error for a null or for a row at or
92/// past `rows`.
93pub fn checked(index: &IdxCa, rows: usize) -> PolarsResult<std::borrow::Cow<'_, [IdxSize]>> {
94    polars_ensure!(
95        index.null_count() == 0,
96        ComputeError: "a row index has a missing row"
97    );
98    if let Some(max) = index.max() {
99        polars_ensure!(
100            (max as usize) < rows,
101            OutOfBounds: "row {max} is past the {rows} rows on hand"
102        );
103    }
104    Ok(match index.cont_slice() {
105        Ok(rows) => std::borrow::Cow::Borrowed(rows),
106        Err(_) => std::borrow::Cow::Owned(index.into_no_null_iter().collect()),
107    })
108}
109
110#[cfg(test)]
111mod tests;
datui_lib/formats/row_index.rs

datui_lib/formats/
row_index.rs
