readstat/
rs_data.rs

1//! Data reading and Arrow [`RecordBatch`] conversion.
2//!
3//! [`ReadStatData`] coordinates the FFI parsing of row values from a `.sas7bdat` file,
4//! accumulating them directly into typed Arrow builders via the `handle_value`
5//! callback, then finishing them into an Arrow `RecordBatch` for downstream writing.
6//! Supports streaming chunks with configurable row offsets and progress tracking.
7
8use arrow::datatypes::Schema;
9use arrow_array::{
10    ArrayRef, RecordBatch, RecordBatchOptions,
11    builder::{
12        Date32Builder, Float32Builder, Float64Builder, Int16Builder, Int32Builder, StringBuilder,
13        Time32MillisecondBuilder, Time32SecondBuilder, Time64MicrosecondBuilder,
14        Time64NanosecondBuilder, TimestampMicrosecondBuilder, TimestampMillisecondBuilder,
15        TimestampNanosecondBuilder, TimestampSecondBuilder,
16    },
17};
18use log::debug;
19use std::{
20    collections::BTreeMap,
21    ffi::CString,
22    os::raw::c_void,
23    sync::{Arc, atomic::AtomicUsize},
24};
25
26use crate::{
27    cb,
28    err::{ReadStatError, check_c_error},
29    progress::ProgressCallback,
30    rs_buffer_io::ReadStatBufferCtx,
31    rs_metadata::{ReadStatMetadata, ReadStatVarMetadata},
32    rs_parser::ReadStatParser,
33    rs_path::ReadStatPath,
34    rs_var::{ReadStatVarFormatClass, ReadStatVarType, ReadStatVarTypeClass},
35};
36
37/// Upper bound on the row capacity pre-allocated for Arrow builders.
38///
39/// The claimed row count comes from an untrusted file header, so the up-front
40/// allocation is capped here; builders grow on demand past this for honest
41/// files. 1,000,000 rows is far beyond the default 10k streaming chunk while
42/// keeping the worst-case empty-builder reservation bounded.
43const MAX_PREALLOC_ROWS: usize = 1_000_000;
44/// Maximum aggregate initial row capacity across all string columns.
45const MAX_STRING_PREALLOC_ROWS: usize = 1_000_000;
46/// Maximum initial value-byte reservation across all string columns.
47const MAX_STRING_PREALLOC_BYTES: usize = 16 * 1024 * 1024;
48/// Maximum share of the aggregate reservation assigned to one string column.
49const MAX_STRING_COLUMN_PREALLOC_BYTES: usize = 1024 * 1024;
50
51fn string_column_byte_budget(string_columns: usize) -> usize {
52    MAX_STRING_PREALLOC_BYTES
53        .checked_div(string_columns)
54        .unwrap_or(0)
55        .min(MAX_STRING_COLUMN_PREALLOC_BYTES)
56}
57
58fn string_column_row_budget(string_columns: usize) -> usize {
59    MAX_STRING_PREALLOC_ROWS
60        .checked_div(string_columns)
61        .unwrap_or(0)
62}
63
64fn string_value_capacity(capacity: usize, storage_width: usize, byte_budget: usize) -> usize {
65    capacity.saturating_mul(storage_width).min(byte_budget)
66}
67
68/// A typed Arrow array builder for a single column.
69///
70/// Each variant wraps the corresponding Arrow builder, pre-sized with capacity
71/// hints from the metadata (row count, string `storage_width`). Values are
72/// appended directly during FFI callbacks, eliminating intermediate allocations.
73pub(crate) enum ColumnBuilder {
74    /// UTF-8 string column.
75    Str(StringBuilder),
76    /// 16-bit signed integer column (covers both SAS Int8 and Int16).
77    Int16(Int16Builder),
78    /// 32-bit signed integer column.
79    Int32(Int32Builder),
80    /// 32-bit floating point column.
81    Float32(Float32Builder),
82    /// 64-bit floating point column.
83    Float64(Float64Builder),
84    /// Date column (days since Unix epoch).
85    Date32(Date32Builder),
86    /// Timestamp with second precision.
87    TimestampSecond(TimestampSecondBuilder),
88    /// Timestamp with millisecond precision.
89    TimestampMillisecond(TimestampMillisecondBuilder),
90    /// Timestamp with microsecond precision.
91    TimestampMicrosecond(TimestampMicrosecondBuilder),
92    /// Timestamp with nanosecond precision.
93    TimestampNanosecond(TimestampNanosecondBuilder),
94    /// Time of day with second precision.
95    Time32Second(Time32SecondBuilder),
96    /// Time of day with millisecond precision.
97    Time32Millisecond(Time32MillisecondBuilder),
98    /// Time of day with microsecond precision.
99    Time64Microsecond(Time64MicrosecondBuilder),
100    /// Time of day with nanosecond precision.
101    Time64Nanosecond(Time64NanosecondBuilder),
102}
103
104impl ColumnBuilder {
105    /// Appends a null value, regardless of the underlying builder type.
106    pub(crate) fn append_null(&mut self) {
107        match self {
108            Self::Str(b) => b.append_null(),
109            Self::Int16(b) => b.append_null(),
110            Self::Int32(b) => b.append_null(),
111            Self::Float32(b) => b.append_null(),
112            Self::Float64(b) => b.append_null(),
113            Self::Date32(b) => b.append_null(),
114            Self::TimestampSecond(b) => b.append_null(),
115            Self::TimestampMillisecond(b) => b.append_null(),
116            Self::TimestampMicrosecond(b) => b.append_null(),
117            Self::TimestampNanosecond(b) => b.append_null(),
118            Self::Time32Second(b) => b.append_null(),
119            Self::Time32Millisecond(b) => b.append_null(),
120            Self::Time64Microsecond(b) => b.append_null(),
121            Self::Time64Nanosecond(b) => b.append_null(),
122        }
123    }
124
125    /// Finishes the builder and returns the completed Arrow array.
126    pub(crate) fn finish(&mut self) -> ArrayRef {
127        match self {
128            Self::Str(b) => Arc::new(b.finish()),
129            Self::Int16(b) => Arc::new(b.finish()),
130            Self::Int32(b) => Arc::new(b.finish()),
131            Self::Float32(b) => Arc::new(b.finish()),
132            Self::Float64(b) => Arc::new(b.finish()),
133            Self::Date32(b) => Arc::new(b.finish()),
134            Self::TimestampSecond(b) => Arc::new(b.finish()),
135            Self::TimestampMillisecond(b) => Arc::new(b.finish()),
136            Self::TimestampMicrosecond(b) => Arc::new(b.finish()),
137            Self::TimestampNanosecond(b) => Arc::new(b.finish()),
138            Self::Time32Second(b) => Arc::new(b.finish()),
139            Self::Time32Millisecond(b) => Arc::new(b.finish()),
140            Self::Time64Microsecond(b) => Arc::new(b.finish()),
141            Self::Time64Nanosecond(b) => Arc::new(b.finish()),
142        }
143    }
144
145    /// Creates a typed builder matching the variable's metadata.
146    ///
147    /// Uses `var_type`, `var_type_class`, and `var_format_class` to select the
148    /// correct builder variant, and pre-sizes it with `capacity` rows.
149    /// For string columns, `storage_width` provides a byte-level capacity hint.
150    fn from_metadata(
151        vm: &ReadStatVarMetadata,
152        capacity: usize,
153        string_row_budget: usize,
154        string_byte_budget: usize,
155    ) -> Self {
156        match vm.var_type_class {
157            ReadStatVarTypeClass::String => {
158                let string_capacity = capacity.min(string_row_budget);
159                Self::Str(StringBuilder::with_capacity(
160                    string_capacity,
161                    string_value_capacity(string_capacity, vm.storage_width, string_byte_budget),
162                ))
163            }
164            ReadStatVarTypeClass::Numeric => {
165                match vm.var_format_class {
166                    Some(ReadStatVarFormatClass::Date) => {
167                        Self::Date32(Date32Builder::with_capacity(capacity))
168                    }
169                    Some(ReadStatVarFormatClass::DateTime) => {
170                        Self::TimestampSecond(TimestampSecondBuilder::with_capacity(capacity))
171                    }
172                    Some(ReadStatVarFormatClass::DateTimeWithMilliseconds) => {
173                        Self::TimestampMillisecond(TimestampMillisecondBuilder::with_capacity(
174                            capacity,
175                        ))
176                    }
177                    Some(ReadStatVarFormatClass::DateTimeWithMicroseconds) => {
178                        Self::TimestampMicrosecond(TimestampMicrosecondBuilder::with_capacity(
179                            capacity,
180                        ))
181                    }
182                    Some(ReadStatVarFormatClass::DateTimeWithNanoseconds) => {
183                        Self::TimestampNanosecond(TimestampNanosecondBuilder::with_capacity(
184                            capacity,
185                        ))
186                    }
187                    Some(ReadStatVarFormatClass::Time) => {
188                        Self::Time32Second(Time32SecondBuilder::with_capacity(capacity))
189                    }
190                    Some(ReadStatVarFormatClass::TimeWithMilliseconds) => {
191                        Self::Time32Millisecond(Time32MillisecondBuilder::with_capacity(capacity))
192                    }
193                    Some(ReadStatVarFormatClass::TimeWithMicroseconds) => {
194                        Self::Time64Microsecond(Time64MicrosecondBuilder::with_capacity(capacity))
195                    }
196                    Some(ReadStatVarFormatClass::TimeWithNanoseconds) => {
197                        Self::Time64Nanosecond(Time64NanosecondBuilder::with_capacity(capacity))
198                    }
199                    None => {
200                        // Plain numeric — dispatch by storage type
201                        match vm.var_type {
202                            ReadStatVarType::Int8 | ReadStatVarType::Int16 => {
203                                Self::Int16(Int16Builder::with_capacity(capacity))
204                            }
205                            ReadStatVarType::Int32 => {
206                                Self::Int32(Int32Builder::with_capacity(capacity))
207                            }
208                            ReadStatVarType::Float => {
209                                Self::Float32(Float32Builder::with_capacity(capacity))
210                            }
211                            _ => Self::Float64(Float64Builder::with_capacity(capacity)),
212                        }
213                    }
214                }
215            }
216        }
217    }
218}
219
220/// Holds parsed row data from a `.sas7bdat` file and converts it to Arrow format.
221///
222/// Values are appended directly into typed Arrow `ColumnBuilder`s during the
223/// `handle_value` callback, then finished into an Arrow [`RecordBatch`] via
224/// `cols_to_batch`. The high-level streaming reader reuses the parse context
225/// while rotating these builders between bounded output batches.
226pub struct ReadStatData {
227    /// Number of variables (columns) in the dataset.
228    pub var_count: i32,
229    /// Per-variable metadata, keyed by variable index.
230    /// Wrapped in `Arc` so parallel chunks share the same metadata without deep cloning.
231    pub vars: Arc<BTreeMap<i32, ReadStatVarMetadata>>,
232    /// Typed Arrow builders — one per variable, pre-sized with capacity hints.
233    pub(crate) builders: Vec<ColumnBuilder>,
234    /// Arrow schema for the dataset.
235    /// Wrapped in `Arc` for cheap sharing across parallel chunks.
236    pub schema: Arc<Schema>,
237    /// The Arrow `RecordBatch` produced after parsing, if available.
238    pub batch: Option<RecordBatch>,
239    /// Number of rows to process in this chunk.
240    pub chunk_rows_to_process: usize,
241    /// Starting row offset for this chunk.
242    pub(crate) chunk_row_start: usize,
243    /// Ending row offset (exclusive) for this chunk.
244    pub(crate) chunk_row_end: usize,
245    /// Number of rows actually processed so far in this chunk.
246    pub(crate) chunk_rows_processed: usize,
247    /// Shared atomic counter of total rows processed across all chunks.
248    pub(crate) total_rows_processed: Option<Arc<AtomicUsize>>,
249    /// Optional progress callback for visual feedback during parsing.
250    pub(crate) progress: Option<Arc<dyn ProgressCallback>>,
251    /// A typed error raised by a value callback that aborted parsing.
252    ///
253    /// Set by `handle_value` (e.g. on date/time overflow or a builder/value
254    /// type mismatch) and surfaced by the parse routines in preference to the
255    /// generic `USER_ABORT` the C library reports for any callback abort.
256    pub(crate) abort_error: Option<ReadStatError>,
257    /// Optional mapping: original var index -> filtered column index.
258    /// Wrapped in `Arc` so parallel chunks share the same filter without deep cloning.
259    pub(crate) column_filter: Option<Arc<BTreeMap<i32, i32>>>,
260    /// Total variable count in the unfiltered dataset.
261    /// Used for row-boundary detection in `handle_value` when filtering is active.
262    /// Defaults to `var_count` when no filter is set.
263    pub(crate) total_var_count: i32,
264}
265
266/// Callback context used by the high-level reader to rotate builders while a
267/// single ReadStat parse is in progress.
268pub(crate) struct StreamingData<'a> {
269    pub(crate) data: ReadStatData,
270    chunk_rows: usize,
271    expected_rows: usize,
272    rows_emitted: usize,
273    sink: &'a mut dyn FnMut(RecordBatch) -> Result<(), ReadStatError>,
274}
275
276impl StreamingData<'_> {
277    pub(crate) fn finish_chunk(&mut self) -> Result<(), ReadStatError> {
278        if self.data.chunk_rows_processed == 0 {
279            return Ok(());
280        }
281        self.data.cols_to_batch()?;
282        let batch = self
283            .data
284            .batch
285            .take()
286            .ok_or_else(|| ReadStatError::Other("no record batch was produced".into()))?;
287        let rows = batch.num_rows();
288        (self.sink)(batch)?;
289        self.rows_emitted += rows;
290        if let Some(progress) = &self.data.progress {
291            progress.inc(rows as u64);
292        }
293        self.data.chunk_rows_processed = 0;
294        if self.rows_emitted < self.expected_rows {
295            let data = std::mem::take(&mut self.data);
296            self.data = data.allocate_builders_with_capacity(self.chunk_rows);
297        }
298        Ok(())
299    }
300
301    pub(crate) fn row_complete(&mut self) -> Result<(), ReadStatError> {
302        if self.data.chunk_rows_processed == self.chunk_rows {
303            self.finish_chunk()?;
304        }
305        Ok(())
306    }
307}
308
309impl Default for ReadStatData {
310    fn default() -> Self {
311        Self::new()
312    }
313}
314
315impl ReadStatData {
316    /// Creates a new `ReadStatData` with default (empty) values.
317    pub fn new() -> Self {
318        Self {
319            // metadata
320            var_count: 0,
321            vars: Arc::new(BTreeMap::new()),
322            // data
323            builders: Vec::new(),
324            schema: Arc::new(Schema::empty()),
325            // record batch
326            batch: None,
327            chunk_rows_to_process: 0,
328            chunk_rows_processed: 0,
329            chunk_row_start: 0,
330            chunk_row_end: 0,
331            // total rows
332            total_rows_processed: None,
333            // progress
334            progress: None,
335            // errors
336            abort_error: None,
337            // column filtering
338            column_filter: None,
339            total_var_count: 0,
340        }
341    }
342
343    /// Allocates typed Arrow builders with capacity for `chunk_rows_to_process`.
344    ///
345    /// Each builder's type is determined by the variable metadata. String builders
346    /// are additionally pre-sized with `storage_width * chunk_rows` bytes.
347    ///
348    /// The capacity hint is clamped to `MAX_PREALLOC_ROWS` (1,000,000 rows) because both the row
349    /// count and per-string `storage_width` originate from untrusted file headers;
350    /// a crafted file claiming billions of rows would otherwise trigger a multi-GB
351    /// up-front allocation (or a multiply overflow) before a single row is parsed.
352    /// Builders grow on demand, so clamping costs honest files nothing.
353    #[must_use]
354    pub fn allocate_builders(self) -> Self {
355        let capacity = self.chunk_rows_to_process;
356        self.allocate_builders_with_capacity(capacity)
357    }
358
359    fn allocate_builders_with_capacity(self, capacity: usize) -> Self {
360        let capacity = capacity
361            .min(self.chunk_rows_to_process)
362            .min(MAX_PREALLOC_ROWS);
363        let string_columns = self
364            .vars
365            .values()
366            .filter(|vm| matches!(vm.var_type_class, ReadStatVarTypeClass::String))
367            .count();
368        let string_byte_budget = string_column_byte_budget(string_columns);
369        let string_row_budget = string_column_row_budget(string_columns);
370        let builders: Vec<ColumnBuilder> = self
371            .vars
372            .values()
373            .map(|vm| {
374                ColumnBuilder::from_metadata(vm, capacity, string_row_budget, string_byte_budget)
375            })
376            .collect();
377        Self { builders, ..self }
378    }
379
380    /// Finishes all builders and assembles the Arrow [`RecordBatch`].
381    ///
382    /// Each builder produces its final array via `finish()`, which is an O(1)
383    /// operation (no data copying). The heavy work was already done during
384    /// `handle_value` when values were appended directly into the builders.
385    pub(crate) fn cols_to_batch(&mut self) -> Result<(), ReadStatError> {
386        let arrays: Vec<ArrayRef> = self
387            .builders
388            .iter_mut()
389            .map(ColumnBuilder::finish)
390            .collect();
391
392        self.batch = Some(if arrays.is_empty() {
393            RecordBatch::try_new_with_options(
394                self.schema.clone(),
395                arrays,
396                &RecordBatchOptions::new().with_row_count(Some(self.chunk_rows_processed)),
397            )?
398        } else {
399            RecordBatch::try_new(self.schema.clone(), arrays)?
400        });
401
402        Ok(())
403    }
404
405    /// Records that a value was observed for `var_index` during parsing.
406    ///
407    /// When `var_index` is the dataset's final variable, the cell marks the end
408    /// of a row, so the per-chunk and shared row counters are advanced. Boundary
409    /// detection uses `total_var_count` (the *unfiltered* variable count) so it
410    /// stays correct even when a column filter skips trailing columns.
411    ///
412    /// Called from the value callback for both stored and filter-skipped cells,
413    /// keeping row-boundary accounting in a single place.
414    pub(crate) fn note_value(&mut self, var_index: i32) {
415        if var_index == self.total_var_count - 1 {
416            self.chunk_rows_processed += 1;
417            if let Some(trp) = &self.total_rows_processed {
418                trp.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
419            }
420        }
421    }
422
423    /// Parses row data from the file and converts it to an Arrow [`RecordBatch`].
424    ///
425    /// # Errors
426    ///
427    /// Returns [`ReadStatError`] if FFI parsing or Arrow conversion fails.
428    pub fn read_data(&mut self, rsp: &ReadStatPath) -> Result<(), ReadStatError> {
429        // parse data and if successful then convert cols into a record batch
430        self.parse_data(rsp)?;
431        self.cols_to_batch()?;
432        self.report_progress();
433        Ok(())
434    }
435
436    /// Parses row data from an in-memory byte slice and converts it to an Arrow [`RecordBatch`].
437    ///
438    /// Equivalent to [`read_data`](ReadStatData::read_data) but reads from a `&[u8]`
439    /// buffer instead of a file path.
440    ///
441    /// # Errors
442    ///
443    /// Returns [`ReadStatError`] if FFI parsing or Arrow conversion fails.
444    pub fn read_data_from_bytes(&mut self, bytes: &[u8]) -> Result<(), ReadStatError> {
445        self.parse_data_from_bytes(bytes)?;
446        self.cols_to_batch()?;
447        self.report_progress();
448        Ok(())
449    }
450
451    /// Parses row data from a memory-mapped `.sas7bdat` file and converts it to an Arrow [`RecordBatch`].
452    ///
453    /// Opens the file at `path` and memory-maps it, avoiding explicit read syscalls.
454    /// Especially beneficial for large files and repeated chunk reads against the
455    /// same file, as the OS manages page caching automatically.
456    ///
457    /// # Safety
458    ///
459    /// Memory mapping is safe as long as the file is not modified or truncated by
460    /// another process while the map is active.
461    ///
462    /// # Errors
463    ///
464    /// Returns [`ReadStatError`] if the file cannot be opened, mapped, or parsed.
465    #[cfg(not(target_arch = "wasm32"))]
466    pub fn read_data_from_mmap(&mut self, path: &std::path::Path) -> Result<(), ReadStatError> {
467        let file = std::fs::File::open(path)?;
468        let mmap = unsafe { memmap2::Mmap::map(&file)? };
469        self.read_data_from_bytes(&mmap)
470    }
471
472    pub(crate) fn visit_data(
473        self,
474        rsp: &ReadStatPath,
475        chunk_rows: usize,
476        sink: &mut dyn FnMut(RecordBatch) -> Result<(), ReadStatError>,
477    ) -> Result<(), ReadStatError> {
478        let ppath = rsp.cstring_path.as_ptr();
479        self.visit_parse(chunk_rows, sink, |mut parser, ctx| {
480            Ok(parser.parse_sas7bdat(ppath, ctx))
481        })
482    }
483
484    pub(crate) fn visit_data_from_bytes(
485        self,
486        bytes: &[u8],
487        chunk_rows: usize,
488        sink: &mut dyn FnMut(RecordBatch) -> Result<(), ReadStatError>,
489    ) -> Result<(), ReadStatError> {
490        let mut buffer_ctx = ReadStatBufferCtx::new(bytes);
491        self.visit_parse(chunk_rows, sink, |parser, ctx| {
492            buffer_ctx
493                .configure_parser(parser)
494                .map(|mut p| p.parse_sas7bdat(std::ptr::null(), ctx))
495        })
496    }
497
498    #[cfg(not(target_arch = "wasm32"))]
499    pub(crate) fn visit_data_from_mmap(
500        self,
501        path: &std::path::Path,
502        chunk_rows: usize,
503        sink: &mut dyn FnMut(RecordBatch) -> Result<(), ReadStatError>,
504    ) -> Result<(), ReadStatError> {
505        let file = std::fs::File::open(path)?;
506        let mmap = unsafe { memmap2::Mmap::map(&file)? };
507        self.visit_data_from_bytes(&mmap, chunk_rows, sink)
508    }
509
510    fn visit_parse(
511        self,
512        chunk_rows: usize,
513        sink: &mut dyn FnMut(RecordBatch) -> Result<(), ReadStatError>,
514        parse: impl FnOnce(
515            ReadStatParser,
516            *mut c_void,
517        ) -> Result<readstat_sys::readstat_error_t, ReadStatError>,
518    ) -> Result<(), ReadStatError> {
519        let offset = self.chunk_row_start.try_into()?;
520        let limit = self.chunk_rows_to_process.try_into()?;
521        let expected_rows = self.chunk_rows_to_process;
522        let mut stream = StreamingData {
523            data: self,
524            chunk_rows,
525            expected_rows,
526            rows_emitted: 0,
527            sink,
528        };
529        let ctx = std::ptr::from_mut(&mut stream).cast::<c_void>();
530        let parser = ReadStatParser::new()?
531            .set_value_handler(Some(cb::handle_streaming_value))?
532            .set_row_limit(Some(limit))?
533            .set_row_offset(Some(offset))?;
534        let result = parse(parser, ctx);
535        if let Some(error) = stream.data.abort_error.take() {
536            return Err(error);
537        }
538        check_c_error(result? as i32)?;
539        stream.finish_chunk()?;
540        if stream.rows_emitted != stream.expected_rows {
541            return Err(ReadStatError::Other(format!(
542                "ReadStat emitted {} rows, expected {}",
543                stream.rows_emitted, stream.expected_rows
544            )));
545        }
546        Ok(())
547    }
548
549    /// Parses row data from the file via FFI callbacks (without Arrow conversion).
550    #[allow(clippy::cast_possible_wrap, clippy::ptr_as_ptr)]
551    pub(crate) fn parse_data(&mut self, rsp: &ReadStatPath) -> Result<(), ReadStatError> {
552        // path as pointer
553        debug!("Path as C string is {:?}", rsp.cstring_path);
554        let ppath = rsp.cstring_path.as_ptr();
555
556        // initialize context
557        let ctx = std::ptr::from_mut::<Self>(self) as *mut c_void;
558
559        // initialize error
560        let error: readstat_sys::readstat_error_t = readstat_sys::readstat_error_e_READSTAT_OK;
561        debug!("Initially, error ==> {error:#?}");
562
563        // setup parser
564        // once call parse_sas7bdat, iteration begins
565        let error = ReadStatParser::new()?
566            // do not set metadata handler nor variable handler as already processed
567            .set_value_handler(Some(cb::handle_value))?
568            .set_row_limit(Some(self.chunk_rows_to_process.try_into()?))?
569            .set_row_offset(Some(self.chunk_row_start.try_into()?))?
570            .parse_sas7bdat(ppath, ctx);
571
572        // A value callback may have aborted with a specific, typed error; prefer
573        // it over the generic `USER_ABORT` the C library reports for any abort.
574        if let Some(e) = self.abort_error.take() {
575            return Err(e);
576        }
577        check_c_error(error as i32)?;
578
579        Ok(())
580    }
581
582    fn report_progress(&self) {
583        if let Some(progress) = &self.progress {
584            progress.inc(self.batch.as_ref().map_or(0, RecordBatch::num_rows) as u64);
585        }
586    }
587
588    #[allow(clippy::cast_possible_wrap, clippy::ptr_as_ptr)]
589    fn parse_data_from_bytes(&mut self, bytes: &[u8]) -> Result<(), ReadStatError> {
590        let mut buffer_ctx = ReadStatBufferCtx::new(bytes);
591
592        // initialize context
593        let ctx = std::ptr::from_mut::<Self>(self) as *mut c_void;
594
595        // initialize error
596        let error: readstat_sys::readstat_error_t = readstat_sys::readstat_error_e_READSTAT_OK;
597        debug!("Initially, error ==> {error:#?}");
598
599        // Dummy path — custom I/O handlers ignore it
600        let dummy_path = CString::new("").expect("empty string is valid C string");
601
602        // setup parser with buffer I/O
603        let error = buffer_ctx
604            .configure_parser(
605                ReadStatParser::new()?
606                    .set_value_handler(Some(cb::handle_value))?
607                    .set_row_limit(Some(self.chunk_rows_to_process.try_into()?))?
608                    .set_row_offset(Some(self.chunk_row_start.try_into()?))?,
609            )?
610            .parse_sas7bdat(dummy_path.as_ptr(), ctx);
611
612        // A value callback may have aborted with a specific, typed error; prefer
613        // it over the generic `USER_ABORT` the C library reports for any abort.
614        if let Some(e) = self.abort_error.take() {
615            return Err(e);
616        }
617        check_c_error(error as i32)?;
618        Ok(())
619    }
620
621    /// Initializes this instance with metadata and chunk boundaries, allocating builders.
622    ///
623    /// Wraps `vars` and `schema` in `Arc` internally. For the parallel read path,
624    /// prefer [`init_shared`](ReadStatData::init_shared) which accepts pre-wrapped
625    /// `Arc`s to avoid repeated deep clones.
626    #[must_use]
627    pub fn init(self, md: ReadStatMetadata, row_start: u32, row_end: u32) -> Self {
628        self.set_metadata(md)
629            .set_chunk_counts(row_start, row_end)
630            .allocate_builders()
631    }
632
633    /// Initializes this instance with a column filter applied, in one step.
634    ///
635    /// Combines [`set_column_filter`](ReadStatData::set_column_filter) and
636    /// [`init`](ReadStatData::init) in the correct order so callers cannot
637    /// accidentally invoke them the wrong way around (which would clobber the
638    /// original variable count needed for row-boundary detection).
639    ///
640    /// `md` must be the **original, unfiltered** metadata and `mapping` the
641    /// result of [`ReadStatMetadata::resolve_selected_columns`]. The filtered
642    /// metadata and the original variable count are derived internally.
643    ///
644    /// ```no_run
645    /// use readstat::{ReadStatPath, ReadStatMetadata, ReadStatData};
646    ///
647    /// # fn main() -> Result<(), readstat::ReadStatError> {
648    /// let rsp = ReadStatPath::new("data.sas7bdat")?;
649    /// let mut md = ReadStatMetadata::new();
650    /// md.read_metadata(&rsp, false)?;
651    ///
652    /// if let Some(mapping) = md.resolve_selected_columns(Some(vec!["name".into(), "age".into()]))? {
653    ///     let row_count = u32::try_from(md.row_count.ok_or(readstat::ReadStatError::RowCountUnavailable)?)?;
654    ///     let mut d = ReadStatData::new().init_filtered(md, &mapping, 0, row_count);
655    ///     d.read_data(&rsp)?;
656    /// }
657    /// # Ok(())
658    /// # }
659    /// ```
660    #[must_use]
661    pub fn init_filtered(
662        self,
663        md: ReadStatMetadata,
664        mapping: &BTreeMap<i32, i32>,
665        row_start: u32,
666        row_end: u32,
667    ) -> Self {
668        let original_var_count = md.var_count;
669        let filtered = md.filter_to_selected_columns(mapping);
670        self.set_column_filter(Some(Arc::new(mapping.clone())), original_var_count)
671            .init(filtered, row_start, row_end)
672    }
673
674    pub(crate) fn init_for_visit(
675        self,
676        md: ReadStatMetadata,
677        mapping: Option<&BTreeMap<i32, i32>>,
678        row_start: u32,
679        row_end: u32,
680        batch_rows: usize,
681    ) -> Self {
682        let data = if let Some(mapping) = mapping {
683            let original_var_count = md.var_count;
684            let filtered = md.filter_to_selected_columns(mapping);
685            self.set_column_filter(Some(Arc::new(mapping.clone())), original_var_count)
686                .set_metadata(filtered)
687        } else {
688            self.set_metadata(md)
689        };
690        data.set_chunk_counts(row_start, row_end)
691            .allocate_builders_with_capacity(batch_rows)
692    }
693
694    /// Initializes this instance with pre-shared metadata and chunk boundaries.
695    ///
696    /// Accepts `Arc`-wrapped `vars` and `schema` for cheap cloning in parallel loops.
697    /// Each call only increments reference counts (atomic +1) instead of deep-cloning
698    /// the entire metadata tree.
699    #[must_use]
700    pub fn init_shared(
701        self,
702        var_count: i32,
703        vars: Arc<BTreeMap<i32, ReadStatVarMetadata>>,
704        schema: Arc<Schema>,
705        row_start: u32,
706        row_end: u32,
707    ) -> Self {
708        let total_var_count = if self.total_var_count != 0 {
709            self.total_var_count
710        } else {
711            var_count
712        };
713        Self {
714            var_count,
715            vars,
716            schema,
717            total_var_count,
718            ..self
719        }
720        .set_chunk_counts(row_start, row_end)
721        .allocate_builders()
722    }
723
724    #[allow(clippy::cast_possible_truncation)]
725    fn set_chunk_counts(self, row_start: u32, row_end: u32) -> Self {
726        // saturating_sub: guard against a caller passing row_end < row_start,
727        // which would underflow-panic in debug and wrap to ~4 billion in
728        // release (then feed an enormous builder pre-allocation).
729        let chunk_rows_to_process = row_end.saturating_sub(row_start) as usize;
730        let chunk_row_start = row_start as usize;
731        let chunk_row_end = row_end as usize;
732        let chunk_rows_processed = 0_usize;
733
734        Self {
735            chunk_rows_to_process,
736            chunk_row_start,
737            chunk_row_end,
738            chunk_rows_processed,
739            ..self
740        }
741    }
742
743    fn set_metadata(self, md: ReadStatMetadata) -> Self {
744        let var_count = md.var_count;
745        let vars = Arc::new(md.vars);
746        let schema = Arc::new(md.schema);
747        // Only set total_var_count from metadata if not already set by set_column_filter
748        let total_var_count = if self.total_var_count != 0 {
749            self.total_var_count
750        } else {
751            var_count
752        };
753        Self {
754            var_count,
755            vars,
756            schema,
757            total_var_count,
758            ..self
759        }
760    }
761
762    /// Sets the shared atomic counter for tracking rows processed across chunks.
763    #[must_use]
764    pub fn set_total_rows_processed(self, total_rows_processed: Arc<AtomicUsize>) -> Self {
765        Self {
766            total_rows_processed: Some(total_rows_processed),
767            ..self
768        }
769    }
770
771    /// Sets the column filter and original (unfiltered) variable count.
772    ///
773    /// Accepts an `Arc`-wrapped filter for cheap sharing across parallel chunks.
774    /// Must be called **before** [`init`](ReadStatData::init) so that
775    /// `total_var_count` is preserved when `set_metadata` runs.
776    #[must_use]
777    pub fn set_column_filter(
778        self,
779        filter: Option<Arc<BTreeMap<i32, i32>>>,
780        total_var_count: i32,
781    ) -> Self {
782        Self {
783            column_filter: filter,
784            total_var_count,
785            ..self
786        }
787    }
788
789    /// Attaches a progress callback for feedback during parsing.
790    ///
791    /// The callback receives progress increments and parsing status updates.
792    /// See [`ProgressCallback`] for the required interface.
793    #[must_use]
794    pub fn set_progress(self, progress: Arc<dyn ProgressCallback>) -> Self {
795        Self {
796            progress: Some(progress),
797            ..self
798        }
799    }
800}
801
802#[cfg(test)]
803mod allocation_tests {
804    use super::*;
805
806    #[test]
807    fn string_byte_hint_is_normal_for_small_columns() {
808        assert_eq!(string_value_capacity(100, 20, 1024 * 1024), 2_000);
809    }
810
811    #[test]
812    fn string_byte_hint_caps_max_sas_width() {
813        assert_eq!(
814            string_value_capacity(1_000_000, 32_767, MAX_STRING_COLUMN_PREALLOC_BYTES),
815            MAX_STRING_COLUMN_PREALLOC_BYTES
816        );
817    }
818
819    #[test]
820    fn string_byte_hint_caps_hostile_overflow() {
821        assert_eq!(
822            string_value_capacity(usize::MAX, usize::MAX, MAX_STRING_COLUMN_PREALLOC_BYTES),
823            MAX_STRING_COLUMN_PREALLOC_BYTES
824        );
825    }
826
827    #[test]
828    fn string_byte_budget_is_bounded_across_many_columns() {
829        let columns = 32_767;
830        let per_column = string_column_byte_budget(columns);
831        assert!(per_column * columns <= MAX_STRING_PREALLOC_BYTES);
832    }
833
834    #[test]
835    fn string_row_budget_is_bounded_across_many_columns() {
836        let columns = 32_767;
837        let per_column = string_column_row_budget(columns);
838        assert!(per_column * columns <= MAX_STRING_PREALLOC_ROWS);
839    }
840}
readstat/rs_data.rs

readstat/
rs_data.rs
