Merge pull request #44 from firelink-data/feature/ipc

Feature/ipc

Author: Ignalina

Cargo.toml

@@ -43,6 +43,8 @@ arrow2 = "0.18.0"
 libc = "0.2.154"
 arrow = "51.0.0"
 parquet = "51.0.0"
+arrow-ipc = "51.0.0"
+
 padder = { version = "1.2.0", features = ["serde"] }
 ordered-channel="1.1.0"
 atomic-counter = "1.0.1"

src/chunked/arrow_converter.rs

@@ -25,6 +25,7 @@
 // Last updated: 2024-05-15
 //
 
+use crate::chunked::trimmer::{trimmer_factory, ColumnTrimmer};
 use arrow::array::{ArrayRef, BooleanBuilder, Int32Builder, Int64Builder, StringBuilder};
 use arrow::record_batch::RecordBatch;
 use log::{debug, info};
@@ -38,22 +39,28 @@ use rayon::prelude::*;
 
 use std::fs;
 use std::fs::File;
+use std::ops::Deref;
 use std::path::PathBuf;
 use std::str::from_utf8_unchecked;
 use std::sync::Arc;
 
-use super::{ColumnBuilder, Converter, FnFindLastLineBreak, FnLineBreakLen};
+use super::{
+    arrow_file_output, trimmer, ColumnBuilder, Converter, FnFindLastLineBreak, FnLineBreakLen,
+    Stats,
+};
+use crate::chunked;
+use crate::chunked::threaded_file_output::{ipc_file_out, output_factory, parquet_file_out};
+pub use crate::cli::Targets;
 use crate::datatype::DataType;
 use crate::schema;
-use crate::trimmer::{trimmer_factory, ColumnTrimmer};
-use crate::{chunked, trimmer};
 use arrow::datatypes::{Field, Schema, SchemaRef};
 use atomic_counter::{AtomicCounter, ConsistentCounter};
 use crossbeam::atomic::AtomicConsume;
 use libc::bsearch;
 use ordered_channel::Sender;
 use parquet::errors::{ParquetError, Result};
 use parquet::file::metadata::FileMetaData;
+use rayon::join;
 use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
 use std::sync::mpsc::{sync_channel, Receiver, RecvError, SyncSender};
 use std::thread;
@@ -71,6 +78,7 @@ pub(crate) struct Slice2Arrow<'a> {
     pub(crate) fn_line_break_len: FnLineBreakLen,
     pub(crate) masterbuilders: MasterBuilders,
     pub(crate) consistent_counter: ConsistentCounter,
+    pub(crate) target: Targets,
 }
 
 pub(crate) struct MasterBuilders {
@@ -93,21 +101,6 @@ impl MasterBuilders {
         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)
-            .open(out_file)
-            .expect("aaa");
-
-        let props = WriterProperties::builder()
-            .set_compression(Compression::SNAPPY)
-            .build();
-
-        let writer: ArrowWriter<File> =
-            ArrowWriter::try_new(_out_file, schema, Some(props.clone())).unwrap();
-        writer
-    }
 
     pub fn builders_factory(out_file: PathBuf, schema_path: PathBuf, instances: i16) -> Self {
         let schema = schema::FixedSchema::from_path(schema_path).unwrap();
@@ -242,40 +235,14 @@ impl<'a> Converter<'a> for Slice2Arrow<'a> {
         (bytes_in, bytes_out, parse_duration, builder_write_duration)
     }
 
-    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(),
+    fn setup(&mut self) -> (Sender<RecordBatch>, JoinHandle<Result<Stats>>) {
+        let o = output_factory(
+            self.target.clone(),
+            self.masterbuilders.schema_factory(),
+            self.masterbuilders.outfile.clone(),
         );
-
-        let (sender, mut receiver) = bounded::<RecordBatch>(100);
-
-        let t: JoinHandle<Result<format::FileMetaData>> = thread::spawn(move || {
-            'outer: loop {
-                let mut 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();
+        self.masterbuilders.sender = Some(o.0.clone());
+        o
     }
 
     fn shutdown(&mut self) {

src/chunked/mod.rs

@@ -27,8 +27,12 @@
 
 use arrow::array::{ArrayRef, RecordBatch};
 use parquet::format;
+use std::fmt::Debug;
 
+use arrow::datatypes::SchemaRef;
+use ordered_channel::Sender;
 use std::fs;
+use std::path::PathBuf;
 use std::sync::mpsc::SyncSender;
 use std::thread::JoinHandle;
 use std::time::Duration;
@@ -39,6 +43,8 @@ use parquet::errors::{ParquetError, Result};
 pub(crate) mod arrow_converter;
 pub(crate) mod residual_slicer;
 pub(crate) mod self_converter;
+pub(crate) mod threaded_file_output;
+mod trimmer;
 
 pub(crate) struct ChunkAndResidue {
     pub(crate) chunk: Box<[u8; SLICER_IN_CHUNK_SIZE]>,
@@ -155,7 +161,7 @@ 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 setup(&mut self) -> JoinHandle<Result<format::FileMetaData>>;
+    fn setup(&mut self) -> (Sender<RecordBatch>, JoinHandle<Result<Stats>>);
     fn shutdown(&mut self);
 }
 
@@ -164,3 +170,11 @@ pub trait ColumnBuilder {
     fn finish(&mut self) -> (&str, ArrayRef);
     //    fn name(&  self) -> &String;
 }
+
+pub trait arrow_file_output {
+    fn setup(
+        &mut self,
+        schema: SchemaRef,
+        outfile: PathBuf,
+    ) -> (Sender<RecordBatch>, JoinHandle<Result<Stats>>);
+}

src/chunked/residual_slicer.rs

@@ -40,7 +40,7 @@ use std::time::{Duration, Instant};
 
 use super::{ChunkAndResidue, Converter, FnFindLastLineBreak, IterRevolver, Slicer, Stats};
 
-pub(crate) const SLICER_IN_CHUNK_SIZE: usize = 1024 * 2024;
+pub(crate) const SLICER_IN_CHUNK_SIZE: usize = 1024 * 2048;
 pub(crate) const SLICER_MAX_RESIDUE_SIZE: usize = SLICER_IN_CHUNK_SIZE;
 
 pub(crate) const IN_MAX_CHUNKS: usize = 2;
@@ -89,7 +89,7 @@ impl<'a> Slicer<'a> for ResidualSlicer<'a> {
             .build_global()
             .unwrap();
 
-        let j = converter.setup();
+        let threaded_writer = converter.setup();
 
         loop {
             let cr = ir.next().unwrap();
@@ -141,8 +141,8 @@ impl<'a> Slicer<'a> for ResidualSlicer<'a> {
         info!("about to shudown converter...");
 
         converter.shutdown();
+        threaded_writer.1.join();
         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",

src/chunked/self_converter.rs

@@ -25,10 +25,12 @@
 // Last updated: 2024-05-15
 //
 
+use crate::chunked::Stats;
 use log::info;
 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 
 use arrow::array::RecordBatch;
+use ordered_channel::Sender;
 use parquet::errors::{ParquetError, Result};
 use parquet::file::metadata::FileMetaData;
 use parquet::format;
@@ -69,7 +71,7 @@ impl<'a> Converter<'a> for SampleSliceAggregator<'a> {
         (bytes_processed, 0, duration, duration)
     }
 
-    fn setup(&mut self) -> JoinHandle<Result<format::FileMetaData>> {
+    fn setup(&mut self) -> (Sender<RecordBatch>, JoinHandle<Result<Stats>>) {
         todo!()
     }