readstat/
api.rs

1//! High-level convenience entry points for the common case.
2//!
3//! [`ReadStatReader`] is the primary reading API. The free functions are concise
4//! equivalents for reading a path with default settings.
5
6use std::{
7    path::Path,
8    sync::{Arc, Mutex},
9};
10
11#[cfg(not(target_arch = "wasm32"))]
12use std::path::PathBuf;
13
14use arrow_array::{RecordBatch, RecordBatchOptions};
15
16use crate::{
17    err::ReadStatError, progress::ProgressCallback, rs_data::ReadStatData,
18    rs_metadata::ReadStatMetadata, rs_path::ReadStatPath,
19};
20
21enum Source {
22    Path(ReadStatPath),
23    Bytes(Arc<[u8]>),
24    #[cfg(not(target_arch = "wasm32"))]
25    Mmap(PathBuf),
26}
27
28struct ReadPlan {
29    metadata: ReadStatMetadata,
30    mapping: Option<std::collections::BTreeMap<i32, i32>>,
31    count: u32,
32}
33
34/// High-level, reusable SAS reader.
35///
36/// Configure row and column selection once, then use [`read`](Self::read),
37/// [`chunks`](Self::chunks), or [`visit`](Self::visit). `visit` holds only the
38/// current chunk and is therefore the bounded-memory option.
39pub struct ReadStatReader {
40    source: Source,
41    metadata: Mutex<Option<ReadStatMetadata>>,
42    offset: u32,
43    limit: Option<u32>,
44    columns: Option<Vec<String>>,
45    chunk_rows: u32,
46    progress: Option<Arc<dyn ProgressCallback>>,
47}
48
49impl ReadStatReader {
50    /// Creates a reader for a filesystem path.
51    pub fn from_path(path: impl AsRef<Path>) -> Result<Self, ReadStatError> {
52        Ok(Self::new(Source::Path(ReadStatPath::new(path)?)))
53    }
54
55    /// Creates a reader owning an in-memory SAS file.
56    #[must_use]
57    pub fn from_bytes(bytes: impl Into<Arc<[u8]>>) -> Self {
58        Self::new(Source::Bytes(bytes.into()))
59    }
60
61    /// Creates a reader which memory maps `path` for each parse.
62    #[cfg(not(target_arch = "wasm32"))]
63    pub fn from_mmap(path: impl Into<PathBuf>) -> Result<Self, ReadStatError> {
64        let path = path.into();
65        let _ = ReadStatPath::new(&path)?;
66        Ok(Self::new(Source::Mmap(path)))
67    }
68
69    fn new(source: Source) -> Self {
70        Self {
71            source,
72            metadata: Mutex::new(None),
73            offset: 0,
74            limit: None,
75            columns: None,
76            chunk_rows: 10_000,
77            progress: None,
78        }
79    }
80
81    /// Selects the half-open row range beginning at `offset`, optionally limited
82    /// to `limit` rows. The range is validated against metadata when reading.
83    #[must_use]
84    pub fn rows(mut self, offset: u32, limit: Option<u32>) -> Self {
85        self.offset = offset;
86        self.limit = limit;
87        self
88    }
89
90    /// Selects columns by name, preserving dataset order.
91    ///
92    /// An empty selection deliberately produces a zero-column batch with the
93    /// requested number of rows.
94    #[must_use]
95    pub fn columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
96        self.columns = Some(columns.into_iter().map(Into::into).collect());
97        self
98    }
99
100    /// Sets rows per yielded chunk. Zero is rejected when reading.
101    #[must_use]
102    pub fn chunk_rows(mut self, rows: u32) -> Self {
103        self.chunk_rows = rows;
104        self
105    }
106
107    /// Attaches a progress callback.
108    #[must_use]
109    pub fn progress(mut self, callback: Arc<dyn ProgressCallback>) -> Self {
110        self.progress = Some(callback);
111        self
112    }
113
114    /// Reads and caches metadata transactionally.
115    ///
116    /// Reusing the same reader for metadata and data uses one cached metadata
117    /// snapshot for planning and avoids a second metadata parse. For path-backed
118    /// sources, callers must still prevent the file from being replaced between
119    /// metadata and data parsing, or between repeated reads with this reader.
120    pub fn metadata(&self) -> Result<ReadStatMetadata, ReadStatError> {
121        let mut cached = self
122            .metadata
123            .lock()
124            .map_err(|_| ReadStatError::Other("reader metadata cache is poisoned".into()))?;
125        if let Some(md) = cached.as_ref() {
126            return Ok(md.clone());
127        }
128
129        let mut md = ReadStatMetadata::new();
130        match &self.source {
131            Source::Path(path) => md.read_metadata(path, false)?,
132            Source::Bytes(bytes) => md.read_metadata_from_bytes(bytes, false)?,
133            #[cfg(not(target_arch = "wasm32"))]
134            Source::Mmap(path) => md.read_metadata_from_mmap(path, false)?,
135        }
136        *cached = Some(md.clone());
137        Ok(md)
138    }
139
140    fn plan(&self) -> Result<ReadPlan, ReadStatError> {
141        if self.chunk_rows == 0 {
142            return Err(ReadStatError::InvalidChunkSize);
143        }
144        let md = self.metadata()?;
145        let total = u32::try_from(md.row_count.ok_or(ReadStatError::RowCountUnavailable)?)?;
146        if self.offset > total {
147            return Err(ReadStatError::InvalidRowRange {
148                offset: self.offset,
149                limit: self.limit,
150                row_count: total,
151            });
152        }
153        let available = total - self.offset;
154        let count = self.limit.unwrap_or(available);
155        if count > available {
156            return Err(ReadStatError::InvalidRowRange {
157                offset: self.offset,
158                limit: self.limit,
159                row_count: total,
160            });
161        }
162        let mapping = md.resolve_selected_columns(self.columns.clone())?;
163        Ok(ReadPlan {
164            metadata: md,
165            mapping,
166            count,
167        })
168    }
169
170    /// Visits each batch without collecting previous chunks.
171    pub fn visit(
172        &self,
173        visitor: impl FnMut(RecordBatch) -> Result<(), ReadStatError>,
174    ) -> Result<(), ReadStatError> {
175        let plan = self.plan()?;
176        self.visit_with_plan(&plan, visitor)
177    }
178
179    fn visit_with_plan(
180        &self,
181        plan: &ReadPlan,
182        mut visitor: impl FnMut(RecordBatch) -> Result<(), ReadStatError>,
183    ) -> Result<(), ReadStatError> {
184        let ReadPlan {
185            metadata: md,
186            mapping,
187            count,
188        } = plan;
189        if let Some(progress) = &self.progress {
190            let label = match &self.source {
191                Source::Path(p) => p.path.to_string_lossy().into_owned(),
192                Source::Bytes(_) => "<bytes>".into(),
193                #[cfg(not(target_arch = "wasm32"))]
194                Source::Mmap(p) => p.to_string_lossy().into_owned(),
195            };
196            progress.parsing_started(&label);
197        }
198        if *count == 0 {
199            return Ok(());
200        }
201        let end = self
202            .offset
203            .checked_add(*count)
204            .ok_or_else(|| ReadStatError::Other("row offset overflow".into()))?;
205        let mut data = ReadStatData::new().init_for_visit(
206            md.clone(),
207            mapping.as_ref(),
208            self.offset,
209            end,
210            self.chunk_rows as usize,
211        );
212        if let Some(progress) = &self.progress {
213            data = data.set_progress(progress.clone());
214        }
215        match &self.source {
216            Source::Path(path) => data.visit_data(path, self.chunk_rows as usize, &mut visitor),
217            Source::Bytes(bytes) => {
218                data.visit_data_from_bytes(bytes, self.chunk_rows as usize, &mut visitor)
219            }
220            #[cfg(not(target_arch = "wasm32"))]
221            Source::Mmap(path) => {
222                data.visit_data_from_mmap(path, self.chunk_rows as usize, &mut visitor)
223            }
224        }
225    }
226
227    /// Collects all chunks.
228    pub fn chunks(&self) -> Result<Vec<RecordBatch>, ReadStatError> {
229        let mut batches = Vec::new();
230        self.visit(|batch| {
231            batches.push(batch);
232            Ok(())
233        })?;
234        Ok(batches)
235    }
236
237    /// Reads the selected rows into one batch.
238    pub fn read(&self) -> Result<RecordBatch, ReadStatError> {
239        let plan = self.plan()?;
240        let ReadPlan {
241            metadata: md,
242            mapping,
243            ..
244        } = &plan;
245        let schema = mapping.as_ref().map_or_else(
246            || md.schema.clone(),
247            |m| md.filter_to_selected_columns(m).schema,
248        );
249        let mut batches = Vec::new();
250        self.visit_with_plan(&plan, |batch| {
251            batches.push(batch);
252            Ok(())
253        })?;
254        if batches.is_empty() {
255            return Ok(RecordBatch::new_empty(Arc::new(schema)));
256        }
257        if schema.fields().is_empty() {
258            return RecordBatch::try_new_with_options(
259                Arc::new(schema),
260                Vec::new(),
261                &RecordBatchOptions::new().with_row_count(Some(plan.count as usize)),
262            )
263            .map_err(Into::into);
264        }
265        arrow::compute::concat_batches(&Arc::new(schema), &batches).map_err(Into::into)
266    }
267}
268
269/// Reads file-level and variable metadata from a `.sas7bdat` file without
270/// loading any row data.
271///
272/// This delegates to [`ReadStatReader::metadata`].
273///
274/// ```no_run
275/// # fn main() -> Result<(), readstat::ReadStatError> {
276/// let md = readstat::read_metadata("data.sas7bdat")?;
277/// println!("{:?} rows x {} columns", md.row_count, md.var_count);
278/// # Ok(())
279/// # }
280/// ```
281///
282/// # Errors
283///
284/// Returns [`ReadStatError`] if the path is invalid or FFI parsing fails.
285pub fn read_metadata<P: AsRef<Path>>(path: P) -> Result<ReadStatMetadata, ReadStatError> {
286    ReadStatReader::from_path(path)?.metadata()
287}
288
289/// Reads every row of a `.sas7bdat` file into a single Arrow [`RecordBatch`].
290///
291/// Best for files that fit comfortably in memory. For large files, use
292/// [`ReadStatReader::visit`] to process bounded chunks.
293///
294/// ```no_run
295/// # fn main() -> Result<(), readstat::ReadStatError> {
296/// let batch = readstat::read_to_batch("data.sas7bdat")?;
297/// println!("{} rows x {} columns", batch.num_rows(), batch.num_columns());
298/// # Ok(())
299/// # }
300/// ```
301///
302/// # Errors
303///
304/// Returns [`ReadStatError`] if the path is invalid, FFI parsing fails, or the
305/// row count cannot be represented (i.e. is negative).
306pub fn read_to_batch<P: AsRef<Path>>(path: P) -> Result<RecordBatch, ReadStatError> {
307    ReadStatReader::from_path(path)?.read()
308}
readstat/api.rs

readstat/
api.rs
