devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1316291629
##########
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:
@metesynnada Actually, I just realized that the fifo test passing is
dependent on leaving the `yield_now().await` call in. That call results in the
serialize tasks being run in sequence rather than parallel though, so
performance is poor. I'm not quite sure how to get the fifio test to pass
without the `yield_now().await` call.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
