Merge pull request #42 from JanKaul/partition-record-batches

avoid collecting the streams of partition streams

iceberg-rust/src/arrow/partition.rs

@@ -42,10 +42,12 @@ pub async fn partition_record_batches(
     partition_spec: &PartitionSpec,
     schema: &Schema,
 ) -> Result<
-    Vec<(
-        Vec<Value>,
-        impl Stream<Item = Result<RecordBatch, ArrowError>> + Send,
-    )>,
+    impl Stream<
+        Item = (
+            Vec<Value>,
+            impl Stream<Item = Result<RecordBatch, ArrowError>> + Send,
+        ),
+    >,
     ArrowError,
 > {
     #[allow(clippy::type_complexity)]
@@ -166,8 +168,7 @@ pub async fn partition_record_batches(
         .into_values()
         .for_each(|sender| sender.close_channel());
     partition_sender.close_channel();
-    let receivers = partition_receiver.collect().await;
-    Ok(receivers)
+    Ok(partition_receiver)
 }
 
 fn distinct_values(array: ArrayRef) -> Result<DistinctValues, ArrowError> {
@@ -310,7 +311,7 @@ mod tests {
         let streams = partition_record_batches(record_batches, &partition_spec, &schema)
             .await
             .unwrap();
-        let output = stream::iter(streams.into_iter())
+        let output = streams
             .then(|s| async move { s.1.collect::<Vec<_>>().await })
             .collect::<Vec<_>>()
             .await;

iceberg-rust/src/arrow/write.rs

@@ -5,7 +5,7 @@
 use futures::{
     channel::mpsc::{channel, unbounded, Receiver, Sender},
     lock::Mutex,
-    stream, SinkExt, StreamExt, TryStreamExt,
+    SinkExt, StreamExt, TryStreamExt,
 };
 use object_store::{buffered::BufWriter, ObjectStore};
 use std::sync::{
@@ -66,7 +66,7 @@ pub async fn write_parquet_partitioned(
     } else {
         let streams = partition_record_batches(batches, partition_spec, schema).await?;
 
-        stream::iter(streams.into_iter())
+        streams
             .map(Ok::<_, ArrowError>)
             .try_for_each_concurrent(None, |(partition_values, batches)| {
                 let arrow_schema = arrow_schema.clone();