feat: zero copy on split rows (#3407)

GreptimeTeam · Feb 28, 2024 · a3533c4 · a3533c4
1 parent 3413fc0
commit a3533c4
Showing 1 changed file with 30 additions and 15 deletions.
diff --git a/src/partition/src/splitter.rs b/src/partition/src/splitter.rs
@@ -79,21 +79,36 @@ impl<'a> SplitReadRowHelper<'a> {
     }
 
     fn split_rows(mut self) -> Result<HashMap<RegionNumber, Rows>> {
-        let request_splits = self
-            .split_to_regions()?
-            .into_iter()
-            .map(|(region_number, row_indexes)| {
-                let rows = row_indexes
-                    .into_iter()
-                    .map(|row_idx| std::mem::take(&mut self.rows[row_idx]))
-                    .collect();
-                let rows = Rows {
-                    schema: self.schema.clone(),
-                    rows,
-                };
-                (region_number, rows)
-            })
-            .collect::<HashMap<_, _>>();
+        let regions = self.split_to_regions()?;
+        let request_splits = if regions.len() == 1 {
+            // fast path, zero copy
+            regions
+                .into_keys()
+                .map(|region_number| {
+                    let rows = std::mem::take(&mut self.rows);
+                    let rows = Rows {
+                        schema: self.schema.clone(),
+                        rows,
+                    };
+                    (region_number, rows)
+                })
+                .collect::<HashMap<_, _>>()
+        } else {
+            regions
+                .into_iter()
+                .map(|(region_number, row_indexes)| {
+                    let rows = row_indexes
+                        .into_iter()
+                        .map(|row_idx| std::mem::take(&mut self.rows[row_idx]))
+                        .collect();
+                    let rows = Rows {
+                        schema: self.schema.clone(),
+                        rows,
+                    };
+                    (region_number, rows)
+                })
+                .collect::<HashMap<_, _>>()
+        };
 
         Ok(request_splits)
     }