Merge pull request #60 from curtisalexander/mem

Mem

Author: curtisalexander

.gitignore

@@ -1,6 +1,7 @@
 # data
 /readstat/tests/data/*.csv
 /readstat/tests/data/*.feather
+/readstat/tests/data/*.ndjson
 /readstat/tests/data/*.parquet
 # environment
 .vscode/

readstat/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "readstat"
-version = "0.9.1"
+version = "0.9.2"
 authors = ["Curtis Alexander <calex@calex.org>"]
 edition = "2021"
 description = "Rust wrapper of the ReadStat C library"

readstat/Cargo.toml

@@ -32,7 +32,7 @@ pub use rs_path::ReadStatPath;
 pub use rs_write::ReadStatWriter;
 
 // Default rows to stream
-const STREAM_ROWS: u32 = 50000;
+const STREAM_ROWS: u32 = 10000;
 
 // StructOpt
 #[derive(StructOpt, Debug)]
@@ -338,7 +338,6 @@ pub fn run(rs: ReadStat) -> Result<(), Box<dyn Error + Send + Sync>> {
                     let offsets = build_offsets(total_rows_to_process, total_rows_to_stream)?;
 
                     // Create channels
-                    //let (s, r) = bounded(channel_capacity);
                     let (s, r) = unbounded();
 
                     // Initialize writing
@@ -415,7 +414,7 @@ pub fn run(rs: ReadStat) -> Result<(), Box<dyn Error + Send + Sync>> {
                     for (i, (d, rsp, pairs_cnt)) in r.iter().enumerate() {
                         wtr.write(&d, &rsp)?;
 
-                        if i == pairs_cnt {
+                        if i == (pairs_cnt - 1) {
                             wtr.finish(&d, &rsp)?;
                         }
 

readstat/src/lib.rs

@@ -160,16 +160,12 @@ impl ReadStatData {
 
         // setup parser
         // once call parse_sas7bdat, iteration begins
-        let mut parser = ReadStatParser::new()
+        let error = ReadStatParser::new()
             // do not set metadata handler nor variable handler as already processed
             .set_value_handler(Some(cb::handle_value))?
             .set_row_limit(Some(self.batch_rows_to_process.try_into().unwrap()))?
-            .set_row_offset(Some(self.batch_row_start.try_into().unwrap()))?;
-
-        let error = parser.parse_sas7bdat(ppath, ctx);
-
-        // drop parser after finished
-        drop(parser);
+            .set_row_offset(Some(self.batch_row_start.try_into().unwrap()))?
+            .parse_sas7bdat(ppath, ctx);
 
         match FromPrimitive::from_i32(error as i32) {
             Some(ReadStatError::READSTAT_OK) => Ok(()),

readstat/src/rs_data.rs

@@ -1,7 +1,6 @@
 // Create a writer struct
 use std::fs::OpenOptions;
 use std::io::stdout;
-use std::sync::Arc;
 use std::error::Error;
 
 use arrow::csv as csv_arrow;
@@ -399,18 +398,18 @@ impl ReadStatWriter {
             self.write_message_for_rows(&d, &rsp);
 
             // setup writer if not already started writing
-            if !self.wrote_start {
+            //if !self.wrote_start {
                 self.wtr = Some(ReadStatWriterFormat::Parquet(ArrowWriter::try_new(
                     f,
-                    Arc::new(d.schema.clone()),
-                    Some(WriterProperties::builder().build()),
+                    d.batch.schema(),
+                    Some(WriterProperties::builder().set_write_batch_size(d.batch_rows_to_process).set_max_row_group_size(d.batch_rows_to_process).build()),
                 )?));
-            };
+//            };
 
             // write
             if let Some(rswf) = &mut self.wtr {
                 match rswf {
-                    ReadStatWriterFormat::Parquet(wtr) => wtr.write(&d.batch)?,
+                    ReadStatWriterFormat::Parquet(wtr) => { wtr.write(&d.batch)? },
                     _ => unreachable!()
                 }
             };