This is an automated email from the ASF dual-hosted git repository.
nju_yaho pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-ballista.git
The following commit(s) were added to refs/heads/master by this push:
new 913f6759 Make fetch shuffle partition data in parallel (#256)
913f6759 is described below
commit 913f675955adaf6335002ba9a5c56beef5cce56e
Author: yahoNanJing <[email protected]>
AuthorDate: Tue Sep 27 22:43:34 2022 +0800
Make fetch shuffle partition data in parallel (#256)
* Make fetch shuffle partition data in parallel
* Add unit test
* Fast abort shuffle reader futures when error occurs
Co-authored-by: yangzhong <[email protected]>
---
ballista/rust/core/Cargo.toml | 3 +
.../core/src/execution_plans/shuffle_reader.rs | 211 ++++++++++++++++++---
ballista/rust/executor/src/flight_service.rs | 9 +-
3 files changed, 196 insertions(+), 27 deletions(-)
diff --git a/ballista/rust/core/Cargo.toml b/ballista/rust/core/Cargo.toml
index b484047e..e7fea396 100644
--- a/ballista/rust/core/Cargo.toml
+++ b/ballista/rust/core/Cargo.toml
@@ -50,6 +50,7 @@ datafusion-proto = "12.0.0"
futures = "0.3"
hashbrown = "0.12"
+itertools = "0.10"
libloading = "0.7.3"
log = "0.4"
object_store = "0.5.0"
@@ -59,9 +60,11 @@ parking_lot = "0.12"
parse_arg = "0.1.3"
prost = "0.11"
prost-types = "0.11"
+rand = "0.8"
serde = { version = "1", features = ["derive"] }
sqlparser = "0.23"
tokio = "1.0"
+tokio-stream = { version = "0.1", features = ["net"] }
tonic = "0.8"
url = "2.2"
uuid = { version = "1.0", features = ["v4"] }
diff --git a/ballista/rust/core/src/execution_plans/shuffle_reader.rs
b/ballista/rust/core/src/execution_plans/shuffle_reader.rs
index 0c153d3e..d0d9a28b 100644
--- a/ballista/rust/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/rust/core/src/execution_plans/shuffle_reader.rs
@@ -16,8 +16,10 @@
// under the License.
use std::any::Any;
+use std::collections::HashMap;
use std::sync::Arc;
+#[cfg(not(test))]
use crate::client::BallistaClient;
use crate::serde::scheduler::{PartitionLocation, PartitionStats};
@@ -25,18 +27,22 @@ use datafusion::arrow::datatypes::SchemaRef;
use datafusion::error::{DataFusionError, Result};
use datafusion::physical_plan::expressions::PhysicalSortExpr;
-use datafusion::physical_plan::metrics::{
- ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
-};
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::{
DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
Statistics,
};
-use futures::{StreamExt, TryStreamExt};
+use futures::{Stream, StreamExt, TryStreamExt};
-use datafusion::arrow::error::ArrowError;
use datafusion::execution::context::TaskContext;
+use datafusion::physical_plan::common::AbortOnDropMany;
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use log::debug;
+use itertools::Itertools;
+use log::{error, info};
+use rand::prelude::SliceRandom;
+use rand::thread_rng;
+use tokio::sync::{mpsc, Semaphore};
+use tokio::task::JoinHandle;
+use tokio_stream::wrappers::ReceiverStream;
/// ShuffleReaderExec reads partitions that have already been materialized by
a ShuffleWriterExec
/// being executed by an executor
@@ -102,29 +108,35 @@ impl ExecutionPlan for ShuffleReaderExec {
fn execute(
&self,
partition: usize,
- _context: Arc<TaskContext>,
+ context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
- debug!("ShuffleReaderExec::execute({})", partition);
-
- let fetch_time =
- MetricBuilder::new(&self.metrics).subset_time("fetch_time",
partition);
-
- let locations = self.partition[partition].clone();
- let stream = locations.into_iter().map(move |p| {
- let fetch_time = fetch_time.clone();
- futures::stream::once(async move {
- let timer = fetch_time.timer();
- let r = fetch_partition(&p).await;
- timer.done();
-
- r.map_err(|e| ArrowError::ExternalError(Box::new(e)))
- })
- .try_flatten()
- });
+ let task_id = context.task_id().unwrap_or_else(||
partition.to_string());
+ info!("ShuffleReaderExec::execute({})", task_id);
+
+ // TODO make the maximum size configurable, or make it depends on
global memory control
+ let max_request_num = 50usize;
+ let mut partition_locations = HashMap::new();
+ for p in &self.partition[partition] {
+ partition_locations
+ .entry(p.executor_meta.id.clone())
+ .or_insert_with(Vec::new)
+ .push(p.clone());
+ }
+ // Sort partitions for evenly send fetching partition requests to
avoid hot executors within one task
+ let mut partition_locations: Vec<PartitionLocation> =
partition_locations
+ .into_values()
+ .flat_map(|ps| ps.into_iter().enumerate())
+ .sorted_by(|(p1_idx, _), (p2_idx, _)| Ord::cmp(p1_idx, p2_idx))
+ .map(|(_, p)| p)
+ .collect();
+ // Shuffle partitions for evenly send fetching partition requests to
avoid hot executors within multiple tasks
+ partition_locations.shuffle(&mut thread_rng());
+ let response_receiver =
+ send_fetch_partitions(partition_locations, max_request_num);
let result = RecordBatchStreamAdapter::new(
Arc::new(self.schema.as_ref().clone()),
- futures::stream::iter(stream).flatten(),
+ response_receiver.try_flatten(),
);
Ok(Box::pin(result))
}
@@ -178,6 +190,67 @@ fn stats_for_partitions(
)
}
+/// Adapter for a tokio ReceiverStream that implements the
SendableRecordBatchStream interface
+struct AbortableReceiverStream {
+ inner: ReceiverStream<Result<SendableRecordBatchStream>>,
+
+ #[allow(dead_code)]
+ drop_helper: AbortOnDropMany<()>,
+}
+
+impl AbortableReceiverStream {
+ /// Construct a new SendableRecordBatchReceiverStream which will send
batches of the specified schema from inner
+ pub fn create(
+ rx: tokio::sync::mpsc::Receiver<Result<SendableRecordBatchStream>>,
+ join_handles: Vec<JoinHandle<()>>,
+ ) -> AbortableReceiverStream {
+ let inner = ReceiverStream::new(rx);
+ Self {
+ inner,
+ drop_helper: AbortOnDropMany(join_handles),
+ }
+ }
+}
+
+impl Stream for AbortableReceiverStream {
+ type Item = Result<SendableRecordBatchStream>;
+
+ fn poll_next(
+ mut self: std::pin::Pin<&mut Self>,
+ cx: &mut std::task::Context<'_>,
+ ) -> std::task::Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+fn send_fetch_partitions(
+ partition_locations: Vec<PartitionLocation>,
+ max_request_num: usize,
+) -> AbortableReceiverStream {
+ let (response_sender, response_receiver) = mpsc::channel(max_request_num);
+ let semaphore = Arc::new(Semaphore::new(max_request_num));
+ let mut join_handles = vec![];
+ for p in partition_locations.into_iter() {
+ let semaphore = semaphore.clone();
+ let response_sender = response_sender.clone();
+ let join_handle = tokio::spawn(async move {
+ // Block if exceeds max request number
+ let permit = semaphore.acquire_owned().await.unwrap();
+ let r = fetch_partition(&p).await;
+ // Block if the channel buffer is full
+ if let Err(e) = response_sender.send(r).await {
+ error!("Fail to send response event to the channel due to {}",
e);
+ }
+ // Increase semaphore by dropping existing permits.
+ drop(permit);
+ });
+ join_handles.push(join_handle);
+ }
+
+ AbortableReceiverStream::create(response_receiver, join_handles)
+}
+
+#[cfg(not(test))]
async fn fetch_partition(
location: &PartitionLocation,
) -> Result<SendableRecordBatchStream> {
@@ -203,9 +276,22 @@ async fn fetch_partition(
.map_err(|e| DataFusionError::Execution(format!("{:?}", e)))
}
+#[cfg(test)]
+async fn fetch_partition(
+ location: &PartitionLocation,
+) -> Result<SendableRecordBatchStream> {
+ tests::fetch_test_partition(location)
+}
+
#[cfg(test)]
mod tests {
use super::*;
+ use crate::serde::scheduler::{ExecutorMetadata, ExecutorSpecification,
PartitionId};
+ use datafusion::arrow::array::Int32Array;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::physical_plan::common;
+ use datafusion::physical_plan::stream::RecordBatchReceiverStream;
#[tokio::test]
async fn test_stats_for_partitions_empty() {
@@ -274,4 +360,79 @@ mod tests {
assert_eq!(result, exptected);
}
+
+ #[tokio::test]
+ async fn test_send_fetch_partitions_1() {
+ test_send_fetch_partitions(1, 10).await;
+ }
+
+ #[tokio::test]
+ async fn test_send_fetch_partitions_n() {
+ test_send_fetch_partitions(4, 10).await;
+ }
+
+ async fn test_send_fetch_partitions(max_request_num: usize, partition_num:
usize) {
+ let schema = get_test_partition_schema();
+ let partition_locations = get_test_partition_locations(partition_num);
+ let response_receiver =
+ send_fetch_partitions(partition_locations, max_request_num);
+
+ let stream = RecordBatchStreamAdapter::new(
+ Arc::new(schema),
+ response_receiver.try_flatten(),
+ );
+
+ let result = common::collect(Box::pin(stream)).await.unwrap();
+ assert_eq!(partition_num, result.len());
+ }
+
+ fn get_test_partition_locations(n: usize) -> Vec<PartitionLocation> {
+ (0..n)
+ .into_iter()
+ .map(|partition_id| PartitionLocation {
+ partition_id: PartitionId {
+ job_id: "job".to_string(),
+ stage_id: 1,
+ partition_id,
+ },
+ executor_meta: ExecutorMetadata {
+ id: format!("exec{}", partition_id),
+ host: "localhost".to_string(),
+ port: 50051,
+ grpc_port: 50052,
+ specification: ExecutorSpecification { task_slots: 12 },
+ },
+ partition_stats: Default::default(),
+ path: format!("/tmp/job/1/{}", partition_id),
+ })
+ .collect()
+ }
+
+ pub(crate) fn fetch_test_partition(
+ location: &PartitionLocation,
+ ) -> Result<SendableRecordBatchStream> {
+ let id_array =
Int32Array::from(vec![location.partition_id.partition_id as i32]);
+ let schema = Arc::new(get_test_partition_schema());
+
+ let batch = RecordBatch::try_new(schema.clone(),
vec![Arc::new(id_array)])?;
+
+ let (tx, rx) = tokio::sync::mpsc::channel(2);
+
+ // task simply sends data in order but in a separate
+ // thread (to ensure the batches are not available without the
+ // DelayedStream yielding).
+ let join_handle = tokio::task::spawn(async move {
+ println!("Sending batch via delayed stream");
+ if let Err(e) = tx.send(Ok(batch)).await {
+ println!("ERROR batch via delayed stream: {}", e);
+ }
+ });
+
+ // returned stream simply reads off the rx stream
+ Ok(RecordBatchReceiverStream::create(&schema, rx, join_handle))
+ }
+
+ fn get_test_partition_schema() -> Schema {
+ Schema::new(vec![Field::new("id", DataType::Int32, false)])
+ }
}
diff --git a/ballista/rust/executor/src/flight_service.rs
b/ballista/rust/executor/src/flight_service.rs
index 82c4f0ae..fde4f045 100644
--- a/ballista/rust/executor/src/flight_service.rs
+++ b/ballista/rust/executor/src/flight_service.rs
@@ -104,10 +104,11 @@ impl FlightService for BallistaFlightService {
let (tx, rx): (FlightDataSender, FlightDataReceiver) =
channel(2);
+ let file_path = path.to_owned();
// Arrow IPC reader does not implement Sync + Send so we need
to use a channel
// to communicate
task::spawn(async move {
- if let Err(e) = stream_flight_data(reader, tx).await {
+ if let Err(e) = stream_flight_data(file_path, reader,
tx).await {
warn!("Error streaming results: {:?}", e);
}
});
@@ -220,6 +221,7 @@ fn create_flight_iter(
}
async fn stream_flight_data<T>(
+ file_path: String,
reader: FileReader<T>,
tx: FlightDataSender,
) -> Result<(), Status>
@@ -242,7 +244,10 @@ where
send_response(&tx, batch).await?;
}
}
- debug!("FetchPartition streamed {} rows", row_count);
+ debug!(
+ "FetchPartition streamed {} rows for file {}",
+ row_count, file_path
+ );
Ok(())
}
