@@ -38,9 +38,12 @@ use crate::slicers::old_slicer::{OldSlicer, IN_MAX_CHUNKS};
3838use crate :: slicers:: Slicer ;
3939use crate :: slicers:: { find_last_nl, line_break_len_cr, ChunkAndResidue } ;
4040use crate :: { error, mocker, schema} ;
41- use clap:: { Parser , Subcommand } ;
41+ use clap:: { value_parser , ArgAction , Parser , Subcommand } ;
4242use log:: info;
4343use parquet:: arrow:: ArrowWriter ;
44+ use crate :: mocker:: Mocker ;
45+ use crate :: error:: Result ;
46+ use crate :: threads:: get_available_threads;
4447
4548#[ derive( Parser ) ]
4649#[ command( author, version, about, long_about = None ) ]
@@ -77,13 +80,32 @@ enum Commands {
7780 /// does testing things
7881 Mock {
7982 /// Sets schema file
80- #[ arg( short, long, value_name = "SCHEMA" ) ]
83+ #[ arg( short = 's' , long= "schema" , value_name = "SCHEMA" , required = true ) ]
8184 schema : PathBuf ,
8285 /// Sets input file
83- #[ arg( short, long, value_name = "FILE" ) ]
84- file : Option < PathBuf > ,
85- #[ arg( short, long, value_name = "n-rows" , default_value = "100" ) ]
86- n_rows : Option < i64 > ,
86+ #[ arg( short='o' , long="output-file" , value_name = "OUTPUT-FILE" , required = false ) ]
87+ output_file : Option < PathBuf > ,
88+ #[ arg( short='n' , long="n-rows" , value_name = "NUM-ROWS" , default_value = "100" , required = false ) ]
89+ n_rows : Option < usize > ,
90+
91+ /// Set the size of the buffer (number of rows).
92+ #[ arg(
93+ long = "buffer-size" ,
94+ value_name = "BUFFER-SIZE" ,
95+ action = ArgAction :: Set ,
96+ required = false ,
97+ ) ]
98+ buffer_size : Option < usize > ,
99+
100+ /// Set the capacity of the thread channel (number of messages).
101+ #[ arg(
102+ long = "thread-channel-capacity" ,
103+ value_name = "THREAD-CHANNEL-CAPACITY" ,
104+ action = ArgAction :: Set ,
105+ required = false ,
106+ ) ]
107+ thread_channel_capacity : Option < usize > ,
108+
87109 } ,
88110 Convert {
89111 /// Sets schema file
@@ -119,7 +141,7 @@ impl Cli {
119141 pub fn run < ' a > (
120142 & self ,
121143 in_buffers : & mut [ ChunkAndResidue ; IN_MAX_CHUNKS ] ,
122- ) -> Result < ( ) , error :: ExecutionError > {
144+ ) -> Result < ( ) > {
123145 let n_logical_threads = num_cpus:: get ( ) ;
124146 let mut n_threads: usize = self . n_threads as usize ;
125147
@@ -139,17 +161,23 @@ impl Cli {
139161 match & self . command {
140162 Some ( Commands :: Mock {
141163 schema,
142- file ,
164+ output_file ,
143165 n_rows,
144- } ) => {
145- print ! ( "target file {:?}" , file. as_ref( ) ) ;
146-
147- mocker:: Mocker :: new (
148- schema:: FixedSchema :: from_path ( schema. into ( ) ) ,
149- file. clone ( ) ,
150- n_threads,
151- )
152- . generate ( n_rows. unwrap ( ) as usize ) ;
166+ buffer_size,
167+ thread_channel_capacity,
168+ } ) => {
169+ print ! ( "target file {:?}" , output_file. as_ref( ) ) ;
170+
171+ Mocker :: builder ( )
172+ . schema ( schema. to_owned ( ) )
173+ . output_file ( output_file. to_owned ( ) )
174+ . num_rows ( * n_rows)
175+ . num_threads ( n_threads)
176+ . buffer_size ( * buffer_size)
177+ . thread_channel_capacity ( * thread_channel_capacity)
178+ . build ( ) ?
179+ . generate ( ) ;
180+
153181 Ok ( ( ) )
154182 }
155183
0 commit comments