Merge pull request #41 from firelink-data/feature/threaded_revolver

Feature/threaded revolver

Author: Ignalina

Cargo.toml

@@ -25,19 +25,19 @@ include = [ "**/*.rs", "Cargo.toml", "LICENSE", "README.md" ]
 default-run = "evolution"
 
 [dependencies]
-chrono = "0.4.31"
-crossbeam = "0.8.2"
+chrono = "0.4.38"
+crossbeam = "0.8.4"
 colored = "2.0.4"
-env_logger = "0.11.1"
+env_logger = "0.11.3"
 half = "2.3.1"
 log = "0.4.20"
 num_cpus = "1.16.0"
 rand = "0.8.5"
-serde_json = "1.0.108"
+serde_json = "1.0.117"
 threadpool = "1.8.1"
-clap = { version = "4.4.8", features = ["default", "derive"] }
-serde = { version = "1.0.193", features = ["derive"] }
-rayon = { version = "1.8.1", optional = true }
+clap = { version = "4.5.4", features = ["default", "derive"] }
+serde = { version = "1.0.202", features = ["derive"] }
+rayon = { version = "1.10.0", optional = true }
 atoi_simd = { version = "0.15.6", optional = true }
 arrow2 = "0.18.0"
 libc = "0.2.154"

src/chunked/arrow_converter.rs

@@ -27,7 +27,7 @@
 
 use arrow::array::{ArrayRef, BooleanBuilder, Int32Builder, Int64Builder, StringBuilder};
 use arrow::record_batch::RecordBatch;
-use log::debug;
+use log::{debug, info};
 use parquet::arrow::ArrowWriter;
 use parquet::basic::Compression;
 use parquet::file::properties::WriterProperties;
@@ -46,26 +46,48 @@ use crate::datatype::DataType;
 use crate::schema;
 use crate::trimmer::{trimmer_factory, ColumnTrimmer};
 use crate::{chunked, trimmer};
+use arrow::datatypes::{Field, Schema, SchemaRef};
+use crossbeam::atomic::AtomicConsume;
+use parquet::errors::{ParquetError, Result};
+use parquet::file::metadata::FileMetaData;
+use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
+use std::sync::mpsc::{sync_channel, Receiver, RecvError, SyncSender};
+use std::thread;
+use std::thread::JoinHandle;
 use std::time::{Duration, Instant};
+use thread::spawn;
+//use crossbeam::channel::{Receiver, Sender};
+
+static GLOBAL_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
 
 pub(crate) struct Slice2Arrow<'a> {
     //    pub(crate) file_out: File,
-    pub(crate) writer: ArrowWriter<File>,
     pub(crate) fn_line_break: FnFindLastLineBreak<'a>,
     pub(crate) fn_line_break_len: FnLineBreakLen,
     pub(crate) masterbuilders: MasterBuilders,
 }
 
 pub(crate) struct MasterBuilders {
     builders: Vec<Vec<Box<dyn Sync + Send + ColumnBuilder>>>,
-    //      schema: arrow_schema::SchemaRef
+    outfile: PathBuf,
+    sender: Option<SyncSender<RecordBatch>>,
 }
 
 unsafe impl Send for MasterBuilders {}
 unsafe impl Sync for MasterBuilders {}
 
 impl MasterBuilders {
-    pub fn writer_factory(&mut self, out_file: &PathBuf) -> ArrowWriter<File> {
+    pub fn schema_factory(&mut self) -> SchemaRef {
+        let b: &mut Vec<Box<dyn Sync + Send + ColumnBuilder>> = self.builders.get_mut(0).unwrap();
+        let mut br: Vec<(&str, ArrayRef)> = vec![];
+        for bb in b.iter_mut() {
+            br.push(bb.finish());
+        }
+
+        let batch = RecordBatch::try_from_iter(br).unwrap();
+        batch.schema()
+    }
+    pub fn writer_factory(out_file: &PathBuf, schema: SchemaRef) -> ArrowWriter<File> {
         let _out_file = fs::OpenOptions::new()
             .create(true)
             .append(true)
@@ -76,19 +98,12 @@ impl MasterBuilders {
             .set_compression(Compression::SNAPPY)
             .build();
 
-        let b: &mut Vec<Box<dyn Sync + Send + ColumnBuilder>> = self.builders.get_mut(0).unwrap();
-        let mut br: Vec<(&str, ArrayRef)> = vec![];
-        for bb in b.iter_mut() {
-            br.push(bb.finish());
-        }
-
-        let batch = RecordBatch::try_from_iter(br).unwrap();
         let writer: ArrowWriter<File> =
-            ArrowWriter::try_new(_out_file, batch.schema(), Some(props.clone())).unwrap();
+            ArrowWriter::try_new(_out_file, schema, Some(props.clone())).unwrap();
         writer
     }
 
-    pub fn builders_factory(schema_path: PathBuf, instances: i16) -> Self {
+    pub fn builders_factory(out_file: PathBuf, schema_path: PathBuf, instances: i16) -> Self {
         let schema = schema::FixedSchema::from_path(schema_path).unwrap();
         let antal_col = schema.num_columns();
         let mut builders: Vec<Vec<Box<dyn ColumnBuilder + Sync + Send>>> = Vec::new();
@@ -152,7 +167,11 @@ impl MasterBuilders {
             }
             builders.push(buildersmut);
         }
-        MasterBuilders { builders }
+        MasterBuilders {
+            builders: builders,
+            outfile: out_file,
+            sender: None,
+        }
     }
 }
 
@@ -197,6 +216,8 @@ impl<'a> Converter<'a> for Slice2Arrow<'a> {
         }
         parse_duration = start_parse.elapsed();
 
+        let mut rb: Vec<RecordBatch> = Vec::new();
+
         for b in self.masterbuilders.builders.iter_mut() {
             let mut br: Vec<(&str, ArrayRef)> = vec![];
 
@@ -205,24 +226,72 @@ impl<'a> Converter<'a> for Slice2Arrow<'a> {
             }
 
             let start_builder_write = Instant::now();
-            let batch = RecordBatch::try_from_iter(br).unwrap();
+            let record_batch = RecordBatch::try_from_iter(br).unwrap();
+            //            rb.push(RecordBatch::try_from_iter(br).unwrap());
 
-            self.writer.write(&batch).expect("Error Writing batch");
-            bytes_out += self.writer.bytes_written();
+            let _ = self
+                .masterbuilders
+                .sender
+                .clone()
+                .unwrap()
+                .send(record_batch);
 
             builder_write_duration += start_builder_write.elapsed();
         }
+        //        let writer: ArrowWriter<File> =
+
         debug!("Batch write: accumulated bytes_written {}", bytes_out);
 
         (bytes_in, bytes_out, parse_duration, builder_write_duration)
     }
 
-    fn finish(&mut self) -> parquet::errors::Result<format::FileMetaData> {
-        self.writer.finish()
+    fn setup(&mut self) -> JoinHandle<Result<format::FileMetaData>> {
+        let schema = self.masterbuilders.schema_factory();
+        let _outfile = self.masterbuilders.outfile.clone();
+        let mut writer = crate::chunked::arrow_converter::MasterBuilders::writer_factory(
+            &_outfile,
+            schema.clone(),
+        );
+        let (sender, receiver) = sync_channel::<RecordBatch>(1);
+
+        let t: JoinHandle<Result<format::FileMetaData>> = thread::spawn(move || {
+            'outer: loop {
+                let message = receiver.recv();
+                match message {
+                    Ok(rb) => {
+                        writer.write(&rb).expect("Error Writing batch");
+                        if (rb.num_rows() == 0) {
+                            break 'outer;
+                        }
+                    }
+                    Err(e) => {
+                        info!("got RecvError in channel , break to outer");
+
+                        break 'outer;
+                    }
+                }
+            }
+            info!("closing the writer for parquet");
+
+            writer.finish()
+        });
+        self.masterbuilders.sender = Some(sender.clone());
+        t
+
+        //        bytes_out += crate::chunked::arrow_converter::GLOBAL_COUNTER.load_consume();
     }
 
-    fn get_finish_bytes_written(&mut self) -> usize {
-        self.writer.bytes_written()
+    fn shutdown(&mut self) {
+        //        converter.shutdown();
+        let schema = Schema::new(vec![Field::new(
+            "id",
+            arrow::datatypes::DataType::Int32,
+            false,
+        )]);
+
+        let emptyrb = arrow::record_batch::RecordBatch::new_empty(Arc::new(schema));
+
+        let _ = &self.masterbuilders.sender.clone().unwrap().send(emptyrb);
     }
 }
 

src/chunked/mod.rs

@@ -25,13 +25,16 @@
 // Last updated: 2024-05-15
 //
 
-use arrow::array::ArrayRef;
+use arrow::array::{ArrayRef, RecordBatch};
 use parquet::format;
 
 use std::fs;
+use std::sync::mpsc::SyncSender;
+use std::thread::JoinHandle;
 use std::time::Duration;
 
 use self::residual_slicer::SLICER_IN_CHUNK_SIZE;
+use parquet::errors::{ParquetError, Result};
 
 pub(crate) mod arrow_converter;
 pub(crate) mod residual_slicer;
@@ -152,8 +155,8 @@ pub(crate) trait Converter<'a> {
 
     //    fn process(& mut self, slices: Vec< &'a[u8]>) -> usize;
     fn process(&mut self, slices: Vec<&'a [u8]>) -> (usize, usize, Duration, Duration);
-    fn finish(&mut self) -> parquet::errors::Result<format::FileMetaData>;
-    fn get_finish_bytes_written(&mut self) -> usize;
+    fn setup(&mut self) -> JoinHandle<Result<format::FileMetaData>>;
+    fn shutdown(&mut self);
 }
 
 pub trait ColumnBuilder {

src/chunked/residual_slicer.rs

@@ -27,10 +27,15 @@
 
 use log::info;
 
+use arrow::array::Int32Array;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::ipc::RecordBatch;
 use std::cmp;
 use std::fs;
 use std::fs::File;
 use std::io::{BufReader, Read};
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use super::{ChunkAndResidue, Converter, FnFindLastLineBreak, IterRevolver, Slicer, Stats};
@@ -84,6 +89,8 @@ impl<'a> Slicer<'a> for ResidualSlicer<'a> {
             .build_global()
             .unwrap();
 
+        let j = converter.setup();
+
         loop {
             let cr = ir.next().unwrap();
 
@@ -131,23 +138,25 @@ impl<'a> Slicer<'a> for ResidualSlicer<'a> {
             parse_duration_tot += parse_duration;
             builder_write_duration_tot += builder_write_duration;
         }
+        info!("about to shudown converter...");
+
+        converter.shutdown();
+        info!("converter has been shutdown");
+        let _ = j.join(); // Make sure writer thread is done.
 
         info!(
             "Bytes in= {}\n out= {}\nparse duration= {:?}\n \n builder write_duration {:?}\n",
             bytes_in, bytes_out, parse_duration_tot, builder_write_duration_tot
         );
 
-        match converter.finish() {
-            Ok(x) => Result::Ok(Stats {
-                bytes_in,
-                bytes_out: converter.get_finish_bytes_written(),
-                num_rows: x.num_rows,
-                read_duration: read_duration_tot,
-                parse_duration: parse_duration_tot,
-                builder_write_duration: builder_write_duration_tot,
-            }),
-            Err(_x) => Result::Err("Could not produce Parquet"),
-        }
+        Ok(Stats {
+            bytes_in,
+            bytes_out: bytes_out,
+            num_rows: 0,
+            read_duration: read_duration_tot,
+            parse_duration: parse_duration_tot,
+            builder_write_duration: builder_write_duration_tot,
+        })
     }
 }
 

src/chunked/self_converter.rs

@@ -28,8 +28,14 @@
 use log::info;
 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 
+use arrow::array::RecordBatch;
+use parquet::errors::{ParquetError, Result};
+use parquet::file::metadata::FileMetaData;
+use parquet::format;
 use std::fs::File;
 use std::io::Write;
+use std::sync::mpsc::SyncSender;
+use std::thread::JoinHandle;
 use std::time::Duration;
 
 use super::{Converter, FnFindLastLineBreak};
@@ -63,11 +69,11 @@ impl<'a> Converter<'a> for SampleSliceAggregator<'a> {
         (bytes_processed, 0, duration, duration)
     }
 
-    fn finish(&mut self) -> Result<parquet::format::FileMetaData, parquet::errors::ParquetError> {
+    fn setup(&mut self) -> JoinHandle<Result<format::FileMetaData>> {
         todo!()
     }
 
-    fn get_finish_bytes_written(&mut self) -> usize {
+    fn shutdown(&mut self) {
         todo!()
     }
 }

src/cli.rs

@@ -271,13 +271,12 @@ impl Cli {
                 let converter_instance: Box<dyn ChunkedConverter> = match converter {
                     Converters::Arrow => {
                         let mut master_builders = MasterBuilders::builders_factory(
+                            out_file.clone().to_path_buf(),
                             schema.to_path_buf(),
                             n_threads as i16,
                         );
-                        let writer: ArrowWriter<File> = master_builders.writer_factory(out_file);
 
                         let s2a: Box<Slice2Arrow> = Box::new(Slice2Arrow {
-                            writer,
                             fn_line_break: find_last_nl,
                             fn_line_break_len: line_break_len_cr,
                             masterbuilders: master_builders,