split: stream data using copy to not overwhelm memory with large files

src/uu/split/src/split.rs

@@ -1147,49 +1147,32 @@ where
 
     for i in 1_u64..=num_chunks {
         let chunk_size = chunk_size_base + (chunk_size_reminder > i - 1) as u64;
-        let buf = &mut Vec::new();
         if num_bytes > 0 {
-            // Read `chunk_size` bytes from the reader into `buf`
+            // Read `chunk_size` bytes from the reader
             // except the last.
             //
             // The last chunk gets all remaining bytes so that if the number
             // of bytes in the input file was not evenly divisible by
             // `num_chunks`, we don't leave any bytes behind.
-            let limit = {
-                if i == num_chunks {
-                    num_bytes
-                } else {
-                    chunk_size
-                }
+            let limit = if i == num_chunks {
+                num_bytes
+            } else {
+                chunk_size
             };
 
-            let n_bytes_read = reader.by_ref().take(limit).read_to_end(buf);
-
-            match n_bytes_read {
-                Ok(n_bytes) => {
-                    num_bytes -= n_bytes as u64;
-                }
-                Err(error) => {
-                    return Err(USimpleError::new(
-                        1,
-                        translate!("split-error-cannot-read-from-input", "input" => settings.input.maybe_quote(), "error" => error),
-                    ));
-                }
-            }
-
-            match kth_chunk {
-                Some(chunk_number) => {
-                    if i == chunk_number {
-                        stdout_writer.write_all(buf)?;
-                        break;
-                    }
+            let n_bytes = match kth_chunk {
+                Some(chunk_number) if i == chunk_number => {
+                    io::copy(&mut reader.by_ref().take(limit), &mut stdout_writer)?;
+                    break;
                 }
+                Some(_) => io::copy(&mut reader.by_ref().take(limit), &mut io::sink())?,
                 None => {
                     let idx = (i - 1) as usize;
                     let writer = out_files.get_writer(idx, settings)?;
-                    writer.write_all(buf)?;
+                    io::copy(&mut reader.by_ref().take(limit), writer)?
                 }
-            }
+            };
+            num_bytes -= n_bytes;
         } else {
             break;
         }

tests/by-util/test_split.rs

@@ -1988,6 +1988,23 @@ fn test_split_separator_same_multiple() {
         .fails();
 }
 
+#[test]
+#[cfg(target_os = "linux")]
+fn test_number_n_chunks_streaming() {
+    let (at, mut ucmd) = at_and_ucmd!();
+
+    // 100MB file, 100MB memory limit, split into 2x50MB chunks
+    let mut f = std::fs::File::create(at.plus("hundred_mb.bin")).unwrap();
+    f.write_all(&vec![0u8; 100 * 1024 * 1024]).unwrap();
+
+    ucmd.args(&["--number=2", "hundred_mb.bin"])
+        .limit(Resource::AS, 100 * 1024 * 1024, 100 * 1024 * 1024)
+        .succeeds();
+
+    assert_eq!(at.metadata("xaa").len(), 50 * 1024 * 1024);
+    assert_eq!(at.metadata("xab").len(), 50 * 1024 * 1024);
+}
+
 #[test]
 fn test_long_lines() {
     let (at, mut ucmd) = at_and_ucmd!();