refactor: Make --part and --parts Option, better error handling

Author: alamb

tpchgen-cli/src/main.rs

@@ -90,13 +90,15 @@ struct Cli {
     #[arg(short = 'T', long = "tables", value_delimiter = ',', value_parser = TableValueParser)]
     tables: Option<Vec<Table>>,
 
-    /// Number of parts to generate (manual parallel generation)
-    #[arg(short, long, default_value_t = -1)]
-    parts: i32,
+    /// Number of partitions to generate (manual parallel generation)
+    #[arg(short, long)]
+    parts: Option<i32>,
 
-    /// Which part to generate (1-based, only relevant if parts > 1)
-    #[arg(long, default_value_t = -1)]
-    part: i32,
+    /// Which partition to generate (1-based)
+    ///
+    /// If not specified, generates all parts
+    #[arg(long)]
+    part: Option<i32>,
 
     /// Output format: tbl, csv, parquet (default: tbl)
     #[arg(short, long, default_value = "tbl")]
@@ -254,14 +256,15 @@ macro_rules! define_generate {
     ($FUN_NAME:ident,  $TABLE:expr, $GENERATOR:ident, $TBL_SOURCE:ty, $CSV_SOURCE:ty, $PARQUET_SOURCE:ty) => {
         async fn $FUN_NAME(&self) -> io::Result<()> {
             let filename = self.output_filename($TABLE);
-            let plan = GenerationPlan::new(
+            let plan = GenerationPlan::try_new(
                 &$TABLE,
                 self.format,
                 self.scale_factor,
                 self.part,
                 self.parts,
                 self.num_threads,
-            );
+            )
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
             let scale_factor = self.scale_factor;
             info!("Writing table {} (SF={scale_factor}) to {filename}", $TABLE);
             debug!("Plan: {plan}");

tpchgen-cli/src/plan.rs

@@ -60,59 +60,98 @@ impl GenerationPlan {
     /// Returns the number of parts to generate
     ///
     /// cli_part and cli_part_count are passed from the CLI arguments
-    /// to specify a particular part or number of parts to generate.
-    pub fn new(
+    /// to specify a particular part or number of partitions to generate.
+    pub fn try_new(
         table: &Table,
         format: OutputFormat,
         scale_factor: f64,
-        cli_part: i32,
-        cli_part_count: i32,
+        cli_part: Option<i32>,
+        cli_part_count: Option<i32>,
         num_threads: usize,
-    ) -> Self {
+    ) -> Result<Self, String> {
         // If a single part is specified, split it into chunks to enable parallel generation.
-        if cli_part != -1 || cli_part_count != -1 {
-            // These tables are small not parameterized by part count,
-            // so we must create only a single part.
-            if table == &Table::Nation || table == &Table::Region {
-                return Self {
-                    part_count: 1,
-                    part_list: vec![1],
-                };
+        match (cli_part, cli_part_count) {
+            (Some(_part), None) => Err(String::from(
+                "The --part option requires the --parts option to be set",
+            )),
+            (None, Some(_part_count)) => {
+                // TODO automatically create multiple files if part_count > 1
+                // and part is not specified
+                Err(String::from(
+                    "The --part_count option requires the --part option to be set",
+                ))
             }
-
-            // sanity check arguments (TODO: real Errors)
-            if cli_part < 1 || cli_part_count < 1 || cli_part > cli_part_count {
-                panic!(
-                    "Invalid CLI part or part count. \
-                    Expect greater than 1 and cli_part <= cli_part_count. \
-                    Got: cli_part={cli_part}, cli_part_count={cli_part_count}",
-                );
+            (Some(part), Some(part_count)) => {
+                Self::try_new_with_parts(table, format, scale_factor, part, part_count, num_threads)
             }
+            (None, None) => Self::try_new_without_parts(table, format, scale_factor),
+        }
+    }
+
+    /// Returns a new `GenerationPlan` when partitioning is specified on the command line
+    fn try_new_with_parts(
+        table: &Table,
+        // todo take into account the output format (e.g. parquet row group size)
+        _format: OutputFormat,
+        scale_factor: f64,
+        cli_part: i32,
+        cli_part_count: i32,
+        num_threads: usize,
+    ) -> Result<Self, String> {
+        if cli_part < 1 {
+            return Err(format!(
+                "Invalid --part. Expected a number greater than zero, got {cli_part}"
+            ));
+        }
+        if cli_part_count < 1 {
+            return Err(format!(
+                "Invalid --part_count. Expected a number greater than zero, got {cli_part_count}"
+            ));
+        }
+        if cli_part > cli_part_count {
+            return Err(format!(
+                    "Invalid --part. Expected at most the value of --parts ({cli_part_count}), got {cli_part}"));
+        }
 
-            let num_chunks = num_threads as i32;
-
-            // The new total number of parts is the original number of parts multiplied by the number of chunks.
-            let new_total_parts = cli_part_count * num_chunks;
-
-            // The new part numbers to generate are the chunks that make up the original part.
-            let start_part = (cli_part - 1) * num_chunks + 1;
-            let end_part = cli_part * num_chunks;
-            let new_parts_to_generate = (start_part..=end_part).collect();
-            debug!(
-                "Generating {} parts for table {:?} with scale factor {}",
-                new_total_parts, table, scale_factor
-            );
-            debug!(
-                "CLI part: {}, CLI part count: {}, num_threads: {}",
-                cli_part, cli_part_count, num_threads
-            );
-            debug!("New parts to generate: {:?}", new_parts_to_generate);
-            return Self {
-                part_count: new_total_parts,
-                part_list: new_parts_to_generate,
-            };
+        // These tables are small they are not parameterized by part count, so
+        // we must create only a single part.
+        if table == &Table::Nation || table == &Table::Region {
+            return Ok(Self {
+                part_count: 1,
+                part_list: vec![1],
+            });
         }
 
+        let num_chunks = num_threads as i32;
+
+        // The new total number of parts is the original number of parts multiplied by the number of chunks.
+        let new_total_parts = cli_part_count * num_chunks;
+
+        // The new part numbers to generate are the chunks that make up the original part.
+        let start_part = (cli_part - 1) * num_chunks + 1;
+        let end_part = cli_part * num_chunks;
+        let new_parts_to_generate = (start_part..=end_part).collect();
+        debug!(
+            "Generating {} parts for table {:?} with scale factor {}",
+            new_total_parts, table, scale_factor
+        );
+        debug!(
+            "CLI part: {}, CLI part count: {}, num_threads: {}",
+            cli_part, cli_part_count, num_threads
+        );
+        debug!("New parts to generate: {:?}", new_parts_to_generate);
+        Ok(Self {
+            part_count: new_total_parts,
+            part_list: new_parts_to_generate,
+        })
+    }
+
+    /// Returns a new `GenerationPlan` when no partitioning is specified on the command line
+    fn try_new_without_parts(
+        table: &Table,
+        format: OutputFormat,
+        scale_factor: f64,
+    ) -> Result<Self, String> {
         // Note use part=1, part_count=1 to calculate the total row count
         // for the table
         //
@@ -157,10 +196,10 @@ impl GenerationPlan {
         let num_parts = num_parts.try_into().unwrap();
         // generating all the parts
 
-        Self {
+        Ok(Self {
             part_count: num_parts,
             part_list: (1..=num_parts).collect(),
-        }
+        })
     }
 }
 
@@ -324,31 +363,36 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(
-        expected = "Invalid CLI part or part count. Expect greater than 1 and cli_part <= cli_part_count. Got: cli_part=0, cli_part_count=10"
-    )]
-    fn sf1_lineitem_cli_parts_invalid_small() {
+    fn sf1_lineitem_cli_invalid_part() {
         Test::new()
             .with_table(Table::Lineitem)
             .with_format(OutputFormat::Tbl)
             .with_scale_factor(1.0)
             .with_cli_part(0) // part 0 of 10 (invalid)
             .with_cli_part_count(10)
-            .assert(40, [13, 14, 15, 16])
+            .assert_err("Invalid --part. Expected a number greater than zero, got 0")
     }
 
     #[test]
-    #[should_panic(
-        expected = "Invalid CLI part or part count. Expect greater than 1 and cli_part <= cli_part_count. Got: cli_part=11, cli_part_count=10"
-    )]
     fn sf1_lineitem_cli_parts_invalid_big() {
         Test::new()
             .with_table(Table::Lineitem)
             .with_format(OutputFormat::Tbl)
             .with_scale_factor(1.0)
             .with_cli_part(11) // part 11 of 10 (invalid)
             .with_cli_part_count(10)
-            .assert(40, [13, 14, 15, 16])
+            .assert_err("Invalid --part. Expected at most the value of --parts (10), got 11");
+    }
+
+    #[test]
+    fn sf1_lineitem_cli_invalid_part_count() {
+        Test::new()
+            .with_table(Table::Lineitem)
+            .with_format(OutputFormat::Tbl)
+            .with_scale_factor(1.0)
+            .with_cli_part(1) // part 0 of 0 (invalid)
+            .with_cli_part_count(0)
+            .assert_err("Invalid --part_count. Expected a number greater than zero, got 0");
     }
 
     #[test]
@@ -393,8 +437,8 @@ mod tests {
         table: Table,
         format: OutputFormat,
         scale_factor: f64,
-        cli_part: i32,
-        cli_part_count: i32,
+        cli_part: Option<i32>,
+        cli_part_count: Option<i32>,
         num_cpus: usize,
     }
 
@@ -410,19 +454,34 @@ mod tests {
             expected_part_count: i32,
             expected_part_numbers: impl IntoIterator<Item = i32>,
         ) {
-            let plan = GenerationPlan::new(
+            let plan = GenerationPlan::try_new(
                 &self.table,
                 self.format,
                 self.scale_factor,
                 self.cli_part,
                 self.cli_part_count,
                 self.num_cpus,
-            );
+            )
+            .unwrap();
             assert_eq!(plan.part_count, expected_part_count);
             let expected_part_numbers: Vec<i32> = expected_part_numbers.into_iter().collect();
             assert_eq!(plan.part_list, expected_part_numbers);
         }
 
+        /// Assert that creating a [`GenerationPlan`] returns the specified error
+        fn assert_err(self, expected_error: &str) {
+            let actual_error = GenerationPlan::try_new(
+                &self.table,
+                self.format,
+                self.scale_factor,
+                self.cli_part,
+                self.cli_part_count,
+                self.num_cpus,
+            )
+            .unwrap_err();
+            assert_eq!(actual_error, expected_error);
+        }
+
         /// Set table
         fn with_table(mut self, table: Table) -> Self {
             self.table = table;
@@ -443,13 +502,13 @@ mod tests {
 
         /// Set CLI part
         fn with_cli_part(mut self, cli_part: i32) -> Self {
-            self.cli_part = cli_part;
+            self.cli_part = Some(cli_part);
             self
         }
 
         /// Set CLI part count
         fn with_cli_part_count(mut self, cli_part_count: i32) -> Self {
-            self.cli_part_count = cli_part_count;
+            self.cli_part_count = Some(cli_part_count);
             self
         }
     }
@@ -460,8 +519,8 @@ mod tests {
                 table: Table::Orders,
                 format: OutputFormat::Tbl,
                 scale_factor: 1.0,
-                cli_part: -1,
-                cli_part_count: -1,
+                cli_part: None,
+                cli_part_count: None,
                 num_cpus: 4, // hard code 4 cores for testing
             }
         }