devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1316129515
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,13 +299,90 @@ pub(crate) async fn create_writer(
}
}
+/// Serializes a single data stream in parallel and writes to an ObjectStore
+/// concurrently. Data order is preserved. In the event of an error,
+/// the ObjectStore writer is returned to the caller in addition to an error,
+/// so that the caller may handle aborting failed writes.
+async fn serialize_rb_stream_to_object_store(
+ mut data_stream: Pin<Box<dyn RecordBatchStream + Send>>,
+ mut serializer: Box<dyn BatchSerializer>,
+ mut writer: AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let mut row_count = 0;
+ // Not using JoinSet here since we want to ulimately write to ObjectStore
preserving file order
+ let mut serialize_tasks: Vec<JoinHandle<Result<(usize, Bytes),
DataFusionError>>> =
+ Vec::new();
+ while let Some(maybe_batch) = data_stream.next().await {
+ let mut serializer_clone = match serializer.duplicate() {
+ Ok(s) => s,
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ serialize_tasks.push(task::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ }));
+ }
+ for serialize_result in serialize_tasks {
+ let result = serialize_result.await;
+ match result {
+ Ok(res) => {
+ let (cnt, bytes) = match res {
+ Ok(r) => r,
+ Err(e) => return Err((writer, e)),
+ };
+ row_count += cnt;
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ }
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ }
+ }
+
+ Ok((serializer, writer, row_count as u64))
+}
Review Comment:
This is fantastic, thank you @metesynnada ! I tested this and with just a
small change it resulted in identical performance with much lower memory usage
due to the backpressure introduced by the channel. The FIFO tests is also
passing now, so win win!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
