harmos_storage/query/
table.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Immutable Paged Query Tables For Derived Evidence
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16/* === Imports === */
17
18// Standard Library Imports
19use std::sync::Arc;
20
21// Third Party Imports
22use arrow_array::RecordBatch;
23use arrow_schema::Schema;
24
25// Local Imports
26use super::spill::{Builder, Pages, Spool};
27use super::transform::{Transform, TransformExecution, invalid};
28use crate::Result;
29
30/* === Constants === */
31
32/* === Macros === */
33
34/* === Type Aliases === */
35
36/* === Structs === */
37
38/// Immutable Derived Rows, Retained As Bounded Arrow Pages Or Private Spill Files.
39#[derive(Clone)]
40pub struct Table {
41    spool: Spool,
42}
43
44/// Appends Producer Pages Without Retaining Earlier Pages.
45pub struct TableWriter {
46    builder: Builder,
47    schema: Arc<Schema>,
48}
49
50/// Retains Shared Arrow Pages Or The Spill File Until The Reader Is Dropped.
51pub struct TableReader {
52    pages: Pages,
53}
54
55/* === Enums === */
56
57/* === Traits === */
58
59/* === Implementations === */
60
61impl TableWriter {
62    pub fn new(schema: Arc<Schema>) -> Result<Self> {
63        crate::storage::arrow::validate_schema(&schema)?;
64        let builder = Builder::new(schema.clone())?;
65        Ok(Self { builder, schema })
66    }
67
68    pub fn push(&mut self, page: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<()> {
69        check()?;
70        if page.schema() != self.schema {
71            return Err(invalid("derived table schema changed"));
72        }
73        self.builder.push(page, check)
74    }
75
76    pub fn finish(self) -> Result<Table> {
77        Ok(Table { spool: self.builder.finish()? })
78    }
79}
80
81impl Table {
82    pub fn len(&self) -> usize {
83        self.spool.rows
84    }
85
86    pub fn is_empty(&self) -> bool {
87        self.len() == 0
88    }
89
90    pub fn schema(&self) -> Arc<Schema> {
91        self.spool.schema.clone()
92    }
93
94    /// Apply One Transform To The Complete Table With Bounded Spill Workspace.
95    pub fn transform(&self, transform: &Transform, check: &mut impl FnMut() -> Result<()>) -> Result<Self> {
96        check()?;
97        let schema = transform.output_schema(&self.spool.schema)?;
98        crate::storage::arrow::validate_schema(&schema)?;
99        let spool = super::reduction::transform(self.spool.clone(), transform, check)?;
100        Ok(Self { spool })
101    }
102
103    pub fn reader(&self) -> Result<TableReader> {
104        Ok(TableReader { pages: self.spool.clone().pages()? })
105    }
106}
107
108impl TableReader {
109    pub fn next(&mut self, check: &mut impl FnMut() -> Result<()>) -> Result<Option<RecordBatch>> {
110        check()?;
111        self.pages.next()
112    }
113}
114
115/* === Public Functions === */
116
117/* === Private Functions === */
118
119/* === Tests === */
120
121#[cfg(test)]
122mod tests {
123    use super::super::transform::{
124        FieldPath,
125        downsample::{Algorithm, Downsample},
126    };
127    use super::*;
128    use arrow_array::Float64Array;
129    use arrow_schema::{DataType, Field};
130
131    #[test]
132    fn derived_table_reduces_exactly_across_pages_and_survives_owner_drop() {
133        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Float64, false), Field::new("y", DataType::Float64, false)]));
134        let mut writer = TableWriter::new(schema.clone()).unwrap();
135        let x: Vec<f64> = (0..10_080).map(f64::from).collect();
136        let y: Vec<f64> = x.iter().map(|x| (x * 0.13).sin()).collect();
137        for (x, y) in x.chunks(120).zip(y.chunks(120)) {
138            let page = RecordBatch::try_new(
139                schema.clone(),
140                vec![Arc::new(Float64Array::from(x.to_vec())), Arc::new(Float64Array::from(y.to_vec()))],
141            )
142            .unwrap();
143            writer.push(&page, &mut || Ok(())).unwrap();
144        }
145        let table = writer.finish().unwrap();
146        assert_eq!(table.len(), x.len());
147        for (algorithm, expected) in [
148            (Algorithm::Lttb, crate::signal::lttb(&x, &y, 31)),
149            (Algorithm::Extrema, crate::signal::extrema(&y, 15)),
150            (Algorithm::Decimate, crate::signal::decimate(x.len(), 31)),
151        ] {
152            let transform = Transform::Downsample(Downsample {
153                max_rows: 31,
154                algorithm,
155                value: Some(FieldPath { fields: vec!["y".into()] }),
156                axis: Some(FieldPath { fields: vec!["x".into()] }),
157            });
158            let reduced = table.transform(&transform, &mut || Ok(())).unwrap();
159            let mut reader = reduced.reader().unwrap();
160            drop(reduced);
161            let mut actual = Vec::new();
162            while let Some(page) = reader.next(&mut || Ok(())).unwrap() {
163                let values = page.column(0).as_any().downcast_ref::<Float64Array>().unwrap();
164                actual.extend(values.values().iter().map(|value| *value as usize));
165            }
166            assert_eq!(actual, expected);
167        }
168        let mut reader = table.reader().unwrap();
169        drop(table);
170        let mut rows = 0;
171        while let Some(page) = reader.next(&mut || Ok(())).unwrap() {
172            assert!(page.num_rows() <= super::super::spill::PAGE_ROWS);
173            rows += page.num_rows();
174        }
175        assert_eq!(rows, x.len());
176    }
177
178    #[test]
179    fn derived_table_rejects_changed_schema_and_polls_cancellation() {
180        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Float64, false)]));
181        let mut writer = TableWriter::new(schema.clone()).unwrap();
182        let wrong = RecordBatch::new_empty(Arc::new(Schema::empty()));
183        assert!(writer.push(&wrong, &mut || Ok(())).is_err());
184        let page = RecordBatch::new_empty(schema);
185        assert!(writer.push(&page, &mut || Err(invalid("cancelled"))).is_err());
186        let table = writer.finish().unwrap();
187        assert!(table.is_empty());
188        let mut reader = table.reader().unwrap();
189        assert!(reader.next(&mut || Err(invalid("cancelled"))).is_err());
190        assert!(reader.next(&mut || Ok(())).unwrap().is_none());
191    }
192}
harmos_storage/query/table.rs

harmos_storage/query/
table.rs
