chore: Remove num partitions from repartitioner (apache#1498)

* chore: Remove num partitions from repartitioner

treat all single partition schemes the same way
remove duplicate messaging by using assert_eq instead of assert

* rebased and resolved conflicts

---------

Co-authored-by: Emily Matheys <[email protected]>

Author: EmilyMatt

native/core/src/execution/shuffle/shuffle_writer.rs

@@ -290,7 +290,6 @@ struct ShuffleRepartitioner {
     buffered_partitions: Vec<PartitionBuffer>,
     /// Partitioning scheme to use
     partitioning: Partitioning,
-    num_output_partitions: usize,
     runtime: Arc<RuntimeEnv>,
     metrics: ShuffleRepartitionerMetrics,
     /// Hashes for each row in the current batch
@@ -315,8 +314,6 @@ impl ShuffleRepartitioner {
         codec: CompressionCodec,
         enable_fast_encoding: bool,
     ) -> Result<Self> {
-        let num_output_partitions = partitioning.partition_count();
-
         let mut hashes_buf = Vec::with_capacity(batch_size);
         let mut partition_ids = Vec::with_capacity(batch_size);
 
@@ -336,7 +333,7 @@ impl ShuffleRepartitioner {
             output_data_file,
             output_index_file,
             schema: Arc::clone(&schema),
-            buffered_partitions: (0..num_output_partitions)
+            buffered_partitions: (0..partitioning.partition_count())
                 .map(|_| {
                     PartitionBuffer::try_new(
                         Arc::clone(&schema),
@@ -348,7 +345,6 @@ impl ShuffleRepartitioner {
                 })
                 .collect::<Result<Vec<_>>>()?,
             partitioning,
-            num_output_partitions,
             runtime,
             metrics,
             hashes_buf,
@@ -401,9 +397,21 @@ impl ShuffleRepartitioner {
         // number of rows those are written to output data file.
         self.metrics.baseline.record_output(input.num_rows());
 
-        let num_output_partitions = self.num_output_partitions;
         match &self.partitioning {
-            Partitioning::Hash(exprs, _) => {
+            any if any.partition_count() == 1 => {
+                let buffered_partitions = &mut self.buffered_partitions;
+
+                assert_eq!(buffered_partitions.len(), 1, "Expected 1 partition");
+
+                // TODO the single partition case could be optimized to avoid appending all
+                // rows from the batch into builders and then recreating the batch
+                // https://github.com/apache/datafusion-comet/issues/1453
+                let indices = (0..input.num_rows()).collect::<Vec<usize>>();
+
+                self.append_rows_to_partition(input.columns(), &indices, 0)
+                    .await?;
+            }
+            Partitioning::Hash(exprs, num_output_partitions) => {
                 let (partition_starts, shuffled_partition_ids): (Vec<usize>, Vec<usize>) = {
                     let mut timer = self.metrics.repart_time.timer();
 
@@ -423,11 +431,11 @@ impl ShuffleRepartitioner {
                         .iter()
                         .enumerate()
                         .for_each(|(idx, hash)| {
-                            partition_ids[idx] = pmod(*hash, num_output_partitions) as u64
+                            partition_ids[idx] = pmod(*hash, *num_output_partitions) as u64
                         });
 
                     // count each partition size
-                    let mut partition_counters = vec![0usize; num_output_partitions];
+                    let mut partition_counters = vec![0usize; *num_output_partitions];
                     partition_ids
                         .iter()
                         .for_each(|partition_id| partition_counters[*partition_id as usize] += 1);
@@ -478,24 +486,6 @@ impl ShuffleRepartitioner {
                     .await?;
                 }
             }
-            Partitioning::UnknownPartitioning(n) if *n == 1 => {
-                let buffered_partitions = &mut self.buffered_partitions;
-
-                assert_eq!(
-                    buffered_partitions.len(),
-                    1,
-                    "Expected 1 partition but got {}",
-                    buffered_partitions.len()
-                );
-
-                // TODO the single partition case could be optimized to avoid appending all
-                // rows from the batch into builders and then recreating the batch
-                // https://github.com/apache/datafusion-comet/issues/1453
-                let indices = (0..input.num_rows()).collect::<Vec<usize>>();
-
-                self.append_rows_to_partition(input.columns(), &indices, 0)
-                    .await?;
-            }
             other => {
                 // this should be unreachable as long as the validation logic
                 // in the constructor is kept up-to-date
@@ -511,8 +501,8 @@ impl ShuffleRepartitioner {
     /// Writes buffered shuffled record batches into Arrow IPC bytes.
     async fn shuffle_write(&mut self) -> Result<SendableRecordBatchStream> {
         let mut elapsed_compute = self.metrics.baseline.elapsed_compute().timer();
-        let num_output_partitions = self.num_output_partitions;
         let buffered_partitions = &mut self.buffered_partitions;
+        let num_output_partitions = buffered_partitions.len();
         let mut output_batches: Vec<Vec<u8>> = vec![vec![]; num_output_partitions];
         let mut offsets = vec![0; num_output_partitions + 1];
         for i in 0..num_output_partitions {