somatize_data/
parquet.rs

1//! A parquet file in a store, read by spans.
2
3use crate::{Frame, Span};
4use arrow_schema::ArrowError;
5use arrow_select::concat::concat_batches;
6use bytes::Bytes;
7use somatize_core::{Ctx, Node, NodeError, Value};
8use somatize_store::{Digest, Store};
9use std::fmt;
10use std::sync::{Arc, OnceLock};
11// The crate and this module have the same name, which is right — one file per
12// type, and the type is `Parquet` — so every path to the crate says so.
13use ::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
14
15/// A parquet file kept in a store, answering spans of rows.
16///
17/// ```ignore
18/// let sms = Parquet::at(store, "data/sms")?;
19/// memory.identify("sms", "Parquet");
20/// memory.freeze("sms", Some(sms.version().to_string()));   // and it cost no bytes
21/// ```
22///
23/// Declaring it resolves the name and stops there: a graph that names a dataset
24/// has not opened it. And it keeps the **digest**, not the name, so a dataset
25/// rebound mid-run does not change what that run is reading — resolving once is
26/// what makes the version true.
27pub struct Parquet {
28    store: Arc<dyn Store>,
29    name: String,
30    version: Digest,
31    file: OnceLock<Bytes>,
32}
33
34impl Parquet {
35    /// The parquet file bound under this name, or why there is none.
36    ///
37    /// One `resolve` and no bytes: what this costs is what makes stating a
38    /// version affordable at all.
39    pub fn at(store: Arc<dyn Store>, name: impl Into<String>) -> Result<Self, ParquetError> {
40        let name = name.into();
41        let bound = store
42            .resolve(&name)
43            .map_err(|e| ParquetError(format!("`{name}` could not be looked up: {e}")))?
44            .ok_or_else(|| ParquetError(format!("nothing is bound to `{name}` in this store")))?;
45        Ok(Self {
46            store,
47            name,
48            version: bound.digest,
49            file: OnceLock::new(),
50        })
51    }
52
53    /// What this dataset is, for the key of everything computed from it: the
54    /// digest of the content, which the store had already worked out. This is
55    /// what lets a source be settled without reading itself.
56    pub fn version(&self) -> &str {
57        self.version.as_str()
58    }
59
60    /// The name it was declared under, which is the graph's word and not the
61    /// data's.
62    pub fn name(&self) -> &str {
63        &self.name
64    }
65
66    /// The rows that span names. Short is not an error: the last span is
67    /// whatever is left, and one past the end is a frame with no rows.
68    pub fn read(&self, span: Span) -> Result<Frame, ParquetError> {
69        let file = self.file()?;
70        let builder = ParquetRecordBatchReaderBuilder::try_new(file.clone())
71            .map_err(|e| self.unreadable(e))?
72            .with_offset(span.at as usize)
73            .with_limit(span.take as usize)
74            .with_batch_size(span.take.max(1) as usize);
75        let schema = builder.schema().clone();
76        let read: Vec<_> = builder
77            .build()
78            .map_err(|e| self.unreadable(e))?
79            .collect::<Result<_, ArrowError>>()
80            .map_err(|e| self.unreadable(e))?;
81        // Concatenated rather than handed over one at a time: a span that
82        // crosses a row group comes back as two batches, and whoever asked for
83        // rows 4096..8192 asked for one frame.
84        let whole = concat_batches(&schema, read.iter()).map_err(|e| self.unreadable(e))?;
85        Ok(Frame::new(whole))
86    }
87
88    /// The bytes, fetched once.
89    fn file(&self) -> Result<&Bytes, ParquetError> {
90        if let Some(held) = self.file.get() {
91            return Ok(held);
92        }
93        let raw = self
94            .store
95            .get(&self.version)
96            .map_err(|e| ParquetError(format!("`{}` could not be read: {e}", self.name)))?
97            .ok_or_else(|| {
98                ParquetError(format!(
99                    "`{}` names {} and there are no such bytes here: a store can be \
100                     swept, and what it says it has is a record, not a promise",
101                    self.name,
102                    self.version.as_str()
103                ))
104            })?;
105        let _ = self.file.set(Bytes::from(raw));
106        Ok(self.file.get().expect("just set"))
107    }
108
109    /// The same complaint however it arrives, with the name in front of it.
110    fn unreadable(&self, why: impl fmt::Display) -> ParquetError {
111        ParquetError(format!("`{}` is not readable as parquet: {why}", self.name))
112    }
113}
114
115impl Node for Parquet {
116    /// A span in, a [`Frame`] out.
117    fn forward(&self, input: &Value, _ctx: &Ctx<'_>) -> Result<Value, NodeError> {
118        let span = Span::of(input).map_err(|e| NodeError::new(e.message()))?;
119        Ok(self
120            .read(span)
121            .map_err(|e| NodeError::new(e.message()))?
122            .value())
123    }
124}
125
126impl fmt::Debug for Parquet {
127    /// Without the bytes, and without the store: what identifies it is the name
128    /// it was declared under and the content it settled on.
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        f.debug_struct("Parquet")
131            .field("name", &self.name)
132            .field("version", &self.version.as_str())
133            .finish()
134    }
135}
136
137/// Why those rows could not be had.
138#[derive(Debug, Clone, PartialEq, Eq)]
139pub struct ParquetError(String);
140
141impl ParquetError {
142    /// The message.
143    pub fn message(&self) -> &str {
144        &self.0
145    }
146}
147
148impl fmt::Display for ParquetError {
149    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150        f.write_str(&self.0)
151    }
152}
153
154impl std::error::Error for ParquetError {}
somatize_data/parquet.rs

somatize_data/
parquet.rs
