You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "devinjdangelo (via GitHub)" <gi...@apache.org> on 2023/08/30 19:40:40 UTC
[GitHub] [arrow-datafusion] devinjdangelo opened a new pull request, #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
devinjdangelo opened a new pull request, #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452
## Which issue does this PR close?
Part of #7079
## Rationale for this change
Serialization of "stateless" file types (where the serialized bytes of each record batch has no dependency on the serialized bytes of any other record batch), can be parallelized efficiently across all available CPU cores for a significant decrease in the time needed to write out the file. This PR uses tokio task spawning to parallelize file serialization.
There is likely a tradeoff between write speed and memory utilization. If the ObjectStore writer cannot keep up with the data being serialized, bytes could accumulate in memory. ObjectStore puts are concurrent but not parallelized so the risk of higher memory usage increases as the number of cores in the system increases.
## What changes are included in this PR?
Spawn a tokio task to serialize each each record batch.
TODOs
- [ ] Test error handling after these changes, ensure writes are still atomic
- [ ] Benchmark write speed and memory consumption for CSV/JSON in this PR vs main
## Are these changes tested?
Yes by existing tests, but more tests needed to verify abort behavior and performance
## Are there any user-facing changes?
No, just faster writes!
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1311919574
##########
datafusion/core/tests/fifo.rs:
##########
@@ -336,6 +336,7 @@ mod unix_test {
/// It tests the INSERT INTO functionality.
#[tokio::test]
+ #[ignore]
Review Comment:
@mustafasrepo Any idea how to fix this test for this PR?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1318861453
##########
datafusion/core/src/datasource/listing/table.rs:
##########
@@ -915,6 +915,7 @@ impl TableProvider for ListingTable {
output_schema: self.schema(),
table_partition_cols: self.options.table_partition_cols.clone(),
writer_mode,
+ unbounded_input: self.options().infinite_source,
Review Comment:
👍🏻
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1320225484
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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(
Review Comment:
I tried to rewrite this as a `futures::stream` computation using `buffered` -- but I got stuck on some "higher-ranked lifetime error" so I think this is about as good as it is going to get
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
Review Comment:
I think this is a real error and would be nice to pass along the message to the caller
```suggestion
Err(e) => {
return Err((
writer,
DataFusionError::Execution(
format!("Error writing to object store: {e}"),
),
))
}
```
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ row_count += cnt;
+ }
+ Ok(Err(e)) => {
+ // Return the writer along with the error
+ return Err((writer, e));
+ }
+ Err(_) => {
+ // Handle task panic or cancellation
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Serialization task panicked or was cancelled".into(),
+ ),
+ ));
+ }
Review Comment:
Likewise here I think this is a real error that could be exposed to the user
```suggestion
Err(e) => {
// Handle task panic or cancellation
return Err((
writer,
DataFusionError::Execution(
format!("Serialization task panicked or was cancelled: {e}")
),
));
}
```
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
devinjdangelo commented on PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#issuecomment-1701413965
@alamb Any tips on how to benchmark this more scientifically? Would adding "insert into" and "copy" queries to our existing benchmarking framework be a reasonable idea?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1316829319
##########
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:
I believe we are facing the same issue at https://github.com/apache/arrow-datafusion/issues/5278. We could apply the same heuristic to https://github.com/apache/arrow-datafusion/issues/6290.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1316839463
##########
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:
Also, is your profiling environment multi-threaded?
##########
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:
Also, is your profiling runtime multi-threaded?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb merged pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1317692081
##########
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:
I will check the entire pull request tomorrow. Thank you for your hard work. I am also pleased that the new mechanism has contributed to the performance.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1320438041
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -137,8 +137,15 @@ impl TableProviderFactory for ListingTableFactory {
let mut statement_options = StatementOptions::from(&cmd.options);
// Extract ListingTable specific options if present or set default
- // Discard unbounded option if present
- statement_options.take_str_option("unbounded");
+ let unbounded = if infinite_source {
Review Comment:
I think it is OK if both methods are supported. That would provide backwards compatibility for anyone who used this syntax in past versions.
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -213,7 +220,8 @@ impl TableProviderFactory for ListingTableFactory {
.with_file_sort_order(cmd.order_exprs.clone())
.with_insert_mode(insert_mode)
.with_single_file(single_file)
- .with_write_options(file_type_writer_options);
+ .with_write_options(file_type_writer_options)
+ .with_infinite_source(unbounded);
Review Comment:
Thanks!
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1314847083
##########
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:
```suggestion
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 (tx, mut rx) =
mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100); // buffer size of 100, adjust as needed
let serialize_task = tokio::spawn(async move {
while let Some(maybe_batch) = data_stream.next().await {
match serializer.duplicate() {
Ok(mut serializer_clone) => {
let handle = tokio::spawn(async move {
let batch = maybe_batch?;
let num_rows = batch.num_rows();
let bytes = serializer_clone.serialize(batch).await?;
Ok((num_rows, bytes))
});
tx.send(handle).await.map_err(|_| {
DataFusionError::Internal(
"Unknown error writing to object store".into(),
)
})?;
yield_now().await;
}
Err(_) => {
return Err(DataFusionError::Internal(
"Unknown error writing to object store".into(),
))
}
}
}
Ok(serializer)
});
let mut row_count = 0;
while let Some(handle) = rx.recv().await {
match handle.await {
Ok(Ok((cnt, bytes))) => {
match writer.write_all(&bytes).await {
Ok(_) => (),
Err(_) => {
return Err((
writer,
DataFusionError::Internal(
"Unknown error writing to object store".into(),
),
))
}
};
row_count += cnt;
}
Ok(Err(e)) => {
// Return the writer along with the error
return Err((writer, e));
}
Err(_) => {
// Handle task panic or cancellation
return Err((
writer,
DataFusionError::Internal(
"Serialization task panicked or was cancelled".into(),
),
));
}
}
}
let serializer = match serialize_task.await {
Ok(Ok(serializer)) => serializer,
Ok(Err(e)) => return Err((writer, e)),
Err(_) => {
return Err((
writer,
DataFusionError::Internal("Unknown error writing to object store".into()),
))
}
};
Ok((serializer, writer, row_count as u64))
}
}
```
How about this?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1317665396
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -330,43 +391,117 @@ pub(crate) async fn stateless_serialize_and_write_files(
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(data_stream, serializer, writer)
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
+ // If we hit this code, ownership of a writer was not joined back to
+ // this thread, so we cannot clean it up (hence any_abort_errors is true)
+ any_errors = true;
+ any_abort_errors = true;
+ }
+ }
+ }
+
+ // Finalize or abort writers as appropriate
+ for mut writer in finished_writers.into_iter() {
+ match any_errors {
+ true => {
+ let abort_result = writer.abort_writer();
+ if abort_result.is_err() {
+ any_abort_errors = true;
+ }
+ }
+ false => {
+ // TODO if we encounter an error during shutdown, delete previously written files?
+ writer.shutdown()
Review Comment:
Yes, I agree that if someone wants atomic commit/rollback they should build that in at a higher level than datafusion -- there isn't much we can do with just the object store API
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1318859784
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ row_count += cnt;
+ }
+ Ok(Err(e)) => {
+ // Return the writer along with the error
+ return Err((writer, e));
+ }
+ Err(_) => {
+ // Handle task panic or cancellation
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Serialization task panicked or was cancelled".into(),
+ ),
+ ));
+ }
+ }
+ }
+
+ let serializer = match serialize_task.await {
+ Ok(Ok(serializer)) => serializer,
+ Ok(Err(e)) => return Err((writer, e)),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal("Unknown error writing to object store".into()),
+ ))
+ }
+ };
+ Ok((serializer, writer, row_count as u64))
+}
+
/// Contains the common logic for serializing RecordBatches and
/// writing the resulting bytes to an ObjectStore.
/// Serialization is assumed to be stateless, i.e.
/// each RecordBatch can be serialized without any
/// dependency on the RecordBatches before or after.
pub(crate) async fn stateless_serialize_and_write_files(
- mut data: Vec<SendableRecordBatchStream>,
+ data: Vec<SendableRecordBatchStream>,
mut serializers: Vec<Box<dyn BatchSerializer>>,
mut writers: Vec<AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>>,
single_file_output: bool,
+ unbounded_input: bool,
) -> Result<u64> {
if single_file_output && (serializers.len() != 1 || writers.len() != 1) {
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(
+ data_stream,
+ serializer,
+ writer,
+ unbounded_input,
+ )
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
+ // If we hit this code, ownership of a writer was not joined back to
+ // this thread, so we cannot clean it up (hence any_abort_errors is true)
+ any_errors = true;
+ any_abort_errors = true;
+ }
+ }
+ }
+
+ // Finalize or abort writers as appropriate
+ for mut writer in finished_writers.into_iter() {
+ match any_errors {
+ true => {
+ let abort_result = writer.abort_writer();
+ if abort_result.is_err() {
+ any_abort_errors = true;
+ }
+ }
+ false => {
+ // TODO if we encounter an error during shutdown, delete previously written files?
+ writer.shutdown()
+ .await
+ .map_err(|_| DataFusionError::Internal("Error encountered while finalizing writes! Partial results may have been written to ObjectStore!".into()))?;
+ }
+ }
+ }
+ }
+ true => {
+ let mut writer = writers.remove(0);
+ let mut serializer = serializers.remove(0);
+ let mut cnt;
+ for data_stream in data.into_iter() {
+ (serializer, writer, cnt) = match serialize_rb_stream_to_object_store(
+ data_stream,
+ serializer,
+ writer,
+ unbounded_input,
+ )
+ .await
+ {
+ Ok((s, w, c)) => (s, w, c),
Review Comment:
Comment on these types of declarations.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1319060191
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -137,8 +137,15 @@ impl TableProviderFactory for ListingTableFactory {
let mut statement_options = StatementOptions::from(&cmd.options);
// Extract ListingTable specific options if present or set default
- // Discard unbounded option if present
- statement_options.take_str_option("unbounded");
+ let unbounded = if infinite_source {
Review Comment:
I see, I think it is a leftover. I guess we can change it into Create Unbounded External if I am not missing anything.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1318878825
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -137,8 +137,15 @@ impl TableProviderFactory for ListingTableFactory {
let mut statement_options = StatementOptions::from(&cmd.options);
// Extract ListingTable specific options if present or set default
- // Discard unbounded option if present
- statement_options.take_str_option("unbounded");
+ let unbounded = if infinite_source {
Review Comment:
That syntax is supported, but the test for FIFO insert passes unbounded as a string option. The logic here was added to handle either case. Create unbounded syntax takes precedence, but if it is not present we check for OPTION (unbounded true).
https://github.com/apache/arrow-datafusion/blob/1dd0e988ceb6ecc1df4053fd7d065c15c6e8b431/datafusion/core/tests/fifo.rs#L399
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1316823526
##########
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:
Thank you for your feedback on the code. I've taken a closer look at the use of `yield_now().await` and its implications on task execution. Let's clarify a few points:
**1- Purpose of `yield_now().await`:** The `yield_now().await` call tells the Tokio runtime, "I'm temporarily done with my current work, and you can run other tasks if they're ready." It doesn't inherently force tasks to run in sequence; instead, it provides an opportunity for the runtime to switch to another task. If no other task is ready, the current task will continue almost immediately.
**2- Parallelism vs. Concurrency:** Tokio's work-stealing scheduler allows tasks to run concurrently by default. True parallel execution requires multiple threads. If your profiling runtime is configured to use a single thread, tasks will run sequentially, even if they're spawned concurrently. The `yield_now().await` call can influence the order of task execution but doesn't force sequential execution by itself.
Based on what you've said, I would suggest checking if the profiling runtime is single-threaded. Can you also provide your configuration regarding this matter?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1320160930
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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(
Review Comment:
I have some ideas how to simplify this code, which I will try out shortly, but I also think it can be merged like this too.
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -213,7 +220,8 @@ impl TableProviderFactory for ListingTableFactory {
.with_file_sort_order(cmd.order_exprs.clone())
.with_insert_mode(insert_mode)
.with_single_file(single_file)
- .with_write_options(file_type_writer_options);
+ .with_write_options(file_type_writer_options)
+ .with_infinite_source(unbounded);
Review Comment:
I took the liberty of doing so in e6f1b2fe5a0219118919da5eaa24ac9293334ff3
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -237,29 +240,11 @@ pub enum FileWriterMode {
pub trait BatchSerializer: Unpin + Send {
/// Asynchronously serializes a `RecordBatch` and returns the serialized bytes.
async fn serialize(&mut self, batch: RecordBatch) -> Result<Bytes>;
-}
-
-/// Checks if any of the passed writers have encountered an error
-/// and if so, all writers are aborted.
-async fn check_for_errors<T, W: AsyncWrite + Unpin + Send>(
- result: Result<T>,
- writers: &mut [AbortableWrite<W>],
-) -> Result<T> {
- match result {
- Ok(value) => Ok(value),
- Err(e) => {
- // Abort all writers before returning the error:
- for writer in writers {
- let mut abort_future = writer.abort_writer();
- if let Ok(abort_future) = &mut abort_future {
- let _ = abort_future.await;
- }
- // Ignore errors that occur during abortion,
- // We do try to abort all writers before returning error.
- }
- // After aborting writers return original error.
- Err(e)
- }
+ /// Duplicates self to support serializing multiple batches in parralell on multiple cores
Review Comment:
```suggestion
/// Duplicates self to support serializing multiple batches in parallel on multiple cores
```
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1311912776
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -330,43 +391,117 @@ pub(crate) async fn stateless_serialize_and_write_files(
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(data_stream, serializer, writer)
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
Review Comment:
Join errors (perhaps because a thread was killed by the OS?) could also result in non atomic writes, since we would have no way to recover ownership of the writer and abort it.
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -330,43 +391,117 @@ pub(crate) async fn stateless_serialize_and_write_files(
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(data_stream, serializer, writer)
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
+ // If we hit this code, ownership of a writer was not joined back to
+ // this thread, so we cannot clean it up (hence any_abort_errors is true)
+ any_errors = true;
+ any_abort_errors = true;
+ }
+ }
+ }
+
+ // Finalize or abort writers as appropriate
+ for mut writer in finished_writers.into_iter() {
+ match any_errors {
+ true => {
+ let abort_result = writer.abort_writer();
+ if abort_result.is_err() {
+ any_abort_errors = true;
+ }
+ }
+ false => {
+ // TODO if we encounter an error during shutdown, delete previously written files?
+ writer.shutdown()
Review Comment:
I think this is the trickiest case to ensure the write is atomic. Suppose we have two writers A and B. Writer A could successfully commit and shutdown. Then, before Writer B can complete, a network or hardware fault could prevent Writer B from either finalizing or Writer A from Aborting.
For this to be atomic, we would need some way to simultaneously commit all or none of our multipart writers. I don't think ObjectStores (S3 ect) support a way to do that.
Downstream table providers could make this atomic in practice via an atomic metadata operation, which is I believe how DeltaLake and friends work.
##########
datafusion/core/tests/fifo.rs:
##########
@@ -336,6 +336,7 @@ mod unix_test {
/// It tests the INSERT INTO functionality.
#[tokio::test]
+ #[ignore]
Review Comment:
This test is deadlocking. I think this has to do with how this test is spawning threads.
Based on the other tests passing, I don't believe this PR has broken anything with FIFO tables.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#issuecomment-1712482287
> @alamb Any tips on how to benchmark this more scientifically? Would adding "insert into" and "copy" queries to our existing benchmarking framework be a reasonable idea?
It would be great to add "write" / "copy" queries into our benchmarks. Since DataFusion only recently started supporting such statements we don't have same infrastructure for benchmarking that we do for queries
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] devinjdangelo commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "devinjdangelo (via GitHub)" <gi...@apache.org>.
devinjdangelo commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1317410490
##########
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:
Thanks for the issue links. I tried calling yield_now().await every N iterations, but the FIFO insert test only passes when N==1 (i.e. yield_now().await on every iteration).
I just pushed up a possible option of calling yield_now() if the input is unbounded, otherwise skipping it. That allows bounded tables to benefit fully from the parallelization but still allowing fifo tests to pass as they currently do.
The fifo test is not a multithreaded runtime, but when I benchmark performance I am running in a multithreaded runtime, yes.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1314542376
##########
datafusion/core/tests/fifo.rs:
##########
@@ -336,6 +336,7 @@ mod unix_test {
/// It tests the INSERT INTO functionality.
#[tokio::test]
+ #[ignore]
Review Comment:
I am currently checking for any potential problems.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1314561329
##########
datafusion/core/tests/fifo.rs:
##########
@@ -336,6 +336,7 @@ mod unix_test {
/// It tests the INSERT INTO functionality.
#[tokio::test]
+ #[ignore]
Review Comment:
It's not the way the threads are started that's the issue, but rather the assumption that if a FIFO file receives 10 rows and the batch size is also 10, it will create a record batch. The current implementation of `serialize_rb_stream_to_object_store` requires access to all data before it can proceed.
```rust
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(),
),
))
}
};
let batch = maybe_batch.unwrap();
print_batches(&[batch.clone()]).unwrap();
serialize_tasks.push(task::spawn(async move {
let num_rows = batch.num_rows();
let bytes = serializer_clone.serialize(batch).await?;
Ok((num_rows, bytes))
}));
}
```
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1316823526
##########
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:
Thank you for your feedback on the code. I've taken a closer look at the use of `yield_now().await` and its implications on task execution. Let's clarify a few points:
**1- Purpose of `yield_now().await`:** The `yield_now().await` call tells the Tokio runtime, "I'm temporarily done with my current work, and you can run other tasks if they're ready." It doesn't inherently force tasks to run in sequence; instead, it provides an opportunity for the runtime to switch to another task. If no other task is ready, the current task will continue almost immediately.
**2- Parallelism vs. Concurrency:** Tokio's work-stealing scheduler allows tasks to run concurrently by default. True parallel execution requires multiple threads. If your profiling runtime is configured to use a single thread, tasks will run sequentially, even if they're spawned concurrently. The `yield_now().await` call can influence the order of task execution but doesn't force sequential execution by itself.
Based on what you've said, I would suggest checking if the profiling runtime is single-threaded. Can you also provide your configuration regarding this matter?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#issuecomment-1712482645
I think all comments have now been addressed. Thank you very much @devinjdangelo and @metesynnada -- this is really cool to see. 🚀
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1314561329
##########
datafusion/core/tests/fifo.rs:
##########
@@ -336,6 +336,7 @@ mod unix_test {
/// It tests the INSERT INTO functionality.
#[tokio::test]
+ #[ignore]
Review Comment:
It's not the way the threads are started that's the issue, but rather the assumption that if a FIFO file receives 10 rows and the batch size is also 10, it will create a record batch. The current implementation of `serialize_rb_stream_to_object_store` requires access to all data before it can proceed.
```rust
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))
}));
}
```
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1314561329
##########
datafusion/core/tests/fifo.rs:
##########
@@ -336,6 +336,7 @@ mod unix_test {
/// It tests the INSERT INTO functionality.
#[tokio::test]
+ #[ignore]
Review Comment:
It's not the way the threads are started that's the issue, but rather the assumption that if a FIFO file receives 10 rows and the batch size is also 10, it will create a record batch. The current implementation of `serialize_rb_stream_to_object_store` requires access to all data before it can proceed.
```rust
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(),
),
))
}
};
let batch = maybe_batch.unwrap();
print_batches(&[batch.clone()]).unwrap();
serialize_tasks.push(task::spawn(async move {
let num_rows = batch.num_rows();
let bytes = serializer_clone.serialize(batch).await?;
Ok((num_rows, bytes))
}));
}
```
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb commented on pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#issuecomment-1712016952
From my point of view this PR is ready to go -- There are still some suggestion from @metesynnada, but I am not sure if they feel any are required prior to merge.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1318827920
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ row_count += cnt;
+ }
+ Ok(Err(e)) => {
+ // Return the writer along with the error
+ return Err((writer, e));
+ }
+ Err(_) => {
+ // Handle task panic or cancellation
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Serialization task panicked or was cancelled".into(),
+ ),
+ ));
+ }
+ }
+ }
+
+ let serializer = match serialize_task.await {
+ Ok(Ok(serializer)) => serializer,
+ Ok(Err(e)) => return Err((writer, e)),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal("Unknown error writing to object store".into()),
+ ))
+ }
+ };
+ Ok((serializer, writer, row_count as u64))
+}
+
/// Contains the common logic for serializing RecordBatches and
/// writing the resulting bytes to an ObjectStore.
/// Serialization is assumed to be stateless, i.e.
/// each RecordBatch can be serialized without any
/// dependency on the RecordBatches before or after.
pub(crate) async fn stateless_serialize_and_write_files(
- mut data: Vec<SendableRecordBatchStream>,
+ data: Vec<SendableRecordBatchStream>,
mut serializers: Vec<Box<dyn BatchSerializer>>,
mut writers: Vec<AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>>,
single_file_output: bool,
+ unbounded_input: bool,
) -> Result<u64> {
if single_file_output && (serializers.len() != 1 || writers.len() != 1) {
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(
+ data_stream,
+ serializer,
+ writer,
+ unbounded_input,
+ )
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
+ // If we hit this code, ownership of a writer was not joined back to
+ // this thread, so we cannot clean it up (hence any_abort_errors is true)
+ any_errors = true;
+ any_abort_errors = true;
+ }
+ }
+ }
+
+ // Finalize or abort writers as appropriate
+ for mut writer in finished_writers.into_iter() {
+ match any_errors {
Review Comment:
Can be `if any_errors { ... } else { ... }`
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ row_count += cnt;
+ }
+ Ok(Err(e)) => {
+ // Return the writer along with the error
+ return Err((writer, e));
+ }
+ Err(_) => {
+ // Handle task panic or cancellation
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Serialization task panicked or was cancelled".into(),
+ ),
+ ));
+ }
+ }
+ }
+
+ let serializer = match serialize_task.await {
+ Ok(Ok(serializer)) => serializer,
+ Ok(Err(e)) => return Err((writer, e)),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal("Unknown error writing to object store".into()),
+ ))
+ }
+ };
+ Ok((serializer, writer, row_count as u64))
+}
+
/// Contains the common logic for serializing RecordBatches and
/// writing the resulting bytes to an ObjectStore.
/// Serialization is assumed to be stateless, i.e.
/// each RecordBatch can be serialized without any
/// dependency on the RecordBatches before or after.
pub(crate) async fn stateless_serialize_and_write_files(
- mut data: Vec<SendableRecordBatchStream>,
+ data: Vec<SendableRecordBatchStream>,
mut serializers: Vec<Box<dyn BatchSerializer>>,
mut writers: Vec<AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>>,
single_file_output: bool,
+ unbounded_input: bool,
) -> Result<u64> {
if single_file_output && (serializers.len() != 1 || writers.len() != 1) {
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
Review Comment:
Can be ```if single_file_output {
...
} else {
...
}```
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -330,43 +391,117 @@ pub(crate) async fn stateless_serialize_and_write_files(
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(data_stream, serializer, writer)
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
Review Comment:
You should return the error after aborting the writer, in case an execution error occurs within the plans.
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
Review Comment:
The function signature is quite complex with nested types. Using type aliases can simplify the signature and make it more readable. Like
```rust
type WriterType = AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>;
type ResultTuple = (Box<dyn BatchSerializer>, WriterType, u64);
type ErrorTuple = (WriterType, DataFusionError);
```
##########
datafusion/core/src/datasource/listing/table.rs:
##########
@@ -915,6 +915,7 @@ impl TableProvider for ListingTable {
output_schema: self.schema(),
table_partition_cols: self.options.table_partition_cols.clone(),
writer_mode,
+ unbounded_input: self.options().infinite_source,
Review Comment:
👍🏻
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -137,8 +137,15 @@ impl TableProviderFactory for ListingTableFactory {
let mut statement_options = StatementOptions::from(&cmd.options);
// Extract ListingTable specific options if present or set default
- // Discard unbounded option if present
- statement_options.take_str_option("unbounded");
+ let unbounded = if infinite_source {
Review Comment:
You can remove this logic since we support unbounded tables with `CREATE UNBOUNDED EXTERNAL TABLE`.
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ row_count += cnt;
+ }
+ Ok(Err(e)) => {
+ // Return the writer along with the error
+ return Err((writer, e));
+ }
+ Err(_) => {
+ // Handle task panic or cancellation
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Serialization task panicked or was cancelled".into(),
+ ),
+ ));
+ }
+ }
+ }
+
+ let serializer = match serialize_task.await {
+ Ok(Ok(serializer)) => serializer,
+ Ok(Err(e)) => return Err((writer, e)),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal("Unknown error writing to object store".into()),
+ ))
+ }
+ };
+ Ok((serializer, writer, row_count as u64))
+}
+
/// Contains the common logic for serializing RecordBatches and
/// writing the resulting bytes to an ObjectStore.
/// Serialization is assumed to be stateless, i.e.
/// each RecordBatch can be serialized without any
/// dependency on the RecordBatches before or after.
pub(crate) async fn stateless_serialize_and_write_files(
- mut data: Vec<SendableRecordBatchStream>,
+ data: Vec<SendableRecordBatchStream>,
mut serializers: Vec<Box<dyn BatchSerializer>>,
mut writers: Vec<AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>>,
single_file_output: bool,
+ unbounded_input: bool,
) -> Result<u64> {
if single_file_output && (serializers.len() != 1 || writers.len() != 1) {
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
let num_partitions = data.len();
- if !single_file_output && (num_partitions != writers.len()) {
+ let num_writers = writers.len();
+ if !single_file_output && (num_partitions != num_writers) {
return internal_err!("single_file_ouput is false, but did not get 1 writer for each output partition!");
}
let mut row_count = 0;
- // Map errors to DatafusionError.
- let err_converter =
- |_| DataFusionError::Internal("Unexpected FileSink Error".to_string());
- // TODO parallelize serialization accross partitions and batches within partitions
- // see: https://github.com/apache/arrow-datafusion/issues/7079
- for (part_idx, data_stream) in data.iter_mut().enumerate().take(num_partitions) {
- let idx = match single_file_output {
- false => part_idx,
- true => 0,
- };
- while let Some(maybe_batch) = data_stream.next().await {
- // Write data to files in a round robin fashion:
- let serializer = &mut serializers[idx];
- let batch = check_for_errors(maybe_batch, &mut writers).await?;
- row_count += batch.num_rows();
- let bytes =
- check_for_errors(serializer.serialize(batch).await, &mut writers).await?;
- let writer = &mut writers[idx];
- check_for_errors(
- writer.write_all(&bytes).await.map_err(err_converter),
- &mut writers,
- )
- .await?;
+ // tracks if any writers encountered an error triggering the need to abort
+ let mut any_errors = false;
+ // tracks the specific error triggering abort
+ let mut triggering_error = None;
+ // tracks if any errors were encountered in the process of aborting writers.
+ // if true, we may not have a guarentee that all written data was cleaned up.
+ let mut any_abort_errors = false;
+ match single_file_output {
+ false => {
+ let mut join_set = JoinSet::new();
+ for (data_stream, serializer, writer) in data
+ .into_iter()
+ .zip(serializers.into_iter())
+ .zip(writers.into_iter())
+ .map(|((a, b), c)| (a, b, c))
+ {
+ join_set.spawn(async move {
+ serialize_rb_stream_to_object_store(
+ data_stream,
+ serializer,
+ writer,
+ unbounded_input,
+ )
+ .await
+ });
+ }
+ let mut finished_writers = Vec::with_capacity(num_writers);
+ while let Some(result) = join_set.join_next().await {
+ match result {
+ Ok(res) => match res {
+ Ok((_, writer, cnt)) => {
+ finished_writers.push(writer);
+ row_count += cnt;
+ }
+ Err((writer, e)) => {
+ finished_writers.push(writer);
+ any_errors = true;
+ triggering_error = Some(e);
+ }
+ },
+ Err(_) => {
+ // Don't panic, instead try to clean up as many writers as possible.
+ // If we hit this code, ownership of a writer was not joined back to
+ // this thread, so we cannot clean it up (hence any_abort_errors is true)
+ any_errors = true;
+ any_abort_errors = true;
+ }
+ }
+ }
+
+ // Finalize or abort writers as appropriate
+ for mut writer in finished_writers.into_iter() {
+ match any_errors {
+ true => {
+ let abort_result = writer.abort_writer();
+ if abort_result.is_err() {
+ any_abort_errors = true;
+ }
+ }
+ false => {
+ // TODO if we encounter an error during shutdown, delete previously written files?
+ writer.shutdown()
+ .await
+ .map_err(|_| DataFusionError::Internal("Error encountered while finalizing writes! Partial results may have been written to ObjectStore!".into()))?;
+ }
+ }
+ }
+ }
+ true => {
+ let mut writer = writers.remove(0);
+ let mut serializer = serializers.remove(0);
+ let mut cnt;
+ for data_stream in data.into_iter() {
+ (serializer, writer, cnt) = match serialize_rb_stream_to_object_store(
+ data_stream,
+ serializer,
+ writer,
+ unbounded_input,
+ )
+ .await
+ {
+ Ok((s, w, c)) => (s, w, c),
Review Comment:
Comment on these types of declarations.
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
Review Comment:
Define the return error as a closure at the beginning of the method to make it common.
##########
datafusion/core/src/datasource/file_format/write.rs:
##########
@@ -315,58 +300,237 @@ 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>>,
+ unbounded_input: bool,
+) -> std::result::Result<
+ (
+ Box<dyn BatchSerializer>,
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ u64,
+ ),
+ (
+ AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>,
+ DataFusionError,
+ ),
+> {
+ let (tx, mut rx) =
+ mpsc::channel::<JoinHandle<Result<(usize, Bytes), DataFusionError>>>(100);
+
+ let serialize_task = tokio::spawn(async move {
+ while let Some(maybe_batch) = data_stream.next().await {
+ match serializer.duplicate() {
+ Ok(mut serializer_clone) => {
+ let handle = tokio::spawn(async move {
+ let batch = maybe_batch?;
+ let num_rows = batch.num_rows();
+ let bytes = serializer_clone.serialize(batch).await?;
+ Ok((num_rows, bytes))
+ });
+ tx.send(handle).await.map_err(|_| {
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ )
+ })?;
+ if unbounded_input {
+ tokio::task::yield_now().await;
+ }
+ }
+ Err(_) => {
+ return Err(DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ))
+ }
+ }
+ }
+ Ok(serializer)
+ });
+
+ let mut row_count = 0;
+ while let Some(handle) = rx.recv().await {
+ match handle.await {
+ Ok(Ok((cnt, bytes))) => {
+ match writer.write_all(&bytes).await {
+ Ok(_) => (),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Unknown error writing to object store".into(),
+ ),
+ ))
+ }
+ };
+ row_count += cnt;
+ }
+ Ok(Err(e)) => {
+ // Return the writer along with the error
+ return Err((writer, e));
+ }
+ Err(_) => {
+ // Handle task panic or cancellation
+ return Err((
+ writer,
+ DataFusionError::Internal(
+ "Serialization task panicked or was cancelled".into(),
+ ),
+ ));
+ }
+ }
+ }
+
+ let serializer = match serialize_task.await {
+ Ok(Ok(serializer)) => serializer,
+ Ok(Err(e)) => return Err((writer, e)),
+ Err(_) => {
+ return Err((
+ writer,
+ DataFusionError::Internal("Unknown error writing to object store".into()),
+ ))
+ }
+ };
+ Ok((serializer, writer, row_count as u64))
+}
+
/// Contains the common logic for serializing RecordBatches and
/// writing the resulting bytes to an ObjectStore.
/// Serialization is assumed to be stateless, i.e.
/// each RecordBatch can be serialized without any
/// dependency on the RecordBatches before or after.
pub(crate) async fn stateless_serialize_and_write_files(
Review Comment:
The blocks inside the match `serialize_rb_stream_to_object_store` can be refactored into separate functions for better readability and maintainability.
```rust
pub(crate) async fn stateless_serialize_and_write_files(
data: Vec<SendableRecordBatchStream>,
mut serializers: Vec<Box<dyn BatchSerializer>>,
mut writers: Vec<AbortableWrite<Box<dyn AsyncWrite + Send + Unpin>>>,
single_file_output: bool,
unbounded_input: bool,
) -> Result<u64> {
if single_file_output {
if serializers.len() != 1 || writers.len() != 1 {
return internal_err!("single_file_output is true, but got more than 1 writer!");
}
return serialize_single_file(data, serializers.remove(0), writers.remove(0), unbounded_input).await;
} else {
if data.len() != writers.len() {
return internal_err!("Did not get 1 writer for each output partition!");
}
return serialize_multiple_files(data, serializers, writers, unbounded_input).await;
}
}
```
##########
datafusion/core/src/datasource/listing/table.rs:
##########
@@ -915,6 +915,7 @@ impl TableProvider for ListingTable {
output_schema: self.schema(),
table_partition_cols: self.options.table_partition_cols.clone(),
writer_mode,
+ unbounded_input: self.options().infinite_source,
Review Comment:
👍🏻
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -213,7 +220,8 @@ impl TableProviderFactory for ListingTableFactory {
.with_file_sort_order(cmd.order_exprs.clone())
.with_insert_mode(insert_mode)
.with_single_file(single_file)
- .with_write_options(file_type_writer_options);
+ .with_write_options(file_type_writer_options)
+ .with_infinite_source(unbounded);
Review Comment:
Please remove the line with `.with_infinite_source(unbounded);` as it duplicates the logic above.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] metesynnada commented on a diff in pull request #7452: Parallelize Stateless (CSV/JSON) File Write Serialization
Posted by "metesynnada (via GitHub)" <gi...@apache.org>.
metesynnada commented on code in PR #7452:
URL: https://github.com/apache/arrow-datafusion/pull/7452#discussion_r1319060191
##########
datafusion/core/src/datasource/listing_table_factory.rs:
##########
@@ -137,8 +137,15 @@ impl TableProviderFactory for ListingTableFactory {
let mut statement_options = StatementOptions::from(&cmd.options);
// Extract ListingTable specific options if present or set default
- // Discard unbounded option if present
- statement_options.take_str_option("unbounded");
+ let unbounded = if infinite_source {
Review Comment:
I see, I think it is a leftover. I guess we can change the test into Create Unbounded External if I am not missing anything.
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org