tessera/readers/
polars.rs

1//! Load Tessera partitions into Polars frames.
2//!
3//! Reads happen directly over the presigned HTTPS URL via range requests —
4//! only the Parquet footer and the row-groups/columns a query touches cross
5//! the wire.
6
7use polars::prelude::*;
8
9use super::ResolvedPartition;
10use crate::error::TesseraError;
11
12/// Substrings that signal a presigned URL was rejected (typically expired).
13const EXPIRY_MARKERS: [&str; 5] = [
14    "403",
15    "expired",
16    "accessdenied",
17    "access denied",
18    "forbidden",
19];
20
21/// Build a (possibly concatenated) `LazyFrame` over the resolved partitions.
22///
23/// For multi-partition reads, a `coin` and `month` column identifying the
24/// source partition are appended so rows stay attributable after concatenation.
25///
26/// # Errors
27///
28/// Returns [`TesseraError::Network`] when a Parquet path fails to register.
29pub fn build_lazyframe(
30    parts: &[ResolvedPartition],
31    columns: Option<&[&str]>,
32) -> Result<LazyFrame, TesseraError> {
33    let multi = parts.len() > 1;
34    let mut frames = Vec::with_capacity(parts.len());
35    for (partition, url) in parts {
36        let mut lazy = LazyFrame::scan_parquet(url.as_str().into(), ScanArgsParquet::default())
37            .map_err(|err| TesseraError::Network(err.to_string()))?;
38        if let Some(columns) = columns {
39            lazy = lazy.select(
40                columns
41                    .iter()
42                    .map(|column| col(*column))
43                    .collect::<Vec<_>>(),
44            );
45        }
46        if multi {
47            lazy = lazy.with_columns(vec![
48                lit(partition.coin.as_str()).alias("coin"),
49                lit(partition.month.as_str()).alias("month"),
50            ]);
51        }
52        frames.push(lazy);
53    }
54    if let [single] = &frames[..] {
55        return Ok(single.clone());
56    }
57    concat(
58        frames,
59        UnionArgs {
60            rechunk: true,
61            parallel: true,
62            to_supertypes: true,
63            ..Default::default()
64        },
65    )
66    .map_err(|err| TesseraError::Network(err.to_string()))
67}
68
69/// Collect a lazy frame, translating presign-expiry failures into a clear error.
70///
71/// Non-expiry Polars failures surface as [`TesseraError::Network`].
72///
73/// # Errors
74///
75/// Returns [`TesseraError::PresignExpired`] when the failure looks like a
76/// rejected presigned URL; other failures map to [`TesseraError::Network`].
77pub fn collect(lazy: LazyFrame) -> Result<DataFrame, TesseraError> {
78    lazy.collect().map_err(|err| {
79        let message = err.to_string().to_lowercase();
80        if EXPIRY_MARKERS.iter().any(|marker| message.contains(marker)) {
81            TesseraError::PresignExpired
82        } else {
83            TesseraError::Network(err.to_string())
84        }
85    })
86}
tessera/readers/polars.rs

tessera/readers/
polars.rs
